Upload folder using huggingface_hub

README.md

@@ -1,3 +1,101 @@
 ---
+library_name: transformers
+tags:
+- text-generation
+- trust-remote-code
+- sentencepiece
+- custom-architecture
+pipeline_tag: text-generation
 license: mit
 ---
+
+# TaoNet-mini-A2
+
+`TaoNet-mini-A2` is packaged as a Hugging Face `trust_remote_code` repository for causal text generation.
+
+## Files
+
+The repo root contains:
+
+- `configuration_taonet.py`
+- `modeling_taonet.py`
+- `tokenization_taonet.py`
+- `taonet_model.py`
+- `mla_components.py`
+- `embeddings.py`
+
+The checkpoint export helper is:
+
+- `export_to_hf.py`
+
+Run that script once on your machine to convert `checkpoints/sft/final_model.pt` into a standard Hugging Face weight file such as `pytorch_model.bin`.
+
+## Usage
+
+```python
+import time
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_NAME = "TaoTern/TaoNet-mini-A2"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True,
+    torch_dtype=dtype,
+).to(device)
+
+
+def generate_text(prompt, max_new_tokens=64, temperature=0.7, top_p=0.85):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = {key: value.to(device) for key, value in inputs.items()}
+
+    start_time = time.time()
+    with torch.inference_mode():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=1.2,
+            do_sample=True,
+            use_cache=False,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    elapsed_time = time.time() - start_time
+
+    new_tokens = outputs.shape[1] - inputs["input_ids"].shape[1]
+    tokens_per_second = new_tokens / elapsed_time if elapsed_time > 0 else 0.0
+    completion = tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[1] :],
+        skip_special_tokens=True,
+    )
+    return completion, tokens_per_second
+
+
+if __name__ == "__main__":
+    text, tps = generate_text("Fruit is now expensive so we should")
+    print(text)
+    print(f"\\nTokens per second: {tps:.2f}")
+```
+
+## Export
+
+When you are ready to package the trained checkpoint into standard HF weights, run:
+
+```bash
+python export_to_hf.py
+```
+
+That script will:
+
+- load `checkpoints/sft/final_model.pt`
+- rebuild the TaoNet config from the checkpoint metadata
+- load the weights into the HF wrapper
+- write standard Hugging Face model files at repo root
+- copy `tokenizer.model`, `tokenizer.vocab`, and `tokenizer.special_tokens.json` to repo root

added_tokens.json

@@ -0,0 +1,6 @@
+{
+  "<BOS>": 8192,
+  "<EOS>": 8193,
+  "<PAD>": 8195,
+  "<UNK>": 8194
+}

config.json

@@ -0,0 +1,59 @@
+{
+  "architectures": [
+    "TaoNetForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "cnn_channels": [
+    32,
+    64,
+    128
+  ],
+  "cnn_kernel_size": 3,
+  "d_embed_rank": 96,
+  "d_latent_kv": 768,
+  "d_rope": 128,
+  "dropout": 0.02,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gamma_activation": "gelu",
+  "gamma_discretization": "bilinear",
+  "gamma_dt_init": 0.01,
+  "gamma_dt_max": 0.1,
+  "gamma_dt_min": 0.001,
+  "gamma_gate": true,
+  "gamma_gate_bias": 2.0,
+  "gamma_hidden_dim": 1536,
+  "gamma_input_gate": true,
+  "gamma_input_gate_bias": 2.0,
+  "gamma_kernel_mode": "auto",
+  "gamma_kernel_threshold": 64,
+  "gamma_layer_scale_init": 0.1,
+  "gamma_prenorm": true,
+  "gamma_residual_scale": 1.0,
+  "gamma_use_D": true,
+  "gamma_use_output_linear": true,
+  "gqa_groups": 1,
+  "head_dim": 128,
+  "hidden_dim": 1024,
+  "hidden_dim_ff": 3072,
+  "image_size": 224,
+  "image_token": "<image>",
+  "init_std": 0.02,
+  "intermediate_dim": 4096,
+  "max_seq_length": 1024,
+  "model_type": "taonet",
+  "num_heads": 8,
+  "num_layers": 16,
+  "pad_token_id": 3,
+  "rope_scale": 40.0,
+  "transformers_version": "5.1.0",
+  "unk_token_id": 0,
+  "use_factorized_embedding": true,
+  "vision_encoder_type": "cnn",
+  "vision_output_dim": 256,
+  "vision_prefix_tokens": 10,
+  "vocab_size": 8192,
+  "yarn_alpha": 1.0,
+  "yarn_enabled": false,
+  "yarn_original_max_seq_length": null
+}

configuration_taonet.py

@@ -0,0 +1,72 @@
+"""Hugging Face configuration for TaoNet."""
+
+from transformers import PretrainedConfig
+
+
+class TaoNetConfig(PretrainedConfig):
+    """Configuration for the TaoNet causal language model."""
+
+    model_type = "taonet"
+
+    def __init__(
+        self,
+        vocab_size=8192,
+        hidden_dim=768,
+        num_layers=12,
+        num_heads=12,
+        head_dim=None,
+        intermediate_dim=None,
+        dropout=0.1,
+        max_seq_length=2048,
+        d_latent_kv=None,
+        d_rope=None,
+        gqa_groups=1,
+        hidden_dim_ff=None,
+        use_factorized_embedding=False,
+        d_embed_rank=96,
+        rope_scale=40.0,
+        yarn_enabled=False,
+        yarn_original_max_seq_length=None,
+        yarn_alpha=1.0,
+        init_std=0.02,
+        pad_token_id=3,
+        bos_token_id=1,
+        eos_token_id=2,
+        unk_token_id=0,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.head_dim = head_dim if head_dim is not None else hidden_dim // num_heads
+        self.intermediate_dim = intermediate_dim if intermediate_dim is not None else hidden_dim * 4
+        self.dropout = dropout
+        self.max_seq_length = max_seq_length
+        self.d_latent_kv = d_latent_kv
+        self.d_rope = d_rope
+        self.gqa_groups = gqa_groups
+        self.hidden_dim_ff = hidden_dim_ff
+        self.use_factorized_embedding = use_factorized_embedding
+        self.d_embed_rank = d_embed_rank
+        self.rope_scale = rope_scale
+        self.yarn_enabled = yarn_enabled
+        self.yarn_original_max_seq_length = yarn_original_max_seq_length
+        self.yarn_alpha = yarn_alpha
+        self.init_std = init_std
+
+    @classmethod
+    def from_taotrain_model_config(cls, model_config, **extra_kwargs):
+        """Build an HF config from a TaoTrain model-config dictionary."""
+        data = dict(model_config)
+        data.pop("architecture_type", None)
+        data.pop("llm_architecture_type", None)
+        data.update(extra_kwargs)
+        return cls(**data)

embeddings.py

@@ -0,0 +1,23 @@
+"""Embedding layers used by TaoNet."""
+
+import torch.nn as nn
+
+
+class FactorizedEmbedding(nn.Module):
+    """Low-rank factorized embedding."""
+
+    def __init__(self, vocab_size, d_model, d_embed_rank=96):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_embed_rank = d_embed_rank
+
+        self.embed = nn.Embedding(vocab_size, d_embed_rank)
+        self.proj = nn.Linear(d_embed_rank, d_model, bias=False)
+
+        nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.proj.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids):
+        x = self.embed(input_ids)
+        return self.proj(x)

export_to_hf.py

@@ -0,0 +1,84 @@
+"""Convert TaoTrain checkpoint + tokenizer assets into a standard HF package."""
+
+import json
+import shutil
+from pathlib import Path
+
+
+def normalize_checkpoint(checkpoint):
+    if isinstance(checkpoint, dict):
+        if "model_state" in checkpoint:
+            return checkpoint["model_state"], checkpoint.get("config", {})
+        if "model_state_dict" in checkpoint:
+            return checkpoint["model_state_dict"], checkpoint.get("config", {})
+    return checkpoint, {}
+
+
+def infer_special_token_paths(repo_dir):
+    subdir_metadata = repo_dir / "tokenizer" / "tokenizer.special_tokens.json"
+    root_metadata = repo_dir / "tokenizer.special_tokens.json"
+    subdir_model = repo_dir / "tokenizer" / "tokenizer.model"
+    subdir_vocab = repo_dir / "tokenizer" / "tokenizer.vocab"
+    if subdir_metadata.exists() and subdir_model.exists() and subdir_vocab.exists():
+        return subdir_model, subdir_metadata, subdir_vocab
+    if root_metadata.exists():
+        return repo_dir / "tokenizer.model", root_metadata, repo_dir / "tokenizer.vocab"
+    return repo_dir / "tokenizer" / "tokenizer.model", subdir_metadata, repo_dir / "tokenizer" / "tokenizer.vocab"
+
+
+def main():
+    import torch
+
+    from configuration_taonet import TaoNetConfig
+    from modeling_taonet import TaoNetForCausalLM
+    from tokenization_taonet import TaoNetTokenizer
+
+    repo_dir = Path(__file__).resolve().parent
+    checkpoint_path = repo_dir / "checkpoints" / "sft" / "final_model.pt"
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    model_state, train_config = normalize_checkpoint(checkpoint)
+    model_config = dict(train_config.get("model", {}))
+
+    metadata_model_path, metadata_path, vocab_path = infer_special_token_paths(repo_dir)
+    with open(metadata_path, "r", encoding="utf-8") as handle:
+        metadata = json.load(handle)
+
+    special_tokens = metadata.get("special_tokens", {})
+    hf_config = TaoNetConfig.from_taotrain_model_config(
+        model_config,
+        vocab_size=model_config.get("vocab_size", sum(1 for _ in open(vocab_path, "r", encoding="utf-8"))),
+        pad_token_id=special_tokens.get("<PAD>", 3),
+        bos_token_id=special_tokens.get("<BOS>", 1),
+        eos_token_id=special_tokens.get("<EOS>", 2),
+        unk_token_id=special_tokens.get("<UNK>", 0),
+    )
+
+    model = TaoNetForCausalLM(hf_config)
+    missing, unexpected = model.model.load_state_dict(model_state, strict=False)
+    model.tie_weights()
+    model.save_pretrained(repo_dir, safe_serialization=False)
+
+    tokenizer = TaoNetTokenizer(
+        vocab_file=str(metadata_model_path),
+        special_tokens_file=str(metadata_path),
+    )
+    tokenizer.save_pretrained(repo_dir)
+
+    shutil.copyfile(metadata_model_path, repo_dir / "tokenizer.model")
+    shutil.copyfile(metadata_path, repo_dir / "tokenizer.special_tokens.json")
+    shutil.copyfile(vocab_path, repo_dir / "tokenizer.vocab")
+
+    if missing:
+        print("Missing keys while loading model state:")
+        for key in missing:
+            print(f"  - {key}")
+    if unexpected:
+        print("Unexpected keys while loading model state:")
+        for key in unexpected:
+            print(f"  - {key}")
+    print(f"Saved Hugging Face package to: {repo_dir}")
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 3,
+  "transformers_version": "5.1.0"
+}

mla_components.py

@@ -0,0 +1,213 @@
+"""MLA attention blocks used by TaoNet."""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary position embedding with optional YaRN scaling."""
+
+    def __init__(
+        self,
+        dim,
+        rope_scale=40.0,
+        max_seq_length=1024,
+        yarn_enabled=False,
+        yarn_original_max_seq_length=None,
+        yarn_alpha=1.0,
+    ):
+        super().__init__()
+        if dim % 2 != 0:
+            raise ValueError("RotaryEmbedding requires an even dimension")
+        self.dim = dim
+        self.rope_scale = rope_scale
+        self.max_seq_length = max_seq_length
+        self.yarn_enabled = yarn_enabled
+        self.yarn_original_max_seq_length = (
+            yarn_original_max_seq_length if yarn_original_max_seq_length is not None else max_seq_length
+        )
+        self.yarn_alpha = yarn_alpha
+
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _apply_yarn_scaling(self, freqs, seq_len):
+        if not self.yarn_enabled:
+            return freqs
+        original_max_seq_length = self.yarn_original_max_seq_length
+        if seq_len <= original_max_seq_length:
+            return freqs
+        target_scale = self.max_seq_length / original_max_seq_length
+        current_ratio = seq_len / original_max_seq_length
+        progress = min(current_ratio / target_scale, 1.0)
+        scale_factor = 1.0 + (target_scale - 1.0) * (progress ** (1.0 / self.yarn_alpha))
+        return freqs / scale_factor
+
+    def forward(self, seq_len, device):
+        t = torch.arange(seq_len, device=device).type_as(self.inv_freq) / self.rope_scale
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        freqs = self._apply_yarn_scaling(freqs, seq_len)
+        return torch.cat((freqs, freqs), dim=-1)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary(x, cos, sin):
+    if cos.dim() == 2:
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+    cos = cos[..., : x.shape[-1]]
+    sin = sin[..., : x.shape[-1]]
+    x_rot = x[..., : cos.shape[-1]]
+    x_base = x[..., cos.shape[-1] :]
+    x_rot = (x_rot * cos) + (rotate_half(x_rot) * sin)
+    if x_base.shape[-1] > 0:
+        return torch.cat([x_rot, x_base], dim=-1)
+    return x_rot
+
+
+class DeepSeekMLA(nn.Module):
+    """DeepSeek-style multi-head latent attention."""
+
+    def __init__(
+        self,
+        d_model,
+        d_latent_kv,
+        n_heads,
+        d_rope,
+        dropout=0.1,
+        gqa_groups=1,
+        rope_scale=40.0,
+        max_seq_length=1024,
+        yarn_enabled=False,
+        yarn_original_max_seq_length=None,
+        yarn_alpha=1.0,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.d_latent_kv = d_latent_kv
+        self.n_heads = n_heads
+        self.d_rope = d_rope
+        self.gqa_groups = gqa_groups
+
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        if d_latent_kv % n_heads != 0:
+            raise ValueError("d_latent_kv must be divisible by n_heads")
+
+        self.d_head_full = d_model // n_heads
+        self.d_head_latent = d_latent_kv // n_heads
+
+        self.norm = nn.LayerNorm(d_model)
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_latent_kv, bias=False)
+        self.v_proj = nn.Linear(d_model, d_latent_kv, bias=False)
+        self.rotary = RotaryEmbedding(
+            d_rope,
+            rope_scale=rope_scale,
+            max_seq_length=max_seq_length,
+            yarn_enabled=yarn_enabled,
+            yarn_original_max_seq_length=yarn_original_max_seq_length,
+            yarn_alpha=yarn_alpha,
+        )
+        self.out_proj = nn.Linear(d_latent_kv, d_model, bias=False)
+        self.head_weights = nn.Parameter(torch.ones(n_heads))
+        self.attn_dropout = nn.Dropout(dropout)
+        self.proj_dropout = nn.Dropout(dropout)
+
+    def forward(self, x, attention_mask=None):
+        batch_size, seq_len, _ = x.shape
+        x_norm = self.norm(x)
+
+        q = self.q_proj(x_norm)
+        k = self.k_proj(x_norm)
+        v = self.v_proj(x_norm)
+
+        q = q.view(batch_size, seq_len, self.n_heads, self.d_head_full).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.n_heads, self.d_head_latent).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.n_heads, self.d_head_latent).transpose(1, 2)
+
+        if self.d_rope > 0:
+            rotary_emb = self.rotary(seq_len, x.device)
+            cos = torch.cos(rotary_emb).unsqueeze(0).unsqueeze(0)
+            sin = torch.sin(rotary_emb).unsqueeze(0).unsqueeze(0)
+            q_rope = apply_rotary(q[..., : self.d_rope], cos, sin)
+            q = torch.cat([q_rope, q[..., self.d_rope :]], dim=-1)
+            k_rope = apply_rotary(k[..., : self.d_rope], cos, sin)
+            k = torch.cat([k_rope, k[..., self.d_rope :]], dim=-1)
+
+        q_for_attn = q[..., : self.d_head_latent]
+        attn_mask_bool = None
+        if attention_mask is not None:
+            if attention_mask.dim() == 2:
+                attn_mask_bool = attention_mask.bool().unsqueeze(1).unsqueeze(1)
+            else:
+                attn_mask_bool = attention_mask.bool()
+
+        dropout_p = self.attn_dropout.p if self.training else 0.0
+        out_heads = F.scaled_dot_product_attention(
+            q_for_attn,
+            k,
+            v,
+            attn_mask=attn_mask_bool,
+            dropout_p=dropout_p,
+            scale=None,
+        )
+        out_concat = out_heads.transpose(1, 2).reshape(batch_size, seq_len, self.d_latent_kv)
+        out = self.out_proj(out_concat)
+        return self.proj_dropout(out)
+
+
+class AttentionBlock(nn.Module):
+    """Attention block with SwiGLU feed-forward network."""
+
+    def __init__(
+        self,
+        d_model,
+        d_latent_kv,
+        n_heads,
+        d_rope,
+        d_ff,
+        dropout=0.1,
+        gqa_groups=1,
+        rope_scale=40.0,
+        max_seq_length=1024,
+        yarn_enabled=False,
+        yarn_original_max_seq_length=None,
+        yarn_alpha=1.0,
+    ):
+        super().__init__()
+        self.mla = DeepSeekMLA(
+            d_model,
+            d_latent_kv,
+            n_heads,
+            d_rope,
+            dropout,
+            gqa_groups,
+            rope_scale=rope_scale,
+            max_seq_length=max_seq_length,
+            yarn_enabled=yarn_enabled,
+            yarn_original_max_seq_length=yarn_original_max_seq_length,
+            yarn_alpha=yarn_alpha,
+        )
+        self.ff_norm = nn.LayerNorm(d_model)
+        self.ff_gate = nn.Linear(d_model, d_ff, bias=False)
+        self.ff_value = nn.Linear(d_model, d_ff, bias=False)
+        self.ff_out = nn.Linear(d_ff, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, attention_mask=None):
+        attn_out = self.mla(x, attention_mask)
+        x = x + self.dropout(attn_out)
+        ff_norm = self.ff_norm(x)
+        ff_gate = self.ff_gate(ff_norm)
+        ff_value = self.ff_value(ff_norm)
+        ff_out = ff_value * F.silu(ff_gate)
+        ff_out = self.ff_out(ff_out)
+        return x + self.dropout(ff_out)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dcf558286bb14902875884b2a18858a2cab7d0341ced85297814b822a369c77
+size 859468256

modeling_taonet.py

@@ -0,0 +1,86 @@
+"""Hugging Face model wrapper for TaoNet."""
+
+from torch import nn
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+from configuration_taonet import TaoNetConfig
+from taonet_model import SimpleLLM, build_runtime_config
+
+
+class TaoNetForCausalLM(PreTrainedModel, GenerationMixin):
+    """Transformers-compatible TaoNet causal LM."""
+
+    config_class = TaoNetConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        runtime_config = build_runtime_config(config)
+        self.model = SimpleLLM(runtime_config)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self):
+        if getattr(self.model, "use_factorized_embedding", False):
+            return self.model.token_embedding.embed
+        return self.model.token_embedding
+
+    def set_input_embeddings(self, value):
+        if getattr(self.model, "use_factorized_embedding", False):
+            self.model.token_embedding.embed = value
+        else:
+            self.model.token_embedding = value
+
+    def get_output_embeddings(self):
+        return self.model.output_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.model.output_head = new_embeddings
+
+    def tie_weights(self, *args, **kwargs):
+        del args, kwargs
+        if not getattr(self.model, "use_factorized_embedding", False):
+            self.model.output_head.weight = self.get_input_embeddings().weight
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        inputs_embeds=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        del kwargs
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,
+            inputs_embeds=inputs_embeds,
+        )
+        loss = None
+        if labels is not None:
+            shift_logits = outputs["logits"][..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+        if not return_dict:
+            return (loss, outputs["logits"])
+        return CausalLMOutput(loss=loss, logits=outputs["logits"])
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids, "attention_mask": attention_mask}

special_tokens_map.json

@@ -0,0 +1,13 @@
+{
+  "unk_token": "<UNK>",
+  "bos_token": "<BOS>",
+  "eos_token": "<EOS>",
+  "pad_token": "<PAD>",
+  "additional_special_tokens": [
+    "\n",
+    "<think>",
+    "<user>",
+    "<assistant>",
+    "<image>"
+  ]
+}

src/taoTrain/__init__.py

@@ -0,0 +1,8 @@
+"""
+TaoTrain: A clean, modular PyTorch LLM training framework.
+
+For quick usage, see: https://github.com/your-org/taoTrain
+"""
+
+__version__ = "0.1.0"
+__author__ = "Felix"

src/taoTrain/benchmarks/__init__.py

@@ -0,0 +1,5 @@
+"""Benchmarking suite."""
+
+from .runner import BenchmarkRunner
+
+__all__ = ["BenchmarkRunner"]

src/taoTrain/benchmarks/runner.py

@@ -0,0 +1,221 @@
+"""Benchmarking suite for evaluating trained models."""
+
+import time
+from pathlib import Path
+from typing import Optional, Dict
+import torch
+from torch.utils.data import DataLoader
+
+from taoTrain.core import BaseModel
+from taoTrain.config import TrainingConfig
+from taoTrain.data.loaders import get_dataloader
+from taoTrain.inference import Inferencer
+
+
+class BenchmarkRunner:
+    """Run benchmarks on a trained model."""
+    
+    def __init__(
+        self,
+        model: BaseModel,
+        device: torch.device,
+        dtype: torch.dtype = torch.float32,
+    ):
+        """
+        Initialize benchmark runner.
+        
+        Args:
+            model: Trained model
+            device: Device for inference
+            dtype: Data type
+        """
+        self.model = model.to(device)
+        self.model.eval()
+        self.device = device
+        self.dtype = dtype
+    
+    @staticmethod
+    def load_from_checkpoint(
+        checkpoint_path: str | Path,
+        device: Optional[torch.device] = None,
+    ) -> "BenchmarkRunner":
+        """Load model from checkpoint."""
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        
+        # Reconstruct model config
+        from taoTrain.config import ModelConfig
+        from taoTrain.models import get_model
+        
+        model_config = ModelConfig(**checkpoint.get("config", {}).get("model", {}))
+        model = get_model(model_config, device=device)
+        model.load_state_dict(checkpoint["model_state_dict"])
+        
+        return BenchmarkRunner(model, device)
+    
+    def benchmark_perplexity(
+        self,
+        dataset: "DataLoader",
+        num_batches: Optional[int] = None,
+    ) -> float:
+        """
+        Compute perplexity on a dataset.
+        
+        Args:
+            dataset: DataLoader for evaluation
+            num_batches: Limit evaluation to N batches
+        
+        Returns:
+            Perplexity (exp of average loss)
+        """
+        total_loss = 0.0
+        total_tokens = 0
+        
+        with torch.no_grad():
+            for batch_idx, batch in enumerate(dataset):
+                if num_batches and batch_idx >= num_batches:
+                    break
+                
+                # Move to device
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch.get("attention_mask")
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(self.device)
+                labels = batch.get("labels")
+                if labels is not None:
+                    labels = labels.to(self.device)
+                
+                # Forward pass
+                with torch.autocast(
+                    device_type="cuda" if self.device.type == "cuda" else "cpu",
+                    dtype=torch.bfloat16 if self.dtype == torch.bfloat16 else torch.float32,
+                ):
+                    outputs = self.model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        labels=labels,
+                    )
+                    loss = outputs.get("loss")
+                
+                if loss is not None:
+                    total_loss += loss.item() * input_ids.shape[0]
+                    total_tokens += input_ids.shape[0]
+        
+        avg_loss = total_loss / total_tokens if total_tokens > 0 else float('inf')
+        perplexity = torch.exp(torch.tensor(avg_loss)).item()
+        
+        return perplexity
+    
+    def benchmark_throughput(
+        self,
+        batch_size: int = 32,
+        seq_length: int = 1024,
+        num_iters: int = 10,
+    ) -> Dict[str, float]:
+        """
+        Benchmark forward pass throughput.
+        
+        Args:
+            batch_size: Batch size
+            seq_length: Sequence length
+            num_iters: Number of iterations
+        
+        Returns:
+            Dict with throughput metrics
+        """
+        # Create dummy batch
+        dummy_input = torch.randint(
+            0, self.model.config.vocab_size,
+            (batch_size, seq_length)
+        ).to(self.device)
+        
+        # Warmup
+        with torch.no_grad():
+            for _ in range(2):
+                _ = self.model(dummy_input)
+        
+        torch.cuda.synchronize() if torch.cuda.is_available() else None
+        
+        # Benchmark forward pass
+        start = time.time()
+        
+        with torch.no_grad():
+            for _ in range(num_iters):
+                _ = self.model(dummy_input)
+        
+        torch.cuda.synchronize() if torch.cuda.is_available() else None
+        
+        elapsed = time.time() - start
+        
+        total_tokens = batch_size * seq_length * num_iters
+        tokens_per_sec = total_tokens / elapsed
+        
+        return {
+            "throughput_tokens_per_sec": tokens_per_sec,
+            "throughput_samples_per_sec": (batch_size * num_iters) / elapsed,
+            "avg_time_per_iter_ms": (elapsed / num_iters) * 1000,
+        }
+    
+    def benchmark_memory(self) -> Dict[str, float]:
+        """
+        Benchmark peak GPU memory usage.
+        
+        Returns:
+            Dict with memory stats
+        """
+        if not torch.cuda.is_available():
+            return {"peak_memory_gb": 0.0}
+        
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.synchronize()
+        
+        # Create dummy batch
+        dummy_input = torch.randint(
+            0, self.model.config.vocab_size,
+            (16, 1024)
+        ).to(self.device)
+        
+        with torch.no_grad():
+            _ = self.model(dummy_input)
+        
+        torch.cuda.synchronize()
+        
+        peak_memory = torch.cuda.max_memory_allocated() / (1024 ** 3)  # GB
+        
+        return {"peak_memory_gb": peak_memory}
+    
+    def run_all_benchmarks(
+        self,
+        dataset: Optional["DataLoader"] = None,
+        batch_size: int = 32,
+        seq_length: int = 1024,
+    ) -> Dict[str, float]:
+        """
+        Run all benchmarks.
+        
+        Args:
+            dataset: DataLoader for perplexity benchmark
+            batch_size: Batch size for throughput benchmark
+            seq_length: Sequence length for throughput benchmark
+        
+        Returns:
+            Dict with all benchmark results
+        """
+        results = {}
+        
+        if dataset is not None:
+            print("Running perplexity benchmark...")
+            ppl = self.benchmark_perplexity(dataset, num_batches=10)
+            results["perplexity"] = ppl
+        
+        print("Running throughput benchmark...")
+        throughput = self.benchmark_throughput(batch_size, seq_length)
+        results.update(throughput)
+        
+        print("Running memory benchmark...")
+        memory = self.benchmark_memory()
+        results.update(memory)
+        
+        return results

src/taoTrain/checkpointing/__init__.py

@@ -0,0 +1,5 @@
+"""Checkpoint management."""
+
+from .checkpoint import CheckpointManager
+
+__all__ = ["CheckpointManager"]

src/taoTrain/checkpointing/checkpoint.py

@@ -0,0 +1,194 @@
+"""Checkpoint management utilities.
+
+Canonical Checkpoint Format (new):
+    {
+        'step': int,                          # Training step number
+        'model_state': Dict[str, Tensor],     # Model state dict
+        'optimizer_state': Dict,              # Optimizer state dict (optional)
+        'config': Dict,                       # TrainingConfig as dict
+        'metrics': Dict[str, float],          # Training metrics
+        'global_step': int,                   # (deprecated, kept for compat) same as step
+        'current_epoch': int,                 # (optional) current epoch number
+        'best_loss': float,                   # (optional) best validation loss
+    }
+
+Legacy Checkpoint Format (old, from BaseTrainer):
+    {
+        'global_step': int,
+        'current_epoch': int,
+        'best_loss': float,
+        'model_state_dict': Dict[str, Tensor],    # ← Note: uses '_dict' suffix
+        'optimizer_state_dict': Dict,
+        'config': Dict,
+    }
+
+The load() function auto-detects and migrates legacy format to canonical format.
+"""
+
+from pathlib import Path
+from typing import Dict, Any, Optional
+import torch
+from taoTrain.config import TrainingConfig
+
+
+class CheckpointManager:
+    """Manage model checkpoints with versioning."""
+    
+    def __init__(
+        self,
+        checkpoint_dir: str | Path,
+        keep_last_n: int = 3,
+        track_best: bool = True,
+    ):
+        """
+        Initialize checkpoint manager.
+        
+        Args:
+            checkpoint_dir: Directory to save checkpoints
+            keep_last_n: Number of recent checkpoints to keep
+            track_best: Whether to track best model
+        """
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.keep_last_n = keep_last_n
+        self.track_best = track_best
+        
+        self.best_metric = None
+        self.best_metric_name = None
+        self.saved_checkpoints = []
+    
+    def save(
+        self,
+        step: int,
+        model_state: Dict[str, Any],
+        optimizer_state: Optional[Dict[str, Any]] = None,
+        config: Optional[TrainingConfig] = None,
+        metrics: Optional[Dict[str, float]] = None,
+        is_best: bool = False,
+    ) -> Path:
+        """
+        Save a checkpoint.
+        
+        Args:
+            step: Training step
+            model_state: Model state dict
+            optimizer_state: Optimizer state dict
+            config: Training config
+            metrics: Metrics dict
+            is_best: Whether this is the best model so far
+        
+        Returns:
+            Path to saved checkpoint
+        """
+        checkpoint = {
+            "step": step,
+            "model_state": model_state,
+            "optimizer_state": optimizer_state,
+            "config": config.to_dict() if config else None,
+            "metrics": metrics or {},
+        }
+        
+        filename = f"checkpoint_step_{step:06d}.pt"
+        if is_best:
+            filename = "best_model.pt"
+        
+        path = self.checkpoint_dir / filename
+        torch.save(checkpoint, path)
+        
+        # Track saved checkpoints
+        if not is_best:
+            self.saved_checkpoints.append((step, path))
+            
+            # Clean up old checkpoints
+            if len(self.saved_checkpoints) > self.keep_last_n:
+                _, old_path = self.saved_checkpoints.pop(0)
+                if old_path.exists():
+                    old_path.unlink()
+        
+        return path
+    
+    def load(
+        self,
+        checkpoint_path: str | Path,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, Any]:
+        """
+        Load a checkpoint with backward-compatible format handling.
+        
+        Auto-detects checkpoint format (canonical or legacy) and normalizes
+        to canonical format in-memory. Legacy checkpoints are migrated without
+        modifying the file.
+        
+        Args:
+            checkpoint_path: Path to checkpoint
+            device: Device to load to
+        
+        Returns:
+            Checkpoint dict in canonical format with 'model_state' key
+        """
+        if device is None:
+            device = torch.device("cpu")
+        
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        
+        # Auto-detect and migrate legacy format to canonical format
+        checkpoint = self._normalize_checkpoint_format(checkpoint)
+        
+        return checkpoint
+    
+    def _normalize_checkpoint_format(self, checkpoint: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Normalize checkpoint to canonical format.
+        
+        Detects if checkpoint is in legacy format (from BaseTrainer with 'model_state_dict')
+        and migrates it to canonical format (with 'model_state').
+        
+        Args:
+            checkpoint: Raw checkpoint dict
+        
+        Returns:
+            Normalized checkpoint dict with canonical keys
+        """
+        # Check if this is a legacy checkpoint (has 'model_state_dict' but not 'model_state')
+        if "model_state_dict" in checkpoint and "model_state" not in checkpoint:
+            # Migrate legacy format to canonical
+            migrated = {
+                "step": checkpoint.get("global_step", 0),
+                "model_state": checkpoint["model_state_dict"],
+                "optimizer_state": checkpoint.get("optimizer_state_dict"),
+                "config": checkpoint.get("config"),
+                "metrics": {},
+                # Keep legacy keys for backward compatibility in code that uses them
+                "global_step": checkpoint.get("global_step", 0),
+                "current_epoch": checkpoint.get("current_epoch", 0),
+                "best_loss": checkpoint.get("best_loss", float('inf')),
+            }
+            print(f"\n✓ [CheckpointManager] Detected legacy checkpoint format. Auto-migrated to canonical format.")
+            return migrated
+        
+        # Already in canonical format or unknown format
+        if "model_state" not in checkpoint:
+            # If neither format detected, ensure model_state is accessible
+            # (might be a raw state_dict)
+            print(f"\n⚠ [CheckpointManager] Checkpoint format unclear. Assuming raw state_dict format.")
+            checkpoint["model_state"] = checkpoint
+        
+        return checkpoint
+    
+    def get_latest(self) -> Optional[Path]:
+        """Get path to latest checkpoint."""
+        if not self.saved_checkpoints:
+            return None
+        return self.saved_checkpoints[-1][1]
+    
+    def get_best(self) -> Optional[Path]:
+        """Get path to best checkpoint."""
+        best_path = self.checkpoint_dir / "best_model.pt"
+        if best_path.exists():
+            return best_path
+        return None
+    
+    def list_checkpoints(self) -> list[Path]:
+        """List all saved checkpoints."""
+        return sorted(self.checkpoint_dir.glob("checkpoint_step_*.pt"))

src/taoTrain/cli.py

@@ -0,0 +1,470 @@
+"""Main CLI entry point."""
+
+import sys
+from pathlib import Path
+from typing import Optional
+import click
+import torch
+
+from taoTrain.config import (
+    load_config,
+    load_tokenizer_config,
+    TrainingModeEnum,
+    TrainingConfig,
+    PretrainConfig,
+    SFTConfig,
+    RLConfig,
+    VLMConfig,
+    VLMSFTConfig,
+)
+from taoTrain.utils import set_seed, get_device
+from taoTrain.core import BaseModel, create_model, create_datasets
+from taoTrain.data import get_dataloader
+from taoTrain.training import PretrainTrainer, SFTTrainer, RLTrainer, VLMTrainer, VLMSFTTrainer
+from taoTrain.benchmarks import BenchmarkRunner
+from taoTrain.tokenizers import TokenizerTrainer
+from taoTrain.data.tokenizer import load_tokenizer_runtime_info
+
+
+@click.group()
+def main():
+    """TaoTrain: A clean, modular PyTorch LLM training framework."""
+    pass
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to training config file (YAML or JSON)",
+)
+def pretrain(config: str):
+    """Pretrain a language model."""
+    _train_command(config, TrainingModeEnum.PRETRAIN)
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to training config file (YAML or JSON)",
+)
+def sft(config: str):
+    """Supervised fine-tune a language model."""
+    _train_command(config, TrainingModeEnum.SFT)
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to training config file (YAML or JSON)",
+)
+def rl(config: str):
+    """Train with reinforcement learning."""
+    _train_command(config, TrainingModeEnum.RL)
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to multimodal VLM config file (YAML or JSON)",
+)
+def vlm(config: str):
+    """Train a multimodal vision-language connector."""
+    _train_command(config, TrainingModeEnum.VLM)
+
+
+@main.command(name="vlm-sft")
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to multimodal VLM SFT config file (YAML or JSON)",
+)
+def vlm_sft(config: str):
+    """Run end-to-end multimodal supervised fine-tuning."""
+    _train_command(config, TrainingModeEnum.VLM_SFT)
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to tokenizer config file (YAML or JSON)",
+)
+def train_tokenizer(config: str):
+    """Train a SentencePiece tokenizer from a YAML/JSON config file."""
+    try:
+        click.echo("🚀 TaoTrain Tokenizer Trainer")
+        click.echo(f"{'=' * 50}")
+        
+        # Load tokenizer config
+        click.echo(f"Loading config from {config}...")
+        tokenizer_config = load_tokenizer_config(config)
+        
+        # Train tokenizer from config
+        result = TokenizerTrainer.train_from_config(tokenizer_config)
+        
+        # Display results
+        click.echo(f"\n{'=' * 50}")
+        click.echo("✅ Tokenizer Training Complete!")
+        click.echo(f"\n📊 Configuration:")
+        click.echo(f"   - Input file: {tokenizer_config.jsonl_path}")
+        click.echo(f"   - Samples: {tokenizer_config.max_samples or 'all'}")
+        click.echo(f"   - Output dir: {result['output_dir']}")
+        click.echo(f"   - Vocab size: {result['vocab_size']}")
+        click.echo(f"   - Model type: {result['model_type']}")
+        
+        if tokenizer_config.special_tokens:
+            click.echo(f"   - Special tokens: {tokenizer_config.special_tokens}")
+        
+        click.echo(f"\n📁 Generated Files:")
+        click.echo(f"   - Model: {result['model_file']}")
+        click.echo(f"   - Vocab: {result['vocab_file']}")
+        if "special_tokens_file" in result:
+            click.echo(f"   - Special token metadata: {result['special_tokens_file']}")
+        
+        click.echo(f"\n📝 Next Steps:")
+        click.echo(f"   1. Use this tokenizer in your pretraining config:")
+        click.echo(f"       dataset:")
+        click.echo(f"         local: true")
+        click.echo(f"         jsonl_path: {tokenizer_config.jsonl_path}")
+        click.echo(f"         tokenizer_path: {result['model_file']}")
+        click.echo(f"")
+        click.echo(f"   2. Run pretraining with:")
+        click.echo(f"       train pretrain --config your_config.yaml")
+        
+    except ImportError as e:
+        click.echo(f"❌ Error: {e}", err=True)
+        sys.exit(1)
+    except FileNotFoundError as e:
+        click.echo(f"❌ File Error: {e}", err=True)
+        sys.exit(1)
+    except ValueError as e:
+        click.echo(f"❌ Validation Error: {e}", err=True)
+        sys.exit(1)
+    except Exception as e:
+        click.echo(f"❌ Unexpected Error: {e}", err=True)
+        click.echo(f"   Please report this issue.", err=True)
+        sys.exit(1)
+
+
+@click.command()
+@click.option(
+    "--jsonl-path",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to JSONL file containing training data",
+)
+@click.option(
+    "--output-dir",
+    type=click.Path(),
+    default="tokenizers",
+    help="Directory to save tokenizer files",
+)
+@click.option(
+    "--vocab-size",
+    type=int,
+    default=50000,
+    help="Vocabulary size for the tokenizer",
+)
+@click.option(
+    "--model-type",
+    type=click.Choice(["unigram", "bpe", "char", "word"]),
+    default="unigram",
+    help="SentencePiece model type",
+)
+@click.option(
+    "--character-coverage",
+    type=float,
+    default=0.9995,
+    help="Character coverage for SentencePiece",
+)
+@click.option(
+    "--tokenizer-prefix",
+    type=str,
+    default=None,
+    help="Prefix for tokenizer output files (default: model_type)",
+)
+def train_tokenizer_command(
+    jsonl_path: str,
+    output_dir: str,
+    vocab_size: int,
+    model_type: str,
+    character_coverage: float,
+    tokenizer_prefix: Optional[str],
+):
+    """Train a SentencePiece tokenizer from JSONL data."""
+    try:
+        click.echo("🚀 TaoTrain Tokenizer Trainer")
+        click.echo(f"{'=' * 50}")
+        
+        # Train tokenizer
+        result = TokenizerTrainer.train_sentencepiece(
+            jsonl_path=jsonl_path,
+            output_dir=output_dir,
+            vocab_size=vocab_size,
+            model_type=model_type,
+            character_coverage=character_coverage,
+            tokenizer_prefix=tokenizer_prefix,
+        )
+        
+        # Display results
+        click.echo(f"\n{'=' * 50}")
+        click.echo("✅ Tokenizer Training Complete!")
+        click.echo(f"\n📊 Configuration:")
+        click.echo(f"   - Input file: {jsonl_path}")
+        click.echo(f"   - Output dir: {result['output_dir']}")
+        click.echo(f"   - Vocab size: {result['vocab_size']}")
+        click.echo(f"   - Model type: {result['model_type']}")
+        
+        click.echo(f"\n📁 Generated Files:")
+        click.echo(f"   - Model: {result['model_file']}")
+        click.echo(f"   - Vocab: {result['vocab_file']}")
+        if "special_tokens_file" in result:
+            click.echo(f"   - Special token metadata: {result['special_tokens_file']}")
+        
+        click.echo(f"\n📝 Next Steps:")
+        click.echo(f"   1. Use this tokenizer in your pretraining config:")
+        click.echo(f"       dataset:")
+        click.echo(f"         local: true")
+        click.echo(f"         jsonl_path: {jsonl_path}")
+        click.echo(f"         tokenizer_path: {result['model_file']}")
+        click.echo(f"")
+        click.echo(f"   2. Run pretraining with:")
+        click.echo(f"       train pretrain --config your_config.yaml")
+        
+    except ImportError as e:
+        click.echo(f"❌ Error: {e}", err=True)
+        sys.exit(1)
+    except FileNotFoundError as e:
+        click.echo(f"❌ File Error: {e}", err=True)
+        sys.exit(1)
+    except ValueError as e:
+        click.echo(f"❌ Validation Error: {e}", err=True)
+        sys.exit(1)
+    except Exception as e:
+        click.echo(f"❌ Unexpected Error: {e}", err=True)
+        click.echo(f"   Please report this issue.", err=True)
+        sys.exit(1)
+
+
+# Keep legacy CLI command as train-tokenizer-legacy for backward compatibility
+main.add_command(train_tokenizer_command, name="train-tokenizer-legacy")
+
+
+def _train_command(config_path: str, mode: TrainingModeEnum):
+    """Internal training command."""
+    try:
+        # Load config
+        click.echo(f"Loading config from {config_path}...")
+        train_config = load_config(config_path, mode)
+
+        _synchronize_model_vocab_with_tokenizer(train_config)
+        
+        # Set seed
+        set_seed(train_config.seed)
+        
+        # Get device
+        device = get_device(train_config.device)
+        click.echo(f"Using device: {device}")
+        
+        # Create model
+        click.echo("Creating model...")
+        model = create_model(train_config, device)
+        total_params, trainable_params = _count_params(model)
+        click.echo(f"  - Total parameters: {total_params:,}")
+        click.echo(f"  - Trainable parameters: {trainable_params:,}")
+        
+        # Load pretrained checkpoint if provided (for SFT/RL)
+        if train_config.checkpoint_path:
+            click.echo(f"Loading pretrained checkpoint from {train_config.checkpoint_path}...")
+            from taoTrain.checkpointing.checkpoint import CheckpointManager
+            checkpoint_manager = CheckpointManager(train_config.checkpoint_dir)
+            checkpoint = checkpoint_manager.load(train_config.checkpoint_path, device=device)
+            
+            # CheckpointManager.load() normalizes format and ensures 'model_state' key exists
+            if "model_state" in checkpoint:
+                load_result = model.load_state_dict(checkpoint["model_state"], strict=False)
+                click.echo("  ✓ Checkpoint loaded successfully")
+                if load_result.missing_keys:
+                    click.echo(f"  Missing keys: {len(load_result.missing_keys)}")
+                    for key in load_result.missing_keys[:10]:
+                        click.echo(f"      - {key}")
+                if load_result.unexpected_keys:
+                    click.echo(f"  Unexpected keys: {len(load_result.unexpected_keys)}")
+                    for key in load_result.unexpected_keys[:10]:
+                        click.echo(f"      - {key}")
+            else:
+                raise KeyError(f"Invalid checkpoint format: 'model_state' key not found. "
+                             f"Available keys: {list(checkpoint.keys())}")
+
+
+        
+        # Create datasets
+        click.echo("Loading datasets...")
+        train_dataset, val_dataset = create_datasets(train_config)
+        click.echo(f"  - Train samples: {len(train_dataset)}")
+        if val_dataset:
+            click.echo(f"  - Val samples: {len(val_dataset)}")
+        
+        # Select trainer
+        if mode == TrainingModeEnum.PRETRAIN:
+            trainer_class = PretrainTrainer
+        elif mode == TrainingModeEnum.SFT:
+            trainer_class = SFTTrainer
+        elif mode == TrainingModeEnum.RL:
+            trainer_class = RLTrainer
+        elif mode == TrainingModeEnum.VLM:
+            trainer_class = VLMTrainer
+        elif mode == TrainingModeEnum.VLM_SFT:
+            trainer_class = VLMSFTTrainer
+        else:
+            raise ValueError(f"Unknown training mode: {mode}")
+        
+        # Create trainer
+        click.echo("Setting up trainer...")
+        trainer = trainer_class(
+            model=model,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            config=train_config,
+            device=device,
+        )
+        
+        # Training loop
+        click.echo("\nStarting training...\n")
+        for epoch in range(train_config.num_epochs):
+            if train_config.max_steps and trainer.global_step >= train_config.max_steps:
+                break
+            
+            epoch_metrics = trainer.train_epoch()
+            click.echo(f"\nEpoch {epoch + 1} complete")
+            click.echo(f"  - Loss: {epoch_metrics.get('loss', 'N/A')}")
+            click.echo(f"  - Learning rate: {epoch_metrics.get('lr', 'N/A')}")
+        
+        # Final checkpoint
+        final_path = Path(train_config.checkpoint_dir) / "final_model.pt"
+        trainer.save_checkpoint(final_path)
+        click.echo(f"\nTraining complete! Final model saved to {final_path}")
+        
+        # Log finish
+        trainer.logger.finish()
+        
+    except Exception as e:
+        click.echo(f"Error during training: {e}", err=True)
+        sys.exit(1)
+
+
+def _synchronize_model_vocab_with_tokenizer(train_config: TrainingConfig) -> None:
+    """Validate tokenizer runtime IDs and align model vocab size with tokenizer artifacts."""
+    dataset_config = train_config.dataset
+    tokenizer_path = dataset_config.tokenizer_path
+    if not tokenizer_path:
+        return
+
+    tokenizer_info = load_tokenizer_runtime_info(tokenizer_path, dataset_config.tokenizer_type)
+    invalid_builtin_ids = {
+        token_name: token_id
+        for token_name, token_id in tokenizer_info.items()
+        if token_name.endswith("_id") and token_id < 0
+    }
+    if invalid_builtin_ids:
+        raise ValueError(
+            f"Tokenizer `{tokenizer_path}` is missing required built-in IDs: {invalid_builtin_ids}. "
+            "Retrain the tokenizer so PAD/BOS/EOS/UNK are real tokenizer IDs before training."
+        )
+
+    tokenizer_vocab_size = tokenizer_info["vocab_size"]
+    model_vocab_size = train_config.model.vocab_size
+    if model_vocab_size != tokenizer_vocab_size:
+        click.echo(
+            "Tokenizer/model vocab mismatch detected. "
+            f"Updating model.vocab_size from {model_vocab_size} to {tokenizer_vocab_size} "
+            f"based on tokenizer `{tokenizer_path}`."
+        )
+        train_config.model.vocab_size = tokenizer_vocab_size
+
+
+@main.command()
+@click.option(
+    "--model",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to model checkpoint",
+)
+@click.option(
+    "--benchmark-type",
+    type=click.Choice(["all", "perplexity", "throughput", "memory"]),
+    default="all",
+    help="Type of benchmark to run",
+)
+@click.option(
+    "--batch-size",
+    type=int,
+    default=32,
+    help="Batch size for benchmarking",
+)
+@click.option(
+    "--seq-length",
+    type=int,
+    default=1024,
+    help="Sequence length for benchmarking",
+)
+def benchmark(model: str, benchmark_type: str, batch_size: int, seq_length: int):
+    """Benchmark a trained model."""
+    try:
+        click.echo(f"Loading model from {model}...")
+        device = get_device("cuda")
+        runner = BenchmarkRunner.load_from_checkpoint(model, device=device)
+        
+        click.echo("Running benchmarks...\n")
+        
+        if benchmark_type == "throughput" or benchmark_type == "all":
+            click.echo("Throughput benchmark:")
+            results = runner.benchmark_throughput(batch_size, seq_length)
+            for key, val in results.items():
+                click.echo(f"  {key}: {val:.2f}")
+        
+        if benchmark_type == "memory" or benchmark_type == "all":
+            click.echo("\nMemory benchmark:")
+            results = runner.benchmark_memory()
+            for key, val in results.items():
+                click.echo(f"  {key}: {val:.2f}")
+        
+        click.echo("\nBenchmarking complete!")
+        
+    except Exception as e:
+        click.echo(f"Error during benchmarking: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.option("--repo", type=str, default=".aim", help="AimStack repository path")
+def view_logs(repo: str):
+    """View training logs with AimStack."""
+    try:
+        import subprocess
+        click.echo(f"Opening AimStack dashboard for repo: {repo}")
+        subprocess.run(["aim", "up", "--repo", repo])
+    except FileNotFoundError:
+        click.echo("Error: 'aim' command not found. Install with: pip install aim", err=True)
+        sys.exit(1)
+
+
+def _count_params(model: BaseModel) -> tuple[int, int]:
+    """Count model parameters."""
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total, trainable
+
+
+if __name__ == "__main__":
+    main()

src/taoTrain/config.py

@@ -0,0 +1,871 @@
+"""Pydantic configuration schemas for TaoTrain."""
+
+from enum import Enum
+from typing import Optional, Literal
+from pathlib import Path
+import json
+from pydantic import BaseModel as PydanticBaseModel, Field, validator
+import yaml
+
+
+# ============================================================================
+# Enums
+# ============================================================================
+
+
+class DataTypeEnum(str, Enum):
+    """Data types for training."""
+    FLOAT32 = "float32"
+    FLOAT16 = "float16"
+    BFLOAT16 = "bfloat16"
+
+
+class OptimizerEnum(str, Enum):
+    """Supported optimizers."""
+    ADAM = "adam"
+    ADAMW = "adamw"
+    SGD = "sgd"
+    HYBRID_MUON_ADAMW = "hybrid_muon_adamw"
+
+
+class ModelArchitectureEnum(str, Enum):
+    """Built-in model architectures."""
+    TRANSFORMER = "transformer"
+    TAONET = "taonet"
+    GAMMA_NET = "gamma_net"
+    MULTIMODAL_WRAPPER = "multimodal_wrapper"
+
+
+class SchedulerEnum(str, Enum):
+    """Supported learning rate schedulers."""
+    LINEAR_WARMUP = "linearWarmup"
+    COSINE_WARMUP = "cosineWarmup"
+    CONSTANT = "constant"
+
+
+class RLMethodEnum(str, Enum):
+    """Supported RL training methods."""
+    PPO = "ppo"
+    DPO = "dpo"
+
+
+class TrainingModeEnum(str, Enum):
+    """Training stages."""
+    PRETRAIN = "pretrain"
+    SFT = "sft"
+    RL = "rl"
+    VLM = "vlm"
+    VLM_SFT = "vlm_sft"
+
+
+# ============================================================================
+# Base Configs
+# ============================================================================
+
+
+class BaseConfig(PydanticBaseModel):
+    """Base Pydantic model with utility methods."""
+    
+    class Config:
+        """Pydantic config."""
+        arbitrary_types_allowed = True
+    
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        data = self.model_dump(mode='json')  # Enums -> strings
+        return data
+    
+    def to_json_str(self) -> str:
+        """Convert to JSON string."""
+        return json.dumps(self.to_dict(), indent=2)
+    
+    def save_yaml(self, path: str | Path) -> None:
+        """Save config to YAML file."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, 'w') as f:
+            yaml.dump(self.to_dict(), f, default_flow_style=False, sort_keys=False)
+    
+    def save_json(self, path: str | Path) -> None:
+        """Save config to JSON file."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, 'w') as f:
+            f.write(self.to_json_str())
+    
+    @classmethod
+    def load_yaml(cls, path: str | Path) -> "BaseConfig":
+        """Load config from YAML file."""
+        with open(path) as f:
+            data = yaml.safe_load(f)
+        return cls(**data)
+    
+    @classmethod
+    def load_json(cls, path: str | Path) -> "BaseConfig":
+        """Load config from JSON file."""
+        with open(path) as f:
+            data = json.load(f)
+        return cls(**data)
+
+
+# ============================================================================
+# Model Config
+# ============================================================================
+
+
+class ModelConfig(BaseConfig):
+    """Configuration for model architecture."""
+    
+    architecture_type: ModelArchitectureEnum = Field(
+        default=ModelArchitectureEnum.TRANSFORMER,
+        description="Type of model architecture"
+    )
+    llm_architecture_type: Optional[ModelArchitectureEnum] = Field(
+        default=None,
+        description="Base language-model architecture when using a multimodal wrapper."
+    )
+    
+    # Transformer-specific
+    vocab_size: int = Field(default=50257, description="Vocabulary size")
+    hidden_dim: int = Field(default=768, description="Hidden dimension")
+    num_layers: int = Field(default=12, description="Number of transformer blocks")
+    num_heads: int = Field(default=12, description="Number of attention heads")
+    head_dim: Optional[int] = Field(
+        default=None,
+        description="Head dimension (defaults to hidden_dim // num_heads)"
+    )
+    intermediate_dim: Optional[int] = Field(
+        default=None,
+        description="FFN intermediate dimension (defaults to 4 * hidden_dim)"
+    )
+    dropout: float = Field(default=0.1, description="Dropout rate")
+    max_seq_length: int = Field(default=2048, description="Maximum sequence length")
+    
+    # TaoNet (DeepSeek MLA) specific
+    d_latent_kv: Optional[int] = Field(
+        default=None,
+        description="KV compression dimension for MLA (defaults to 3/4 * hidden_dim). Only used for taonet architecture."
+    )
+    d_rope: Optional[int] = Field(
+        default=None,
+        description="RoPE dimension per head (defaults to hidden_dim // num_heads). Only used for taonet architecture."
+    )
+    gqa_groups: int = Field(
+        default=1,
+        description="Grouped Query Attention groups (1 = standard MLA, >1 = GQA). Only used for taonet architecture."
+    )
+    hidden_dim_ff: Optional[int] = Field(
+        default=None,
+        description="Feed-forward intermediate dimension (defaults to 4 * hidden_dim)."
+    )
+    use_factorized_embedding: bool = Field(
+        default=False,
+        description="Use low-rank factorized embedding instead of standard embedding (reduces params). Only for taonet."
+    )
+    d_embed_rank: int = Field(
+        default=96,
+        description="Rank dimension for factorized embedding. Only used if use_factorized_embedding=True."
+    )
+    vision_encoder_type: str = Field(
+        default="cnn",
+        description="Vision encoder type for multimodal models."
+    )
+    image_size: int = Field(
+        default=224,
+        description="Square image size for multimodal image preprocessing."
+    )
+    vision_output_dim: int = Field(
+        default=256,
+        description="Output feature dimension from the vision encoder before projection."
+    )
+    vision_prefix_tokens: int = Field(
+        default=10,
+        description="Number of visual prefix tokens projected into the LLM sequence."
+    )
+    image_token: str = Field(
+        default="<image>",
+        description="Tokenizer special token reserved for image placeholders."
+    )
+    cnn_channels: list[int] = Field(
+        default_factory=lambda: [32, 64, 128],
+        description="Per-stage channel sizes for the in-repo CNN encoder."
+    )
+    cnn_kernel_size: int = Field(
+        default=3,
+        description="Kernel size for CNN encoder convolutions."
+    )
+
+    # GammaSpaceModel-specific
+    gamma_hidden_dim: int = Field(
+        default=1536,
+        description="GammaSpaceBlock hidden_dim (state-space hidden size). Only used for gamma_net."
+    )
+    gamma_dt_min: float = Field(
+        default=1e-3,
+        description="GammaSpaceBlock dt_min. Only used for gamma_net."
+    )
+    gamma_dt_max: float = Field(
+        default=1e-1,
+        description="GammaSpaceBlock dt_max. Only used for gamma_net."
+    )
+    gamma_dt_init: float = Field(
+        default=1e-2,
+        description="GammaSpaceBlock dt_init. Only used for gamma_net."
+    )
+    gamma_discretization: str = Field(
+        default="bilinear",
+        description="GammaSpaceBlock discretization mode: bilinear, zoh, or euler. Only used for gamma_net."
+    )
+    gamma_prenorm: bool = Field(
+        default=True,
+        description="GammaSpaceBlock prenorm flag. Only used for gamma_net."
+    )
+    gamma_residual_scale: float = Field(
+        default=1.0,
+        description="GammaSpaceBlock residual_scale. Only used for gamma_net."
+    )
+    gamma_activation: str = Field(
+        default="gelu",
+        description="GammaSpaceBlock activation. Only used for gamma_net."
+    )
+    gamma_gate: bool = Field(
+        default=True,
+        description="GammaSpaceBlock gate flag. Only used for gamma_net."
+    )
+    gamma_use_D: bool = Field(
+        default=True,
+        description="GammaSpaceBlock use_D flag. Only used for gamma_net."
+    )
+    gamma_kernel_mode: str = Field(
+        default="auto",
+        description="GammaSpaceBlock kernel_mode: auto, recurrent, or conv. Only used for gamma_net."
+    )
+    gamma_kernel_threshold: int = Field(
+        default=64,
+        description="GammaSpaceBlock kernel_threshold. Only used for gamma_net."
+    )
+    gamma_use_output_linear: bool = Field(
+        default=True,
+        description="GammaSpaceBlock use_output_linear flag. Only used for gamma_net."
+    )
+    gamma_gate_bias: float = Field(
+        default=2.0,
+        description="GammaSpaceBlock gate_bias. Only used for gamma_net."
+    )
+    gamma_input_gate: bool = Field(
+        default=True,
+        description="GammaSpaceBlock input_gate flag. Only used for gamma_net."
+    )
+    gamma_input_gate_bias: float = Field(
+        default=2.0,
+        description="GammaSpaceBlock input_gate_bias. Only used for gamma_net."
+    )
+    gamma_layer_scale_init: float = Field(
+        default=0.1,
+        description="GammaSpaceBlock layer_scale_init. Only used for gamma_net."
+    )
+    
+    # YaRN (Yet another RoPE eXtension) for context length extension
+    rope_scale: float = Field(
+        default=40.0,
+        description="Base RoPE scale factor (default: 40.0). Controls position frequency base."
+    )
+    yarn_enabled: bool = Field(
+        default=False,
+        description="Enable YaRN (Yet another RoPE eXtension) for context length interpolation."
+    )
+    yarn_original_max_seq_length: Optional[int] = Field(
+        default=None,
+        description="Original trained context length that YaRN extends from. Defaults to max_seq_length when unset."
+    )
+    yarn_alpha: float = Field(
+        default=1.0,
+        description="YaRN interpolation smoothness (1.0=smooth, <1.0=aggressive, >1.0=conservative). Only used if yarn_enabled=True."
+    )
+    
+    # Initializations
+    init_std: float = Field(default=0.02, description="Weight initialization standard deviation")
+    
+    @validator("head_dim", always=True)
+    def validate_head_dim(cls, v, values):
+        """Validate head dimension."""
+        if v is None and 'hidden_dim' in values:
+            return values['hidden_dim'] // values.get('num_heads', 12)
+        return v
+    
+    @validator("intermediate_dim", always=True)
+    def validate_intermediate_dim(cls, v, values):
+        """Validate intermediate dimension."""
+        if v is None and 'hidden_dim' in values:
+            return 4 * values['hidden_dim']
+        return v
+
+
+# ============================================================================
+# Dataset Config
+# ============================================================================
+
+
+class DatasetConfig(BaseConfig):
+    """Configuration for dataset loading."""
+    
+    # Local vs HuggingFace dataset selection
+    local: bool = Field(default=False, description="Use local JSONL dataset instead of HuggingFace")
+    
+    # HuggingFace dataset fields
+    dataset_name: Optional[str] = Field(default=None, description="HuggingFace dataset name (e.g., 'wikitext', 'openwebtext')")
+    split: str = Field(default="train", description="Dataset split to use")
+    config: Optional[str] = Field(default=None, description="Dataset config if multi-config (e.g., 'wikitext-103')")
+    
+    # Local JSONL dataset fields
+    jsonl_path: Optional[str] = Field(default=None, description="Path to local JSONL dataset file")
+    text_field: str = Field(default="text", description="Name of text field in JSONL")
+    image_path_column: str = Field(
+        default="image_path",
+        description="Column containing a local image path for multimodal JSONL datasets."
+    )
+    image_path_aliases: list[str] = Field(
+        default_factory=lambda: ["image", "image_path", "image_file", "file_name"],
+        description="Fallback column names to try when a multimodal JSONL record does not contain image_path_column."
+    )
+    caption_prompt: str = Field(
+        default="Describe the image.",
+        description="Prompt paired with caption-style multimodal records that only provide image + text."
+    )
+    
+    # Text column name varies by dataset
+    text_column: str = Field(default="text", description="Name of text column in dataset")
+    
+    # Preprocessing
+    max_samples: Optional[int] = Field(
+        default=None,
+        description="Limit dataset to N samples (useful for debugging)"
+    )
+    cache_dir: str = Field(default=".cache/datasets", description="HuggingFace cache directory")
+    
+    # For SFT/RL datasets with instruction-response format
+    instruction_column: Optional[str] = Field(default=None, description="Instruction column for SFT")
+    response_column: Optional[str] = Field(default=None, description="Response column for SFT")
+    prompt_column: Optional[str] = Field(default=None, description="Prompt column for RL")
+    
+    # Instruction template
+    instruction_template: Optional[str] = Field(
+        default=None,
+        description="Template for combining instruction and response. E.g., '{instruction}\\n{response}'"
+    )
+    
+    # Tokenizer configuration
+    tokenizer_type: Optional[str] = Field(
+        default=None,
+        description="Tokenizer type: 'huggingface' or 'sentencepiece'. If None, defaults based on tokenizer_path."
+    )
+    tokenizer_path: Optional[str] = Field(
+        default=None,
+        description="Path to saved tokenizer (for SentencePiece: .model file, for HuggingFace: model name or local path)"
+    )
+    
+    # Chunked loading for large JSONL files
+    enable_streaming: bool = Field(
+        default=True,
+        description="Enable streaming/chunked loading for large JSONL files to reduce memory usage"
+    )
+    chunk_size_gb: float = Field(
+        default=5.0,
+        description="Approximate chunk size in GB (ignored if samples_per_chunk is set)"
+    )
+    samples_per_chunk: Optional[int] = Field(
+        default=1000,
+        description="Number of samples per chunk (takes precedence over chunk_size_gb). Default: 1000 samples"
+    )
+    
+    # Chunk caching
+    enable_chunk_metadata_cache: bool = Field(
+        default=True,
+        description="Enable caching of chunk metadata (file scan results) to avoid re-scanning large JSONL files"
+    )
+    enable_chunk_data_cache: bool = Field(
+        default=False,
+        description="Enable caching of actual chunk data as separate files for faster loading (uses more disk space)"
+    )
+    chunk_cache_dir: str = Field(
+        default=".cache/chunks",
+        description="Directory to store chunk metadata and data cache files"
+    )
+    
+    # Tokenization parallelization
+    tokenizer_threads: int = Field(
+        default=1,
+        description="Number of background threads for tokenization (1-32 recommended). Higher values speed up tokenization but increase memory usage."
+    )
+    
+    @validator('jsonl_path', always=True)
+    def validate_dataset_source(cls, v, values):
+        """Validate that either local JSONL or HuggingFace dataset is specified."""
+        local = values.get('local', False)
+        dataset_name = values.get('dataset_name')
+        
+        if local and not v:
+            raise ValueError("jsonl_path must be provided when local=True")
+        if not local and not dataset_name:
+            raise ValueError("dataset_name must be provided when local=False (HuggingFace dataset)")
+        
+        return v
+    
+    @validator('tokenizer_threads')
+    def validate_tokenizer_threads(cls, v):
+        """Validate tokenizer_threads is a positive integer."""
+        if v < 1:
+            raise ValueError("tokenizer_threads must be at least 1")
+        if v > 128:
+            raise ValueError("tokenizer_threads should not exceed 128 (recommended: 1-32)")
+        return v
+
+
+# ============================================================================
+# Tokenizer Config
+# ============================================================================
+
+
+class TokenizerConfig(BaseConfig):
+    """Configuration for tokenizer training."""
+    
+    # Dataset source
+    jsonl_path: str = Field(description="Path to JSONL file containing training data")
+    text_field: str = Field(default="text", description="Field name in JSONL for text data")
+    
+    # Training configuration
+    vocab_size: int = Field(default=50000, description="Vocabulary size")
+    model_type: str = Field(default="unigram", description="SentencePiece model type (unigram, bpe, char, word)")
+    character_coverage: float = Field(
+        default=0.9995,
+        description="Character coverage for SentencePiece training"
+    )
+    output_dir: str = Field(default="tokenizers", description="Directory to save trained tokenizer")
+    tokenizer_prefix: Optional[str] = Field(
+        default=None,
+        description="Prefix for tokenizer output files (default: model_type)"
+    )
+    
+    # Custom special tokens registered as SentencePiece user-defined symbols.
+    special_tokens: Optional[list[str]] = Field(
+        default=None,
+        description=(
+            "Custom special tokens such as <think>, <user>, <assistant>, or <image>. "
+            "Built-in SentencePiece tokens are managed by the tokenizer itself and should not be listed here."
+        )
+    )
+    
+    # Data sampling
+    max_samples: Optional[int] = Field(
+        default=None,
+        description="Limit training to first N samples from JSONL (useful for quick testing)"
+    )
+    
+    # Tokenizer metadata
+    tokenizer_name: Optional[str] = Field(
+        default=None,
+        description="Optional name for the tokenizer"
+    )
+
+    @validator("special_tokens", pre=True)
+    def normalize_special_tokens(cls, value):
+        """Normalize special token config to a list of custom token strings."""
+        if value is None:
+            return None
+
+        if isinstance(value, dict):
+            return [str(token) for token in value.keys() if str(token) not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}]
+
+        if isinstance(value, (list, tuple)):
+            return [str(token) for token in value]
+
+        raise ValueError("special_tokens must be a list of token strings")
+
+    @validator("special_tokens")
+    def validate_special_tokens(cls, value):
+        """Validate custom special token registry."""
+        if value is None:
+            return value
+
+        builtin_tokens = {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
+        seen: set[str] = set()
+        normalized: list[str] = []
+        for token in value:
+            if token in builtin_tokens:
+                continue
+            if token in seen:
+                continue
+            seen.add(token)
+            normalized.append(token)
+        return normalized
+
+
+# ============================================================================
+# Training Config
+# ============================================================================
+
+
+class OptimizerConfig(BaseConfig):
+    """Optimizer configuration."""
+    
+    optimizer_type: OptimizerEnum = Field(default=OptimizerEnum.ADAMW, description="Optimizer type")
+    learning_rate: float = Field(default=1e-4, description="Peak learning rate (for Muon 2D weights)")
+    adamw_lr: Optional[float] = Field(
+        default=None,
+        description="Learning rate for AdamW (1D parameters). If None, defaults to learning_rate / 10. Used in hybrid_muon_adamw optimizer."
+    )
+    weight_decay: float = Field(default=1e-2, description="Weight decay (L2 regularization)")
+    betas: tuple[float, float] = Field(default=(0.9, 0.999), description="Adam betas")
+    eps: float = Field(default=1e-8, description="Optimizer epsilon")
+    
+    @validator('adamw_lr', always=True)
+    def set_default_adamw_lr(cls, v, values):
+        """Set default adamw_lr as 1/10 of learning_rate if not specified."""
+        if v is None and 'learning_rate' in values:
+            return values['learning_rate'] / 10
+        return v
+
+
+class SchedulerConfig(BaseConfig):
+    """Learning rate scheduler configuration."""
+    
+    scheduler_type: SchedulerEnum = Field(default=SchedulerEnum.LINEAR_WARMUP, description="Scheduler type")
+    warmup_steps: int = Field(default=0, description="Number of warmup steps (takes precedence over warmup_ratio)")
+    warmup_ratio: float = Field(default=0.1, description="Warmup as fraction of total steps (used if warmup_steps=0)")
+    
+    # Cosine scheduler specific
+    num_cycles: float = Field(default=0.5, description="Number of cycles for cosine schedule")
+    last_epoch: int = Field(default=-1, description="Last epoch for scheduler")
+    
+    # TaoNet 3-phase scheduler (warmup -> steady -> cosine decay)
+    steady_ratio: float = Field(
+        default=0.0,
+        description="Fraction of training steps at peak LR before cosine decay (0.0 = no steady phase). Only for cosineWarmup."
+    )
+    min_lr_ratio: float = Field(
+        default=0.0,
+        description="Minimum LR as fraction of peak LR at end of training (0.0 = decay to 0). Only for cosineWarmup."
+    )
+    
+    @validator('warmup_ratio')
+    def validate_warmup_ratio(cls, v):
+        """Validate warmup ratio is between 0 and 1."""
+        if not 0 <= v <= 1:
+            raise ValueError("warmup_ratio must be between 0 and 1")
+        return v
+    
+    @validator('steady_ratio')
+    def validate_steady_ratio(cls, v):
+        """Validate steady ratio is between 0 and 1."""
+        if not 0 <= v <= 1:
+            raise ValueError("steady_ratio must be between 0 and 1")
+        return v
+    
+    @validator('min_lr_ratio')
+    def validate_min_lr_ratio(cls, v):
+        """Validate min_lr_ratio is between 0 and 1."""
+        if not 0 <= v <= 1:
+            raise ValueError("min_lr_ratio must be between 0 and 1")
+        return v
+    
+    @validator('warmup_steps')
+    def validate_warmup_steps(cls, v):
+        """Validate warmup steps is non-negative."""
+        if v < 0:
+            raise ValueError("warmup_steps must be non-negative")
+        return v
+
+
+class TrainingConfig(BaseConfig):
+    """Base training configuration shared across all modes."""
+    
+    # Data and model
+    model: ModelConfig = Field(default_factory=ModelConfig, description="Model configuration")
+    dataset: DatasetConfig = Field(description="Dataset configuration")
+    
+    # Training hyperparameters
+    batch_size: int = Field(default=32, description="Batch size per device")
+    num_epochs: int = Field(default=3, description="Number of training epochs")
+    max_steps: Optional[int] = Field(
+        default=None,
+        description="Maximum steps (overrides num_epochs if set)"
+    )
+    gradient_accumulation_steps: int = Field(
+        default=1,
+        description="Gradient accumulation steps"
+    )
+    max_grad_norm: float = Field(default=1.0, description="Gradient clipping max norm")
+    
+    # Optimizer
+    optimizer: OptimizerConfig = Field(
+        default_factory=OptimizerConfig,
+        description="Optimizer configuration"
+    )
+    
+    # Scheduler
+    scheduler: SchedulerConfig = Field(
+        default_factory=SchedulerConfig,
+        description="Learning rate scheduler configuration"
+    )
+    
+    # Data type and device
+    dtype: DataTypeEnum = Field(
+        default=DataTypeEnum.BFLOAT16,
+        description="Training data type"
+    )
+    device: str = Field(default="cuda", description="Device to train on (cuda, cpu)")
+    seed: int = Field(default=42, description="Random seed")
+    
+    # Checkpointing
+    checkpoint_dir: str = Field(default="checkpoints", description="Directory to save checkpoints")
+    checkpoint_path: Optional[str] = Field(
+        default=None,
+        description="Path to load pretrained checkpoint (for SFT/RL). If provided, loads weights before training starts."
+    )
+    save_every_steps: int = Field(default=500, description="Save checkpoint every N steps")
+    keep_last_n_checkpoints: int = Field(default=3, description="Keep only last N checkpoints")
+    save_best_model: bool = Field(default=True, description="Save best model based on validation loss")
+    
+    # Validation
+    eval_every_steps: int = Field(default=500, description="Evaluate every N steps")
+    eval_samples: int = Field(default=1000, description="Number of validation samples")
+    
+    # Logging
+    log_every_steps: int = Field(default=10, description="Log metrics every N steps")
+    aim_repo: str = Field(default=".aim", description="AimStack repository path")
+    
+    # Misc
+    num_workers: int = Field(default=0, description="Number of DataLoader workers")
+    pin_memory: bool = Field(default=True, description="Pin memory for DataLoader")
+    use_compile: bool = Field(default=False, description="Use torch.compile (experimental)")
+    
+    # Mode
+    mode: TrainingModeEnum = Field(default=TrainingModeEnum.PRETRAIN, description="Training mode")
+
+
+# ============================================================================
+# Stage-Specific Configs
+# ============================================================================
+
+
+class PretrainConfig(TrainingConfig):
+    """Configuration for pretraining."""
+    
+    mode: Literal[TrainingModeEnum.PRETRAIN] = TrainingModeEnum.PRETRAIN
+    
+    # Pretraining-specific
+    sequence_length: int = Field(default=1024, description="Sequence length for pretraining")
+
+
+class SFTConfig(TrainingConfig):
+    """Configuration for supervised fine-tuning."""
+    
+    mode: Literal[TrainingModeEnum.SFT] = TrainingModeEnum.SFT
+    
+    # SFT-specific
+    response_loss_only: bool = Field(
+        default=True,
+        description="Only compute loss on response/assistant tokens (not instruction/user tokens). Uses -100 label masking."
+    )
+    
+    # Multi-turn conversation role tokens
+    user_token: str = Field(
+        default="<user>",
+        description="Special token representing user/instruction role in conversations"
+    )
+    assistant_token: str = Field(
+        default="<assistant>",
+        description="Special token representing assistant/response role in conversations"
+    )
+
+
+class RLConfig(TrainingConfig):
+    """Configuration for reinforcement learning training."""
+    
+    mode: Literal[TrainingModeEnum.RL] = TrainingModeEnum.RL
+    
+    # RL-specific
+    rl_method: RLMethodEnum = Field(
+        default=RLMethodEnum.PPO,
+        description="RL training method (PPO or DPO)"
+    )
+    
+    # Reward model
+    reward_model_path: str = Field(description="Path to trained reward model checkpoint")
+    
+    # PPO-specific
+    ppo_epochs: int = Field(default=4, description="PPO inner epochs")
+    ppo_clip_ratio: float = Field(default=0.2, description="PPO clipping ratio")
+    entropy_coeff: float = Field(default=0.01, description="Entropy bonus coefficient")
+    value_loss_coeff: float = Field(default=1.0, description="Value function loss coefficient")
+    
+    # DPO-specific (Direct Preference Optimization)
+    dpo_beta: float = Field(default=0.1, description="DPO inverse temperature (beta)")
+    
+    # Prompt distribution
+    prompt_dataset: Optional[DatasetConfig] = Field(
+        default=None,
+        description="Separate dataset for prompts (if different from main dataset)"
+    )
+    generation_max_length: int = Field(
+        default=256,
+        description="Maximum length for generated responses during RL"
+    )
+
+
+class VLMConfig(TrainingConfig):
+    """Configuration for multimodal vision-language connector training."""
+
+    mode: Literal[TrainingModeEnum.VLM] = TrainingModeEnum.VLM
+
+    response_loss_only: bool = Field(
+        default=True,
+        description="Only compute loss on assistant/response tokens for multimodal samples."
+    )
+    user_token: str = Field(
+        default="<user>",
+        description="Special token representing user/instruction role in multimodal conversations."
+    )
+    assistant_token: str = Field(
+        default="<assistant>",
+        description="Special token representing assistant/response role in multimodal conversations."
+    )
+    freeze_llm: bool = Field(
+        default=True,
+        description="Freeze the LLM except for any explicitly unfrozen trailing layers."
+    )
+    unfreeze_last_n_layers: int = Field(
+        default=0,
+        description="Number of final LLM blocks to unfreeze during VLM connector training."
+    )
+    vision_learning_rate: float = Field(
+        default=1e-4,
+        description="Learning rate for the vision encoder and multimodal projector."
+    )
+    llm_learning_rate: float = Field(
+        default=5e-5,
+        description="Learning rate for the unfrozen LLM parameter subset."
+    )
+    vision_prefix_tokens: int = Field(
+        default=10,
+        description="Number of visual prefix tokens to inject in place of a single image placeholder token."
+    )
+    image_token: str = Field(
+        default="<image>",
+        description="Tokenizer token used as the multimodal image placeholder."
+    )
+    image_size: int = Field(
+        default=224,
+        description="Square image size used for multimodal preprocessing."
+    )
+
+    @validator("image_token")
+    def validate_image_token_matches_model_config(cls, v, values):
+        """Validate duplicated image token settings stay aligned."""
+        model = values.get("model")
+        if model is not None and getattr(model, "image_token", None) != v:
+            raise ValueError(
+                f"image_token must match model.image_token for VLM configs. "
+                f"Got top-level `{v}` and model `{model.image_token}`."
+            )
+        return v
+
+    @validator("vision_prefix_tokens")
+    def validate_vision_prefix_tokens_match_model_config(cls, v, values):
+        """Validate duplicated visual prefix settings stay aligned."""
+        model = values.get("model")
+        if model is not None and getattr(model, "vision_prefix_tokens", None) != v:
+            raise ValueError(
+                "vision_prefix_tokens must match model.vision_prefix_tokens for VLM configs. "
+                f"Got top-level `{v}` and model `{model.vision_prefix_tokens}`."
+            )
+        return v
+
+    @validator("image_size")
+    def validate_image_size_matches_model_config(cls, v, values):
+        """Validate duplicated image size settings stay aligned."""
+        model = values.get("model")
+        if model is not None and getattr(model, "image_size", None) != v:
+            raise ValueError(
+                f"image_size must match model.image_size for VLM configs. "
+                f"Got top-level `{v}` and model `{model.image_size}`."
+            )
+        return v
+
+    @validator("unfreeze_last_n_layers")
+    def validate_unfreeze_last_n_layers(cls, v):
+        """Validate unfrozen layer count."""
+        if v < 0:
+            raise ValueError("unfreeze_last_n_layers must be non-negative")
+        return v
+
+    @validator("vision_learning_rate", "llm_learning_rate")
+    def validate_multimodal_learning_rates(cls, v):
+        """Validate multimodal learning rates."""
+        if v <= 0:
+            raise ValueError("multimodal learning rates must be greater than 0")
+        return v
+
+    @validator("vision_prefix_tokens")
+    def validate_vision_prefix_tokens(cls, v):
+        """Validate visual prefix length."""
+        if v < 1:
+            raise ValueError("vision_prefix_tokens must be at least 1")
+        return v
+
+    @validator("image_size")
+    def validate_image_size(cls, v):
+        """Validate image size."""
+        if v < 8:
+            raise ValueError("image_size must be at least 8")
+        return v
+
+
+class VLMSFTConfig(VLMConfig):
+    """Configuration for end-to-end multimodal supervised fine-tuning."""
+
+    mode: Literal[TrainingModeEnum.VLM_SFT] = TrainingModeEnum.VLM_SFT
+    freeze_llm: bool = Field(
+        default=False,
+        description="Whether to freeze the LLM during end-to-end multimodal SFT."
+    )
+
+
+# ============================================================================
+# Factory function
+# ============================================================================
+
+
+def load_config(path: str | Path, mode: TrainingModeEnum | str) -> TrainingConfig:
+    """Load config file and return appropriate config class."""
+    if isinstance(mode, str):
+        mode = TrainingModeEnum(mode)
+    
+    config_map = {
+        TrainingModeEnum.PRETRAIN: PretrainConfig,
+        TrainingModeEnum.SFT: SFTConfig,
+        TrainingModeEnum.RL: RLConfig,
+        TrainingModeEnum.VLM: VLMConfig,
+        TrainingModeEnum.VLM_SFT: VLMSFTConfig,
+    }
+    
+    config_class = config_map[mode]
+    
+    path = Path(path)
+    if path.suffix == '.yaml' or path.suffix == '.yml':
+        return config_class.load_yaml(path)
+    elif path.suffix == '.json':
+        return config_class.load_json(path)
+    else:
+        raise ValueError(f"Unsupported config file format: {path.suffix}")
+
+
+def load_tokenizer_config(path: str | Path) -> TokenizerConfig:
+    """Load tokenizer config from YAML or JSON file."""
+    path = Path(path)
+    if path.suffix == '.yaml' or path.suffix == '.yml':
+        return TokenizerConfig.load_yaml(path)
+    elif path.suffix == '.json':
+        return TokenizerConfig.load_json(path)
+    else:
+        raise ValueError(f"Unsupported config file format: {path.suffix}")

src/taoTrain/core/__init__.py

@@ -0,0 +1,5 @@
+"""Base classes for models, trainers, and datasets."""
+
+from .base import BaseModel, BaseTrainer, BaseDataset, create_model, create_datasets
+
+__all__ = ["BaseModel", "BaseTrainer", "BaseDataset", "create_model", "create_datasets"]

src/taoTrain/core/base.py

@@ -0,0 +1,279 @@
+"""Base classes for models, trainers, and datasets."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Optional, Any, Iterator
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset as TorchDataset
+from taoTrain.config import TrainingConfig, ModelConfig, TrainingModeEnum
+
+
+# ============================================================================
+# Base Model
+# ============================================================================
+
+
+class BaseModel(nn.Module, ABC):
+    """Abstract base class for language models."""
+    
+    def __init__(self, config: Any):
+        """Initialize model with config."""
+        super().__init__()
+        self.config = config
+    
+    @abstractmethod
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Forward pass.
+        
+        Args:
+            input_ids: Shape (batch_size, seq_length)
+            attention_mask: Shape (batch_size, seq_length), optional
+            labels: Shape (batch_size, seq_length), optional (for loss computation)
+        
+        Returns:
+            Dict with keys:
+                - 'logits': Shape (batch_size, seq_length, vocab_size)
+                - 'loss': Scalar (if labels provided)
+        """
+        pass
+    
+    def count_parameters(self) -> int:
+        """Count total trainable parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+    
+    def get_num_layers(self) -> int:
+        """Get number of layers (for model architecture)."""
+        return self.config.num_layers
+
+    def get_optimizer_param_groups(self, config: TrainingConfig) -> Optional[list[dict[str, Any]]]:
+        """Optional custom optimizer param groups for advanced training setups."""
+        return None
+
+
+# ============================================================================
+# Base Dataset
+# ============================================================================
+
+
+class BaseDataset(TorchDataset, ABC):
+    """Abstract base class for datasets."""
+    
+    def __init__(self, config: "TrainingConfig"):
+        """Initialize dataset."""
+        self.config = config
+        self.data = None
+    
+    @abstractmethod
+    def __len__(self) -> int:
+        """Return dataset size."""
+        pass
+    
+    @abstractmethod
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        """
+        Get a single sample.
+        
+        Returns:
+            Dict with keys:
+                - 'input_ids': 1D tensor of token IDs
+                - 'attention_mask': 1D tensor of attention mask
+                - 'labels': 1D tensor of labels (optional)
+        """
+        pass
+    
+    def load_dataset(self) -> None:
+        """Load dataset from HuggingFace or other source."""
+        pass
+    
+    def preprocess(self) -> None:
+        """Preprocess dataset (tokenization, etc)."""
+        pass
+
+
+# ============================================================================
+# Base Trainer
+# ============================================================================
+
+
+class BaseTrainer(ABC):
+    """Abstract base class for trainers."""
+    
+    def __init__(
+        self,
+        model: BaseModel,
+        train_dataset: BaseDataset,
+        val_dataset: Optional[BaseDataset],
+        config: TrainingConfig,
+        device: torch.device,
+    ):
+        """Initialize trainer."""
+        self.model = model.to(device)
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.config = config
+        self.device = device
+        
+        # Training state
+        self.global_step = 0
+        self.current_epoch = 0
+        self.best_loss = float('inf')
+        
+        # Logging
+        self.logger = None
+        
+        # Optimizer and scheduler (to be set up by subclass)
+        self.optimizer = None
+        self.scheduler = None
+    
+    @abstractmethod
+    def training_step(self, batch: dict[str, torch.Tensor]) -> dict[str, float]:
+        """
+        Single training step.
+        
+        Args:
+            batch: Training batch with input_ids, attention_mask, labels, etc.
+        
+        Returns:
+            Dict with metrics (e.g., {'loss': 0.5, 'accuracy': 0.8})
+        """
+        pass
+    
+    @abstractmethod
+    def validation_step(self, batch: dict[str, torch.Tensor]) -> dict[str, float]:
+        """
+        Single validation step.
+        
+        Args:
+            batch: Validation batch
+        
+        Returns:
+            Dict with validation metrics
+        """
+        pass
+    
+    @abstractmethod
+    def train_epoch(self) -> dict[str, float]:
+        """
+        Train for one epoch.
+        
+        Returns:
+            Dict with epoch-level metrics
+        """
+        pass
+    
+    @abstractmethod
+    def validate(self) -> dict[str, float]:
+        """
+        Run validation on the entire validation set.
+        
+        Returns:
+            Dict with validation metrics
+        """
+        pass
+    
+    def save_checkpoint(self, path: str | Path) -> None:
+        """
+        Save checkpoint in canonical format.
+        
+        Uses canonical checkpoint format:
+        {
+            'step': int,
+            'model_state': state_dict,
+            'optimizer_state': state_dict,
+            'config': dict,
+            'metrics': dict,
+            'global_step': int,           # Legacy compat
+            'current_epoch': int,         # Legacy compat
+            'best_loss': float,           # Legacy compat
+        }
+        
+        Args:
+            path: Path to save checkpoint
+        """
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Save in canonical format
+        checkpoint = {
+            # Canonical format keys
+            'step': self.global_step,
+            'model_state': self.model.state_dict(),
+            'optimizer_state': self.optimizer.state_dict() if self.optimizer else None,
+            'config': self.config.to_dict(),
+            'metrics': {},
+            # Legacy format keys (for backward compatibility with code that reads them)
+            'global_step': self.global_step,
+            'current_epoch': self.current_epoch,
+            'best_loss': self.best_loss,
+        }
+        
+        torch.save(checkpoint, path)
+    
+    def load_checkpoint(self, path: str | Path) -> None:
+        """
+        Load checkpoint (handles both canonical and legacy formats).
+        
+        Args:
+            path: Path to checkpoint
+        """
+        path = Path(path)
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        # Try canonical keys first, fall back to legacy keys
+        model_state_key = 'model_state' if 'model_state' in checkpoint else 'model_state_dict'
+        optimizer_state_key = 'optimizer_state' if 'optimizer_state' in checkpoint else 'optimizer_state_dict'
+        
+        self.model.load_state_dict(checkpoint[model_state_key])
+        if self.optimizer and checkpoint.get(optimizer_state_key):
+            self.optimizer.load_state_dict(checkpoint[optimizer_state_key])
+        
+        # Try canonical 'step' first, fall back to legacy 'global_step'
+        self.global_step = checkpoint.get('step', checkpoint.get('global_step', 0))
+        self.current_epoch = checkpoint.get('current_epoch', 0)
+        self.best_loss = checkpoint.get('best_loss', float('inf'))
+    
+    def _get_lr(self) -> float:
+        """Get current learning rate from optimizer."""
+        for param_group in self.optimizer.param_groups:
+            return param_group['lr']
+        return 0.0
+
+
+# ============================================================================
+# Utility functions
+# ============================================================================
+
+
+def create_model(config: TrainingConfig, device: torch.device) -> BaseModel:
+    """Create model from config (calls registry)."""
+    from taoTrain.models import get_model
+    if config.mode in {TrainingModeEnum.VLM, TrainingModeEnum.VLM_SFT}:
+        return get_model(config, device=device)
+    return get_model(config.model, device=device)
+
+
+def create_datasets(
+    config: TrainingConfig,
+) -> tuple[BaseDataset, Optional[BaseDataset]]:
+    """Create train and validation datasets using factory pattern."""
+    # Import here to avoid circular imports
+    from taoTrain.data import DatasetFactory
+    
+    # Create train dataset
+    train_dataset = DatasetFactory.create_dataset(config, split="train")
+    
+    # Create validation dataset (only for HuggingFace datasets with explicit validation split)
+    val_dataset = None
+    if not config.dataset.local and hasattr(config.dataset, "validation_split"):
+        val_dataset = DatasetFactory.create_dataset(config, split="validation")
+    
+    return train_dataset, val_dataset

src/taoTrain/data/__init__.py

@@ -0,0 +1,47 @@
+"""Dataset implementations and loaders."""
+
+# HuggingFace-based datasets
+from .hf_base import BaseHFDataset
+from .hf_pretrain import PretrainDataset
+from .hf_sft import SFTDataset
+from .hf_rl import RLDataset
+
+# JSONL-based datasets (async-only)
+from .jsonl_base import BaseJSONLDataset
+from .pretrain_jsonl import PretrainJSONLDataset
+from .sft_jsonl import SFTJSONLDataset
+from .rl_jsonl import RLJSONLDataset
+from .vlm_jsonl import VLMJSONLDataset, VLMSFTJSONLDataset
+
+# Utilities
+from .tokenizer import SentencePieceTokenizerWrapper
+from .sft_utils import parse_sft_record, build_sft_sequence_tokens, apply_response_masking
+from .loaders import get_dataloader
+from .async_loader import AsyncBatchIterator
+from .tokenization_queue import TokenizationQueue
+from .factory import DatasetFactory
+
+__all__ = [
+    # HuggingFace datasets
+    "BaseHFDataset",
+    "PretrainDataset",
+    "SFTDataset", 
+    "RLDataset",
+    # JSONL datasets
+    "BaseJSONLDataset",
+    "PretrainJSONLDataset",
+    "SFTJSONLDataset",
+    "RLJSONLDataset",
+    "VLMJSONLDataset",
+    "VLMSFTJSONLDataset",
+    # Utilities
+    "SentencePieceTokenizerWrapper",
+    "parse_sft_record",
+    "build_sft_sequence_tokens",
+    "apply_response_masking",
+    # Data loading
+    "get_dataloader",
+    "AsyncBatchIterator",
+    "TokenizationQueue",
+    "DatasetFactory",
+]