Upload folder using huggingface_hub

training/aws_lora_base_tools_s3.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""SageMaker entrypoint: LoRA tool-use SFT para VectraYX Base 260M - S3 ONLY.
+
+Igual que aws_lora_nano_tools_s3.py pero con checkpoint y config de Base 260M.
+
+Hyperparameters via env:
+    CORPUS_NAME    = "v3_bash" (default)
+    EPOCHS         = "5"
+    LR             = "2e-4"
+    LORA_RANK      = "16"
+    LORA_ALPHA     = "32"
+    SEED           = "42"
+"""
+import os, sys, json, subprocess, shutil
+from pathlib import Path
+
+S3_BUCKET = "s3://vectrayx-sagemaker-792811916323"
+SM_OUTPUT = Path(os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+WD = Path("/opt/ml/code/work")
+ENV = {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
+
+# Base 260M — checkpoint post-P3 (phase3_last.pt)
+BASE_CKPT = f"{S3_BUCKET}/checkpoints/vectrayx-base-20260506-1901/phase3_last.pt"
+BASE_CFG  = "base.json"
+BASE_BATCH = 8
+BASE_ACCUM = 8   # effective batch = 64
+
+
+def die(m): print(f"\n[FATAL] {m}", flush=True); sys.exit(1)
+
+
+def s3_download(src, dst):
+    dst = Path(dst)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    r = subprocess.run(["aws", "s3", "cp", src, str(dst)],
+                       capture_output=True, text=True)
+    if r.returncode != 0:
+        die(f"s3 download failed: {src}\n{r.stderr}")
+    print(f"[s3] ✓ {src} ({dst.stat().st_size/1e6:.1f}MB)", flush=True)
+
+
+def sh(cmd, cwd=None):
+    print(f"$ {cmd}", flush=True)
+    r = subprocess.run(cmd, shell=True, env={**os.environ, **ENV},
+                       cwd=str(cwd or WD))
+    if r.returncode != 0:
+        die(f"Failed: {cmd}")
+
+
+def main():
+    corpus_name = os.environ.get("CORPUS_NAME", "v3_bash")
+    epochs      = int(os.environ.get("EPOCHS", "5"))
+    lr          = float(os.environ.get("LR", "2e-4"))
+    lora_rank   = int(os.environ.get("LORA_RANK", "16"))
+    lora_alpha  = float(os.environ.get("LORA_ALPHA", "32"))
+    seed        = int(os.environ.get("SEED", "42"))
+
+    WD.mkdir(parents=True, exist_ok=True)
+    SM_OUTPUT.mkdir(parents=True, exist_ok=True)
+
+    print(f"[config] model=base corpus={corpus_name} epochs={epochs} lr={lr} "
+          f"lora_rank={lora_rank} lora_alpha={lora_alpha} seed={seed}", flush=True)
+
+    # 1. Deps
+    subprocess.run([sys.executable, "-m", "pip", "install", "-q",
+                    "sentencepiece", "tokenizers"], check=True)
+
+    # 2. Código training_v2 (incluye finetune_lora_tools.py y utils.py corregidos)
+    print("[code] Downloading training_v2 from S3...", flush=True)
+    subprocess.run(["aws", "s3", "cp",
+                    f"{S3_BUCKET}/code/training_v2.tar.gz",
+                    "/tmp/tv2.tar.gz"], check=True)
+    sh("tar xzf /tmp/tv2.tar.gz", cwd=WD)
+    print(f"[code] ✓ training_v2 extracted", flush=True)
+
+    # 3. Tokenizer (mismo que Nano — BPE 16384)
+    s3_download(f"{S3_BUCKET}/tokenizers/vectrayx_bpe.model", WD/"tokenizer.model")
+
+    # 4. Checkpoint Base 260M (post-P3, pre-SFT)
+    s3_download(BASE_CKPT, WD/"resume.pt")
+
+    # 5. Corpus tool-use
+    s3_download(f"{S3_BUCKET}/training-data/tool_sft_{corpus_name}.jsonl",
+                WD/"tool_sft.jsonl")
+
+    # 6. Eval data — b4_tooluse_v2 con bash básico (60%)
+    eval_dir = WD / "eval_data"
+    for b in ["b1_cveqa", "b2_classification", "b3_commands", "b5_conversational"]:
+        try:
+            s3_download(f"{S3_BUCKET}/eval-data/{b}.jsonl",
+                        eval_dir / f"{b}.jsonl")
+        except Exception:
+            print(f"[s3] skip (optional) {b}.jsonl", flush=True)
+    s3_download(f"{S3_BUCKET}/eval-data/b4_tooluse_v2.jsonl",
+                eval_dir / "b4_tooluse.jsonl")
+
+    # 7. LoRA fine-tune sobre Base 260M
+    out_dir = WD / "checkpoints/lora_tool_sft"
+    sh(f"{sys.executable} -m training_v2.train.finetune_lora_tools "
+       f"--config {WD}/training_v2/configs/{BASE_CFG} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--resume {WD}/resume.pt "
+       f"--tool-corpus {WD}/tool_sft.jsonl "
+       f"--out {out_dir} "
+       f"--lora-rank {lora_rank} "
+       f"--lora-alpha {lora_alpha} "
+       f"--batch-size {BASE_BATCH} "
+       f"--grad-accum {BASE_ACCUM} "
+       f"--epochs {epochs} "
+       f"--lr {lr} "
+       f"--seed {seed}")
+
+    # 8. Copiar artefactos
+    shutil.copy(out_dir / "final.pt",           SM_OUTPUT / "final.pt")
+    shutil.copy(out_dir / "final_lora_only.pt", SM_OUTPUT / "final_lora_only.pt")
+    shutil.copy(WD / f"training_v2/configs/{BASE_CFG}", SM_OUTPUT / "model_config.json")
+
+    # 9. Benchmark B1–B5
+    sh(f"{sys.executable} -m training_v2.eval.benchmark "
+       f"--checkpoint {out_dir}/final.pt "
+       f"--config {WD}/training_v2/configs/{BASE_CFG} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--data-dir {eval_dir} "
+       f"--out {SM_OUTPUT}/bench_lora_tools.json")
+
+    # 10. Manifest
+    manifest = {
+        "model": "base",
+        "method": "lora",
+        "corpus": corpus_name,
+        "lora_rank": lora_rank,
+        "lora_alpha": lora_alpha,
+        "epochs": epochs,
+        "lr": lr,
+        "seed": seed,
+        "resume_from": BASE_CKPT,
+        "effective_batch": BASE_BATCH * BASE_ACCUM,
+    }
+    (SM_OUTPUT / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"[done] LoRA tool-SFT Base 260M → {SM_OUTPUT}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/aws_lora_nano_tools_s3.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""SageMaker entrypoint: LoRA tool-use SFT para VectraYX Nano - S3 ONLY.
+
+Hyperparameters via env:
+    CORPUS_NAME    = "v3_bash" (default)
+    EPOCHS         = "5"
+    LR             = "2e-4"
+    LORA_RANK      = "16"
+    LORA_ALPHA     = "32"
+    SEED           = "42"
+"""
+import os, sys, json, subprocess, shutil
+from pathlib import Path
+
+S3_BUCKET = "s3://vectrayx-sagemaker-792811916323"
+SM_OUTPUT = Path(os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+WD = Path("/opt/ml/code/work")
+ENV = {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
+
+# Nano config — checkpoint post-SFT mixto
+NANO_CKPT = f"{S3_BUCKET}/checkpoints/nano_sft_v5.pt"
+NANO_CFG  = "nano.json"
+NANO_BATCH = 16
+NANO_ACCUM = 4   # effective batch = 64
+
+
+def die(m): print(f"\n[FATAL] {m}", flush=True); sys.exit(1)
+
+
+def s3_download(src, dst):
+    dst = Path(dst)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    r = subprocess.run(["aws", "s3", "cp", src, str(dst)],
+                       capture_output=True, text=True)
+    if r.returncode != 0:
+        die(f"s3 download failed: {src}\n{r.stderr}")
+    print(f"[s3] ✓ {src} ({dst.stat().st_size/1e6:.1f}MB)", flush=True)
+
+
+def sh(cmd, cwd=None):
+    print(f"$ {cmd}", flush=True)
+    r = subprocess.run(cmd, shell=True, env={**os.environ, **ENV},
+                       cwd=str(cwd or WD))
+    if r.returncode != 0:
+        die(f"Failed: {cmd}")
+
+
+def main():
+    corpus_name = os.environ.get("CORPUS_NAME", "v3_bash")
+    epochs      = int(os.environ.get("EPOCHS", "5"))
+    lr          = float(os.environ.get("LR", "2e-4"))
+    lora_rank   = int(os.environ.get("LORA_RANK", "16"))
+    lora_alpha  = float(os.environ.get("LORA_ALPHA", "32"))
+    seed        = int(os.environ.get("SEED", "42"))
+
+    WD.mkdir(parents=True, exist_ok=True)
+    SM_OUTPUT.mkdir(parents=True, exist_ok=True)
+
+    print(f"[config] corpus={corpus_name} epochs={epochs} lr={lr} "
+          f"lora_rank={lora_rank} lora_alpha={lora_alpha} seed={seed}", flush=True)
+
+    # 1. Deps
+    subprocess.run([sys.executable, "-m", "pip", "install", "-q",
+                    "sentencepiece", "tokenizers"], check=True)
+
+    # 2. Código training_v2
+    print("[code] Downloading training_v2 from S3...", flush=True)
+    subprocess.run(["aws", "s3", "cp",
+                    f"{S3_BUCKET}/code/training_v2.tar.gz",
+                    "/tmp/tv2.tar.gz"], check=True)
+    sh("tar xzf /tmp/tv2.tar.gz", cwd=WD)
+    print(f"[code] ✓ training_v2 extracted", flush=True)
+
+    # 3. Tokenizer
+    s3_download(f"{S3_BUCKET}/tokenizers/vectrayx_bpe.model", WD/"tokenizer.model")
+
+    # 4. Checkpoint base Nano (post-SFT mixto)
+    s3_download(NANO_CKPT, WD/"resume.pt")
+
+    # 5. Corpus tool-use
+    s3_download(f"{S3_BUCKET}/training-data/tool_sft_{corpus_name}.jsonl",
+                WD/"tool_sft.jsonl")
+
+    # 6. Eval data — b4_tooluse_v2 tiene 50 preguntas con bash básico
+    eval_dir = WD / "eval_data"
+    for b in ["b1_cveqa", "b2_classification", "b3_commands",
+              "b5_conversational"]:
+        try:
+            s3_download(f"{S3_BUCKET}/eval-data/{b}.jsonl",
+                        eval_dir / f"{b}.jsonl")
+        except Exception:
+            print(f"[s3] skip (optional) {b}.jsonl", flush=True)
+    # B4 v2 — benchmark ampliado con bash básico (60%) + MCP (40%)
+    s3_download(f"{S3_BUCKET}/eval-data/b4_tooluse_v2.jsonl",
+                eval_dir / "b4_tooluse.jsonl")  # mismo nombre para que benchmark.py lo encuentre
+
+    # 7. LoRA fine-tune
+    out_dir = WD / "checkpoints/lora_tool_sft"
+    sh(f"{sys.executable} -m training_v2.train.finetune_lora_tools "
+       f"--config {WD}/training_v2/configs/{NANO_CFG} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--resume {WD}/resume.pt "
+       f"--tool-corpus {WD}/tool_sft.jsonl "
+       f"--out {out_dir} "
+       f"--lora-rank {lora_rank} "
+       f"--lora-alpha {lora_alpha} "
+       f"--batch-size {NANO_BATCH} "
+       f"--grad-accum {NANO_ACCUM} "
+       f"--epochs {epochs} "
+       f"--lr {lr} "
+       f"--seed {seed}")
+
+    # 8. Copiar artefactos al output
+    shutil.copy(out_dir / "final.pt",          SM_OUTPUT / "final.pt")
+    shutil.copy(out_dir / "final_lora_only.pt", SM_OUTPUT / "final_lora_only.pt")
+    shutil.copy(WD / f"training_v2/configs/{NANO_CFG}", SM_OUTPUT / "model_config.json")
+
+    # 9. Benchmark B1–B5 (usa final.pt merged)
+    sh(f"{sys.executable} -m training_v2.eval.benchmark "
+       f"--checkpoint {out_dir}/final.pt "
+       f"--config {WD}/training_v2/configs/{NANO_CFG} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--data-dir {eval_dir} "
+       f"--out {SM_OUTPUT}/bench_lora_tools.json")
+
+    # 10. Manifest
+    manifest = {
+        "model": "nano",
+        "method": "lora",
+        "corpus": corpus_name,
+        "lora_rank": lora_rank,
+        "lora_alpha": lora_alpha,
+        "epochs": epochs,
+        "lr": lr,
+        "seed": seed,
+        "resume_from": NANO_CKPT,
+        "effective_batch": NANO_BATCH * NANO_ACCUM,
+    }
+    (SM_OUTPUT / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"[done] LoRA tool-SFT Nano → {SM_OUTPUT}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/aws_tool_sft_train_s3.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""SageMaker entrypoint: tool-use mini-SFT focalizado (Nano o Base) - S3 ONLY.
+
+Hyperparameters via env:
+    MODEL          = "nano" | "base"
+    CORPUS_NAME    = "v1" | "v2"
+    EPOCHS         = "2"
+    LR             = "1e-5"
+    SEED           = "42"
+"""
+import os, sys, json, subprocess, shutil
+from pathlib import Path
+
+S3_BUCKET = "s3://vectrayx-sagemaker-792811916323"
+SM_OUTPUT = Path(os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+WD = Path("/opt/ml/code/work")
+ENV = {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
+
+MODEL_CFG = {
+    "nano": {
+        "config":   "nano.json",
+        "ckpt_src": f"{S3_BUCKET}/checkpoints/nano_sft_v5.pt",
+        "batch":    16,
+        "accum":    4,
+    },
+    "base": {
+        "config":   "base.json",
+        "ckpt_src": f"{S3_BUCKET}/checkpoints/vectrayx-base-20260506-1901/phase3_last.pt",
+        "batch":    8,
+        "accum":    8,
+    },
+}
+
+
+def die(m): print(f"\n[FATAL] {m}", flush=True); sys.exit(1)
+
+
+def s3_download(src, dst):
+    """Download from S3 using AWS CLI."""
+    dst = Path(dst)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    r = subprocess.run(["aws", "s3", "cp", src, str(dst)], 
+                      capture_output=True, text=True)
+    if r.returncode != 0:
+        die(f"s3 download failed: {src}\n{r.stderr}")
+    print(f"[s3] ✓ {src} ({dst.stat().st_size/1e6:.1f}MB)", flush=True)
+
+
+def sh(cmd, cwd=None):
+    print(f"$ {cmd}", flush=True)
+    r = subprocess.run(cmd, shell=True, env={**os.environ, **ENV}, cwd=str(cwd or WD))
+    if r.returncode != 0: die(f"Failed: {cmd}")
+
+
+def main():
+    model_name = os.environ.get("MODEL", "nano")
+    corpus_name = os.environ.get("CORPUS_NAME", "v1")
+    epochs = int(os.environ.get("EPOCHS", "2"))
+    lr = float(os.environ.get("LR", "1e-5"))
+    seed = int(os.environ.get("SEED", "42"))
+
+    if model_name not in MODEL_CFG: die(f"Unknown MODEL={model_name}")
+    cfg = MODEL_CFG[model_name]
+
+    WD.mkdir(parents=True, exist_ok=True)
+    SM_OUTPUT.mkdir(parents=True, exist_ok=True)
+
+    # 1. Deps
+    subprocess.run([sys.executable, "-m", "pip", "install", "-q",
+                    "sentencepiece", "tokenizers"], check=True)
+
+    # 2. Download and extract training_v2 code
+    print("[code] Downloading training_v2 from S3...", flush=True)
+    subprocess.run(["aws", "s3", "cp", 
+                   "s3://vectrayx-sagemaker-792811916323/code/training_v2.tar.gz",
+                   "/tmp/tv2.tar.gz"], check=True)
+    sh("tar xzf /tmp/tv2.tar.gz", cwd=WD)
+    print(f"[code] ✓ training_v2 extracted to {WD}", flush=True)
+
+    # 3. Tokenizer
+    s3_download(f"{S3_BUCKET}/tokenizers/vectrayx_bpe.model", WD/"tokenizer.model")
+
+    # 4. Checkpoint inicial
+    s3_download(cfg["ckpt_src"], WD/"resume.pt")
+
+    # 5. Tool SFT corpus
+    s3_download(f"{S3_BUCKET}/training-data/tool_sft_{corpus_name}.jsonl",
+                WD/"tool_sft.jsonl")
+
+    # 6. Eval data
+    eval_dir = WD / "eval_data"
+    for b in ["b1_cveqa", "b2_classification", "b3_commands",
+              "b4_tooluse", "b5_conversational"]:
+        try:
+            s3_download(f"{S3_BUCKET}/eval-data/{b}.jsonl",
+                       eval_dir/f"{b}.jsonl")
+        except:
+            print(f"[s3] skip (optional) {b}.jsonl", flush=True)
+
+    # 7. Mini-SFT focalizado
+    out_dir = WD / "checkpoints/tool_sft"
+    sh(f"{sys.executable} -m training_v2.train.finetune_tools "
+       f"--config {WD}/training_v2/configs/{cfg['config']} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--resume {WD}/resume.pt "
+       f"--tool-corpus {WD}/tool_sft.jsonl "
+       f"--out {out_dir} "
+       f"--batch-size {cfg['batch']} --grad-accum {cfg['accum']} "
+       f"--epochs {epochs} --lr {lr} --seed {seed}")
+
+    # 8. Copiar checkpoint final
+    shutil.copy(out_dir/"final.pt", SM_OUTPUT/"final.pt")
+    shutil.copy(WD/f"training_v2/configs/{cfg['config']}",
+                SM_OUTPUT/"model_config.json")
+
+    # 9. Bench B1–B5
+    sh(f"{sys.executable} -m training_v2.eval.benchmark "
+       f"--checkpoint {out_dir}/final.pt "
+       f"--config {WD}/training_v2/configs/{cfg['config']} "
+       f"--tokenizer {WD}/tokenizer.model "
+       f"--data-dir {eval_dir} "
+       f"--out {SM_OUTPUT}/bench_tool_sft.json")
+
+    # 10. Manifest
+    manifest = {
+        "model": model_name,
+        "corpus": corpus_name,
+        "epochs": epochs, "lr": lr, "seed": seed,
+        "resume_from": cfg["ckpt_src"],
+    }
+    (SM_OUTPUT/"manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"[done] tool-SFT {model_name}/{corpus_name}/seed={seed} → {SM_OUTPUT}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/finetune_lora_tools.py

@@ -0,0 +1,367 @@
+"""LoRA tool-use SFT para VectraYX Nano.
+
+Aplica LoRA sobre las proyecciones de atención (wq, wk, wv, wo) del modelo
+custom VectraYXNano. Congela todos los pesos base y solo entrena los adaptadores.
+
+Ventaja sobre full fine-tune:
+- Solo ~0.5% de parámetros entrenables (~200K vs 42M)
+- Menos riesgo de catastrofic forgetting en B1/B2/B5
+- SmolLM2-135M logra B4=0.16 con LoRA — probamos si Nano puede hacer lo mismo
+
+Run example:
+    python -m training_v2.train.finetune_lora_tools \
+        --config training_v2/configs/nano.json \
+        --tokenizer models/vectrayx_bpe.model \
+        --resume checkpoints/nano_sft_v5.pt \
+        --tool-corpus corpus/tool_sft_v2_simple.jsonl \
+        --out checkpoints/nano_lora_tools \
+        --lora-rank 16 --lora-alpha 32 \
+        --batch-size 16 --grad-accum 4 --epochs 5 --lr 2e-4
+"""
+
+import argparse
+import json
+import math
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from training_v2.data.sft_dataset import SFTDataset
+from training_v2.model.transformer import VectraYXNano, ModelConfig
+from training_v2.train.utils import (
+    cosine_with_warmup, log_jsonl,
+)
+
+
+# ---------------------------------------------------------------------------
+# LoRA implementation
+# ---------------------------------------------------------------------------
+
+class LoRALinear(nn.Module):
+    """Reemplaza un nn.Linear con LoRA: W' = W + (B @ A) * scale."""
+
+    def __init__(self, linear: nn.Linear, rank: int, alpha: float):
+        super().__init__()
+        self.linear = linear          # pesos base — CONGELADOS
+        self.rank = rank
+        self.scale = alpha / rank
+
+        in_f = linear.in_features
+        out_f = linear.out_features
+
+        # A: inicialización kaiming, B: ceros (LoRA paper §4)
+        self.lora_A = nn.Parameter(torch.empty(rank, in_f))
+        self.lora_B = nn.Parameter(torch.zeros(out_f, rank))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+
+        # Congelar pesos base
+        for p in self.linear.parameters():
+            p.requires_grad_(False)
+
+    def forward(self, x):
+        base = self.linear(x)
+        # Asegurar que lora_A y lora_B estén en el mismo device que x
+        lora = (x @ self.lora_A.to(x.device).T) @ self.lora_B.to(x.device).T
+        return base + lora * self.scale
+
+
+def inject_lora(model: nn.Module, rank: int, alpha: float,
+                target_modules=("wq", "wk", "wv", "wo")) -> int:
+    """Inyecta LoRA en todas las capas de atención del modelo.
+    
+    Retorna el número de parámetros entrenables.
+    """
+    replaced = 0
+    for name, module in model.named_modules():
+        for attr_name in target_modules:
+            if hasattr(module, attr_name):
+                original = getattr(module, attr_name)
+                if isinstance(original, nn.Linear):
+                    setattr(module, attr_name, LoRALinear(original, rank, alpha))
+                    replaced += 1
+
+    # Congelar todo excepto LoRA
+    for name, param in model.named_parameters():
+        if "lora_A" not in name and "lora_B" not in name:
+            param.requires_grad_(False)
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    print(f"[lora] Inyectado en {replaced} módulos | "
+          f"Entrenables: {trainable/1e3:.1f}K / {total/1e6:.2f}M "
+          f"({trainable/total*100:.2f}%)")
+    return trainable
+
+
+def save_lora_checkpoint(path: Path, model: nn.Module, optimizer, step: int,
+                         extra: dict = None):
+    """Guarda solo los pesos LoRA (no el modelo base)."""
+    lora_state = {k: v for k, v in model.state_dict().items()
+                  if "lora_A" in k or "lora_B" in k}
+    torch.save({
+        "lora_state_dict": lora_state,
+        "optimizer_state_dict": optimizer.state_dict() if optimizer else None,
+        "step": step,
+        **(extra or {}),
+    }, path)
+    print(f"[save] LoRA checkpoint → {path} ({path.stat().st_size/1e6:.1f}MB)")
+
+
+def load_lora_checkpoint(path: Path, model: nn.Module, optimizer=None,
+                         map_location="cpu"):
+    """Carga pesos LoRA en el modelo."""
+    ckpt = torch.load(path, map_location=map_location)
+    missing, unexpected = model.load_state_dict(ckpt["lora_state_dict"], strict=False)
+    lora_keys = [k for k in ckpt["lora_state_dict"]]
+    print(f"[load] LoRA: {len(lora_keys)} keys loaded, "
+          f"{len(missing)} missing, {len(unexpected)} unexpected")
+    if optimizer and ckpt.get("optimizer_state_dict"):
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    return ckpt.get("step", 0)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", required=True)
+    p.add_argument("--tokenizer", required=True)
+    p.add_argument("--resume", required=True, help="checkpoint base a fine-tunear")
+    p.add_argument("--tool-corpus", required=True, help="tool-use JSONL corpus")
+    p.add_argument("--out", required=True)
+    # LoRA
+    p.add_argument("--lora-rank", type=int, default=16,
+                   help="LoRA rank r (default 16)")
+    p.add_argument("--lora-alpha", type=float, default=32.0,
+                   help="LoRA alpha (default 32, scale=alpha/rank=2)")
+    p.add_argument("--lora-targets", nargs="+",
+                   default=["wq", "wk", "wv", "wo"],
+                   help="Módulos de atención a inyectar LoRA")
+    # Training
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--grad-accum", type=int, default=4)
+    p.add_argument("--epochs", type=int, default=5)
+    p.add_argument("--lr", type=float, default=2e-4,
+                   help="LR más alto que full FT (LoRA converge más rápido)")
+    p.add_argument("--weight-decay", type=float, default=0.01)
+    p.add_argument("--grad-clip", type=float, default=1.0)
+    p.add_argument("--warmup-frac", type=float, default=0.05)
+    p.add_argument("--num-workers", type=int, default=2)
+    p.add_argument("--log-every", type=int, default=10)
+    p.add_argument("--save-every", type=int, default=200)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--dtype", default="bfloat16",
+                   choices=["bfloat16", "float16", "float32"])
+    p.add_argument("--max-steps", type=int, default=None)
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # 1. Cargar modelo base
+    cfg = ModelConfig.from_json(args.config)
+    model = VectraYXNano(cfg).to(args.device)
+    total_params = model.num_params()
+    print(f"[model] {total_params/1e6:.2f}M params (base)")
+
+    # Cargar checkpoint base (full weights) usando load_checkpoint de utils
+    from training_v2.train.utils import load_checkpoint as _load_ckpt
+    _load_ckpt(args.resume, model, optimizer=None, map_location=args.device)
+    print(f"[resume] {args.resume}")
+
+    # 2. Inyectar LoRA
+    trainable = inject_lora(model, rank=args.lora_rank, alpha=args.lora_alpha,
+                            target_modules=args.lora_targets)
+    # Mover parámetros LoRA al mismo device que el modelo
+    model = model.to(args.device)
+
+    # 3. Tokenizer
+    sp = spm.SentencePieceProcessor()
+    sp.load(args.tokenizer)
+    pad_id = sp.pad_id() if sp.pad_id() >= 0 else 0
+
+    # 4. Dataset
+    block_size = cfg.max_seq_len
+    tool_corpus = Path(args.tool_corpus)
+    if not tool_corpus.exists():
+        raise FileNotFoundError(f"Tool corpus not found: {tool_corpus}")
+
+    dataset = SFTDataset([tool_corpus], sp, block_size, pad_id=pad_id, seed=args.seed)
+    print(f"[dataset] {len(dataset)} ejemplos de {tool_corpus.name}")
+
+    # 5. Output dir
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    log_path = out_dir / "train_log.jsonl"
+
+    # 6. Optimizer — solo parámetros LoRA
+    lora_params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.AdamW(lora_params, lr=args.lr,
+                                  weight_decay=args.weight_decay,
+                                  betas=(0.9, 0.95))
+
+    # 7. AMP
+    dtype = {"bfloat16": torch.bfloat16,
+             "float16": torch.float16,
+             "float32": torch.float32}[args.dtype]
+    use_amp = args.device == "cuda" and dtype != torch.float32
+
+    # 8. Training loop
+    def collate(batch):
+        xs = torch.stack([b[0] for b in batch])
+        ys = torch.stack([b[1] for b in batch])
+        ms = torch.stack([b[2] for b in batch])
+        return xs, ys, ms
+
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, collate_fn=collate, pin_memory=True,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    steps_per_epoch = max(1, len(loader) // args.grad_accum)
+    total_steps = steps_per_epoch * args.epochs
+    if args.max_steps:
+        total_steps = min(total_steps, args.max_steps)
+    warmup = max(20, int(args.warmup_frac * total_steps))
+
+    print(f"\n[train] LoRA rank={args.lora_rank} alpha={args.lora_alpha} "
+          f"scale={args.lora_alpha/args.lora_rank:.1f}")
+    print(f"[train] epochs={args.epochs} steps/epoch≈{steps_per_epoch} "
+          f"total={total_steps} warmup={warmup}")
+    print(f"[train] lr={args.lr} batch={args.batch_size} accum={args.grad_accum} "
+          f"effective_batch={args.batch_size * args.grad_accum}")
+
+    model.train()
+    t_start = time.time()
+    step = 0
+    running_loss = 0.0
+    running_n = 0
+
+    for ep in range(args.epochs):
+        print(f"\n=== epoch {ep+1}/{args.epochs} (LoRA tool-SFT) ===")
+        data_iter = iter(loader)
+
+        for _ in range(steps_per_epoch):
+            if args.max_steps and step >= args.max_steps:
+                break
+
+            cur_lr = cosine_with_warmup(step, warmup, total_steps, args.lr)
+            for g in optimizer.param_groups:
+                g["lr"] = cur_lr
+
+            optimizer.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+
+            for _micro in range(args.grad_accum):
+                try:
+                    xs, ys, ms = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(loader)
+                    xs, ys, ms = next(data_iter)
+
+                xs = xs.to(args.device, non_blocking=True)
+                ys = ys.to(args.device, non_blocking=True)
+                ms = ms.to(args.device, non_blocking=True)
+
+                with torch.amp.autocast("cuda", dtype=dtype, enabled=use_amp):
+                    _, loss = model(xs, targets=ys, loss_mask=ms)
+                    loss = loss / args.grad_accum
+                loss.backward()
+                loss_accum += loss.item() * args.grad_accum
+
+            gnorm = torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
+            optimizer.step()
+            step += 1
+            running_loss += loss_accum / args.grad_accum
+            running_n += 1
+
+            if step % args.log_every == 0:
+                elapsed = time.time() - t_start
+                avg = running_loss / running_n
+                print(f"[lora ep{ep+1} step {step:>4}/{total_steps}] "
+                      f"loss={avg:.4f} lr={cur_lr:.2e} "
+                      f"gnorm={gnorm:.2f} {elapsed/60:.1f}min")
+                log_jsonl(log_path, {"epoch": ep+1, "step": step, "loss": avg,
+                                     "lr": cur_lr, "gnorm": float(gnorm)})
+                running_loss = 0.0
+                running_n = 0
+
+            if step % args.save_every == 0:
+                save_lora_checkpoint(out_dir / "last_lora.pt", model, optimizer,
+                                     step, {"epoch": ep+1})
+
+        if args.max_steps and step >= args.max_steps:
+            break
+
+        save_lora_checkpoint(out_dir / f"epoch{ep+1}_lora.pt", model, optimizer,
+                             step, {"epoch": ep+1})
+        print(f"[save] epoch{ep+1}_lora.pt")
+
+    # Guardar checkpoint final con pesos COMPLETOS (base + LoRA merged)
+    # Estrategia: construir state_dict manualmente fusionando LoRA
+    print("\n[merge] Mergeando LoRA en pesos base...")
+
+    # Primero recolectar todos los módulos LoRA con sus rutas
+    lora_modules = {}
+    for mod_name, mod in model.named_modules():
+        if isinstance(mod, LoRALinear):
+            lora_modules[mod_name] = mod
+
+    # Construir state_dict fusionado
+    merged_state = {}
+    for param_name, param in model.named_parameters():
+        # Detectar si este parámetro pertenece a un LoRALinear
+        is_lora_internal = False
+        for lora_path in lora_modules:
+            if param_name.startswith(lora_path + ".lora_"):
+                is_lora_internal = True  # saltar lora_A y lora_B
+                break
+            if param_name == lora_path + ".linear.weight":
+                # Fusionar con LoRA
+                lora_mod = lora_modules[lora_path]
+                fused = param.data + (lora_mod.lora_B.data @ lora_mod.lora_A.data) * lora_mod.scale
+                # Guardar con nombre limpio (sin .linear)
+                clean = lora_path + ".weight"
+                merged_state[clean] = fused
+                is_lora_internal = True
+                break
+            if param_name == lora_path + ".linear.bias":
+                clean = lora_path + ".bias"
+                merged_state[clean] = param.data
+                is_lora_internal = True
+                break
+        if not is_lora_internal:
+            merged_state[param_name] = param.data
+
+    print(f"[merge] {len(merged_state)} keys en merged state_dict")
+
+    # Guardar solo LoRA ANTES de modificar el modelo
+    save_lora_checkpoint(out_dir / "final_lora_only.pt", model, optimizer,
+                         step, {"done": True, "lora_rank": args.lora_rank,
+                                "lora_alpha": args.lora_alpha})
+
+    # Guardar merged (full model) para benchmark — usar clave "model" que espera load_checkpoint
+    # strict=False en benchmark porque lm_head comparte pesos con tok_emb (tie_embeddings)
+    torch.save({"model": merged_state, "step": step,
+                "lora_rank": args.lora_rank, "lora_alpha": args.lora_alpha,
+                "merged": True, "tie_embeddings": True},
+               out_dir / "final.pt")
+    print(f"[done] final.pt (merged) → {out_dir}")
+    print(f"[done] final_lora_only.pt (adapter only) → {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

training/finetune_sft.py

@@ -0,0 +1,255 @@
+"""SFT fine-tuning with assistant-only loss masking and an internal mini-curriculum.
+
+Mini-curriculum (within SFT):
+  Epoch 1-2:  60% conversational (OASST1 ES + sft_conv) + 40% CVE Q&A
+  Epoch 3:    add tool-use (50% conv + 25% CVE + 25% tool_use)
+
+This avoids drowning the chat behavior in JSON tool-call patterns the way SFT v3 did.
+
+Run example:
+    python -m training_v2.train.finetune_sft \
+        --config training_v2/configs/nano.json \
+        --tokenizer training_v2/tokenizer/out/vectrayx_bpe.model \
+        --resume training_v2/checkpoints/phase3/last.pt \
+        --out training_v2/checkpoints/sft_v4 \
+        --batch-size 16 --grad-accum 4 --epochs 3 --lr 2e-5
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import sentencepiece as spm
+import torch
+from torch.utils.data import DataLoader, ConcatDataset
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from training_v2.data.sft_dataset import SFTDataset
+from training_v2.model.transformer import VectraYXNano, ModelConfig
+from training_v2.train.utils import (
+    cosine_with_warmup, make_optimizer, save_checkpoint, load_checkpoint, log_jsonl,
+)
+
+
+SFT_FILES = {
+    "conversational": [
+        "corpus/sft_conversational.jsonl",
+        "sft_v2_data/oasst1_es.jsonl",
+    ],
+    "cve_qa": [
+        "corpus/sft_v2_dataset.jsonl",
+    ],
+    "tool_use": [
+        "corpus/tooluse_dataset.jsonl",
+    ],
+}
+
+
+def load_sft_corpus_config(path):
+    global SFT_FILES
+    cfg = json.loads(Path(path).read_text())
+    SFT_FILES = {
+        "conversational": cfg.get("sft_conversational", SFT_FILES["conversational"]),
+        "cve_qa":         cfg.get("sft_cve_qa",         SFT_FILES["cve_qa"]),
+        "tool_use":       cfg.get("sft_tool_use",        SFT_FILES["tool_use"]),
+    }
+
+
+def discover(paths, root):
+    found = []
+    for rel in paths:
+        full = Path(root) / rel
+        if full.exists():
+            found.append(full)
+        else:
+            print(f"  [skip missing] {full}")
+    return found
+
+
+def build_dataset(args, sp, include_tools):
+    block_size = ModelConfig.from_json(args.config).max_seq_len
+    pad_id = sp.pad_id() if sp.pad_id() >= 0 else 0
+
+    conv = discover(SFT_FILES["conversational"], args.corpus_root)
+    cve = discover(SFT_FILES["cve_qa"], args.corpus_root)
+    tools = discover(SFT_FILES["tool_use"], args.corpus_root)
+
+    parts = []
+    if conv:
+        parts.append(("conv", SFTDataset(conv, sp, block_size, pad_id=pad_id, seed=args.seed)))
+    if cve:
+        parts.append(("cve", SFTDataset(cve, sp, block_size, pad_id=pad_id, seed=args.seed + 1)))
+    if include_tools and tools:
+        parts.append(("tools", SFTDataset(tools, sp, block_size, pad_id=pad_id, seed=args.seed + 2)))
+    return parts, pad_id
+
+
+def make_loader(parts, weights, batch_size, num_workers):
+    """Weighted sampling across the named parts."""
+    sizes = [len(d) for _, d in parts]
+    names = [n for n, _ in parts]
+    datasets = [d for _, d in parts]
+    big = ConcatDataset(datasets)
+
+    offsets = np.cumsum([0] + sizes)
+    weight_per_idx = np.zeros(offsets[-1], dtype=np.float64)
+    for i, n in enumerate(names):
+        w = weights.get(n, 1.0) / max(1, sizes[i])
+        weight_per_idx[offsets[i]:offsets[i + 1]] = w
+    sampler = torch.utils.data.WeightedRandomSampler(
+        weights=weight_per_idx,
+        num_samples=int(sum(sizes)),
+        replacement=True,
+    )
+
+    def collate(batch):
+        xs = torch.stack([b[0] for b in batch], 0)
+        ys = torch.stack([b[1] for b in batch], 0)
+        ms = torch.stack([b[2] for b in batch], 0)
+        return xs, ys, ms
+
+    return DataLoader(
+        big, batch_size=batch_size, sampler=sampler,
+        num_workers=num_workers, collate_fn=collate, pin_memory=True,
+        persistent_workers=num_workers > 0,
+    )
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", required=True)
+    p.add_argument("--tokenizer", required=True)
+    p.add_argument("--resume", required=True, help="pre-training checkpoint to fine-tune")
+    p.add_argument("--out", required=True)
+    p.add_argument("--corpus-root", default=".")
+    p.add_argument("--corpus-config", default=None)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--grad-accum", type=int, default=4)
+    p.add_argument("--epochs", type=int, default=3)
+    p.add_argument("--lr", type=float, default=2e-5)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--grad-clip", type=float, default=1.0)
+    p.add_argument("--warmup-frac", type=float, default=0.03)
+    p.add_argument("--num-workers", type=int, default=2)
+    p.add_argument("--log-every", type=int, default=20)
+    p.add_argument("--save-every", type=int, default=500)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    args = p.parse_args()
+
+    if args.corpus_config:
+        load_sft_corpus_config(args.corpus_config)
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    cfg = ModelConfig.from_json(args.config)
+    model = VectraYXNano(cfg).to(args.device)
+    print(f"[model] {model.num_params()/1e6:.2f}M params")
+    load_checkpoint(args.resume, model, optimizer=None, map_location=args.device)
+    print(f"[resume] {args.resume}")
+
+    sp = spm.SentencePieceProcessor()
+    sp.load(args.tokenizer)
+    parts, pad_id = build_dataset(args, sp, include_tools=True)
+    if not parts:
+        raise RuntimeError("no SFT files found")
+
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    log_path = out_dir / "train_log.jsonl"
+
+    optimizer = make_optimizer(model, lr=args.lr, weight_decay=args.weight_decay)
+
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    use_amp = args.device == "cuda" and dtype != torch.float32
+
+    epoch_plans = [
+        {"conv": 1.00, "cve": 0.00, "tools": 0.0},   # epoch 1: SOLO conversacional
+        {"conv": 0.70, "cve": 0.30, "tools": 0.00},  # epoch 2: + CVE Q&A
+        {"conv": 0.55, "cve": 0.30, "tools": 0.15},  # epoch 3: + tool use
+    ]
+
+    total_steps = 0
+    for ep in range(args.epochs):
+        weights = epoch_plans[min(ep, len(epoch_plans) - 1)]
+        print(f"\n=== epoch {ep+1}/{args.epochs} | mix={weights} ===")
+        loader = make_loader(parts, weights, args.batch_size, args.num_workers)
+        steps_per_epoch = max(1, len(loader) // args.grad_accum)
+        total_steps += steps_per_epoch
+    warmup = max(50, int(args.warmup_frac * total_steps))
+    print(f"[sft] total_steps≈{total_steps}  warmup={warmup}")
+
+    model.train()
+    t_start = time.time()
+    step = 0
+    running_loss = 0.0
+    running_n = 0
+
+    for ep in range(args.epochs):
+        weights = epoch_plans[min(ep, len(epoch_plans) - 1)]
+        loader = make_loader(parts, weights, args.batch_size, args.num_workers)
+        data_iter = iter(loader)
+        steps_per_epoch = max(1, len(loader) // args.grad_accum)
+
+        for _ in range(steps_per_epoch):
+            cur_lr = cosine_with_warmup(step, warmup, total_steps, args.lr)
+            for g in optimizer.param_groups:
+                g["lr"] = cur_lr
+
+            optimizer.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+            for _micro in range(args.grad_accum):
+                try:
+                    xs, ys, ms = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(loader)
+                    xs, ys, ms = next(data_iter)
+                xs = xs.to(args.device, non_blocking=True)
+                ys = ys.to(args.device, non_blocking=True)
+                ms = ms.to(args.device, non_blocking=True)
+                with torch.amp.autocast("cuda", dtype=dtype, enabled=use_amp):
+                    _, loss = model(xs, targets=ys, loss_mask=ms)
+                    loss = loss / args.grad_accum
+                loss.backward()
+                loss_accum += loss.item() * args.grad_accum
+
+            gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+            step += 1
+            running_loss += loss_accum / args.grad_accum
+            running_n += 1
+
+            if step % args.log_every == 0:
+                elapsed = time.time() - t_start
+                avg = running_loss / running_n
+                print(f"[sft ep{ep+1} step {step:>5}/{total_steps}] loss={avg:.4f} "
+                      f"lr={cur_lr:.2e} gnorm={gnorm:.2f} elapsed={elapsed/60:.1f}min")
+                log_jsonl(log_path, {"epoch": ep + 1, "step": step, "loss": avg,
+                                     "lr": cur_lr, "gnorm": float(gnorm)})
+                running_loss = 0.0
+                running_n = 0
+
+            if step % args.save_every == 0:
+                save_checkpoint(out_dir / "last.pt", model, optimizer,
+                                {"step": step}, step,
+                                extra={"epoch": ep + 1, "weights": weights})
+
+        save_checkpoint(out_dir / f"epoch{ep+1}.pt", model, optimizer,
+                        {"step": step}, step,
+                        extra={"epoch": ep + 1, "weights": weights})
+        print(f"[save] {out_dir}/epoch{ep+1}.pt")
+
+    save_checkpoint(out_dir / "final.pt", model, optimizer, {"step": step}, step,
+                    extra={"done": True})
+    print(f"[done] SFT → {out_dir}/final.pt")
+
+
+if __name__ == "__main__":
+    main()

training/finetune_tools.py

@@ -0,0 +1,188 @@
+"""Tool-use focused SFT for VectraYX Nano/Base.
+
+This is a simplified version of finetune_sft.py that trains ONLY on tool-use examples.
+The goal is to test the hypothesis that B4=0.000 is due to diluted tool-call gradients
+in the mixed SFT corpus, not a capacity gate.
+
+Run example:
+    python -m training_v2.train.finetune_tools \
+        --config training_v2/configs/nano.json \
+        --tokenizer models/vectrayx_bpe.model \
+        --resume checkpoints/nano_final.pt \
+        --tool-corpus /tmp/tool_sft_v1.jsonl \
+        --out checkpoints/tool_sft_nano \
+        --batch-size 16 --grad-accum 4 --epochs 2 --lr 1e-5
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import sentencepiece as spm
+import torch
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from training_v2.data.sft_dataset import SFTDataset
+from training_v2.model.transformer import VectraYXNano, ModelConfig
+from training_v2.train.utils import (
+    cosine_with_warmup, make_optimizer, save_checkpoint, load_checkpoint, log_jsonl,
+)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", required=True)
+    p.add_argument("--tokenizer", required=True)
+    p.add_argument("--resume", required=True, help="checkpoint to fine-tune from")
+    p.add_argument("--tool-corpus", required=True, help="tool-use JSONL corpus")
+    p.add_argument("--out", required=True)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--grad-accum", type=int, default=4)
+    p.add_argument("--epochs", type=int, default=2)
+    p.add_argument("--lr", type=float, default=1e-5)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--grad-clip", type=float, default=1.0)
+    p.add_argument("--warmup-frac", type=float, default=0.03)
+    p.add_argument("--num-workers", type=int, default=2)
+    p.add_argument("--log-every", type=int, default=20)
+    p.add_argument("--save-every", type=int, default=500)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    p.add_argument("--max-steps", type=int, default=None, help="for testing")
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # Load model
+    cfg = ModelConfig.from_json(args.config)
+    model = VectraYXNano(cfg).to(args.device)
+    print(f"[model] {model.num_params()/1e6:.2f}M params")
+    load_checkpoint(args.resume, model, optimizer=None, map_location=args.device)
+    print(f"[resume] {args.resume}")
+
+    # Load tokenizer
+    sp = spm.SentencePieceProcessor()
+    sp.load(args.tokenizer)
+    pad_id = sp.pad_id() if sp.pad_id() >= 0 else 0
+
+    # Build tool-only dataset
+    block_size = cfg.max_seq_len
+    tool_corpus = Path(args.tool_corpus)
+    if not tool_corpus.exists():
+        raise FileNotFoundError(f"Tool corpus not found: {tool_corpus}")
+    
+    dataset = SFTDataset([tool_corpus], sp, block_size, pad_id=pad_id, seed=args.seed)
+    print(f"[dataset] {len(dataset)} tool-use examples from {tool_corpus}")
+
+    # Setup output
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    log_path = out_dir / "train_log.jsonl"
+
+    # Optimizer
+    optimizer = make_optimizer(model, lr=args.lr, weight_decay=args.weight_decay)
+
+    # AMP setup
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    use_amp = args.device == "cuda" and dtype != torch.float32
+
+    # Training loop
+    def collate(batch):
+        xs = torch.stack([b[0] for b in batch], 0)
+        ys = torch.stack([b[1] for b in batch], 0)
+        ms = torch.stack([b[2] for b in batch], 0)
+        return xs, ys, ms
+
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, collate_fn=collate, pin_memory=True,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    steps_per_epoch = max(1, len(loader) // args.grad_accum)
+    total_steps = steps_per_epoch * args.epochs
+    if args.max_steps:
+        total_steps = min(total_steps, args.max_steps)
+    warmup = max(50, int(args.warmup_frac * total_steps))
+    print(f"[train] epochs={args.epochs} steps_per_epoch≈{steps_per_epoch} total_steps={total_steps} warmup={warmup}")
+
+    model.train()
+    t_start = time.time()
+    step = 0
+    running_loss = 0.0
+    running_n = 0
+
+    for ep in range(args.epochs):
+        print(f"\n=== epoch {ep+1}/{args.epochs} (tool-only) ===")
+        data_iter = iter(loader)
+
+        for _ in range(steps_per_epoch):
+            if args.max_steps and step >= args.max_steps:
+                break
+
+            cur_lr = cosine_with_warmup(step, warmup, total_steps, args.lr)
+            for g in optimizer.param_groups:
+                g["lr"] = cur_lr
+
+            optimizer.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+            for _micro in range(args.grad_accum):
+                try:
+                    xs, ys, ms = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(loader)
+                    xs, ys, ms = next(data_iter)
+                xs = xs.to(args.device, non_blocking=True)
+                ys = ys.to(args.device, non_blocking=True)
+                ms = ms.to(args.device, non_blocking=True)
+                
+                with torch.amp.autocast("cuda", dtype=dtype, enabled=use_amp):
+                    _, loss = model(xs, targets=ys, loss_mask=ms)
+                    loss = loss / args.grad_accum
+                loss.backward()
+                loss_accum += loss.item() * args.grad_accum
+
+            gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+            step += 1
+            running_loss += loss_accum / args.grad_accum
+            running_n += 1
+
+            if step % args.log_every == 0:
+                elapsed = time.time() - t_start
+                avg = running_loss / running_n
+                print(f"[tool-sft ep{ep+1} step {step:>5}/{total_steps}] loss={avg:.4f} "
+                      f"lr={cur_lr:.2e} gnorm={gnorm:.2f} elapsed={elapsed/60:.1f}min")
+                log_jsonl(log_path, {"epoch": ep + 1, "step": step, "loss": avg,
+                                     "lr": cur_lr, "gnorm": float(gnorm)})
+                running_loss = 0.0
+                running_n = 0
+
+            if step % args.save_every == 0:
+                save_checkpoint(out_dir / "last.pt", model, optimizer,
+                                {"step": step}, step,
+                                extra={"epoch": ep + 1, "tool_only": True})
+
+        if args.max_steps and step >= args.max_steps:
+            break
+
+        save_checkpoint(out_dir / f"epoch{ep+1}.pt", model, optimizer,
+                        {"step": step}, step,
+                        extra={"epoch": ep + 1, "tool_only": True})
+        print(f"[save] {out_dir}/epoch{ep+1}.pt")
+
+    save_checkpoint(out_dir / "final.pt", model, optimizer, {"step": step}, step,
+                    extra={"done": True, "tool_only": True})
+    print(f"[done] {out_dir}/final.pt")
+
+
+if __name__ == "__main__":
+    main()

training/pretrain.py

@@ -0,0 +1,228 @@
+"""Curriculum pre-training driver for VectraYX-Nano v2.
+
+Phase 1: 100% conversational  (LR 3e-4 from scratch)
+Phase 2: 75% tech + 25% conv  (LR 1.5e-4, resumed from phase 1)
+Phase 3: 70% tools + 20% tech + 10% conv  (LR 8e-5, resumed from phase 2)
+
+Run example:
+    python -m training_v2.train.pretrain \
+        --config training_v2/configs/nano.json \
+        --bins training_v2/data/bins \
+        --out training_v2/checkpoints \
+        --phase 1 --max-steps 8000 --batch-size 16 --grad-accum 8
+
+Then:
+    --phase 2 --resume training_v2/checkpoints/phase1/last.pt
+    --phase 3 --resume training_v2/checkpoints/phase2/last.pt
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from training_v2.data.curriculum_dataset import (
+    MixedCurriculumDataset, make_phase_mix, load_phase_summary,
+)
+from training_v2.model.transformer import VectraYXNano, ModelConfig
+from training_v2.train.utils import (
+    cosine_with_warmup, make_optimizer, save_checkpoint, load_checkpoint, log_jsonl,
+)
+
+
+PHASE_LR = {1: 3.0e-4, 2: 1.5e-4, 3: 8.0e-5}
+PHASE_WARMUP_FRAC = {1: 0.05, 2: 0.02, 3: 0.02}
+
+
+def build_dataloader(args, mix, block_size):
+    phase_dirs = {
+        "phase1_conv": Path(args.bins) / "phase1_conv",
+        "phase2_tech": Path(args.bins) / "phase2_tech",
+        "phase3_tools": Path(args.bins) / "phase3_tools",
+    }
+    ds = MixedCurriculumDataset(
+        phase_dirs={k: v for k, v in phase_dirs.items() if mix.get(k, 0) > 0},
+        weights=mix,
+        block_size=block_size,
+        dtype=np.uint16,
+        seed=args.seed,
+    )
+
+    def collate(batch):
+        xs = torch.stack([b[0] for b in batch], 0)
+        ys = torch.stack([b[1] for b in batch], 0)
+        return xs, ys
+
+    return DataLoader(
+        ds,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        collate_fn=collate,
+        pin_memory=True,
+        persistent_workers=args.num_workers > 0,
+    )
+
+
+def estimate_phase_tokens(phase_idx, mix, summary):
+    total = 0.0
+    for k, w in mix.items():
+        n = summary.get(k, {}).get("n_tokens", 0)
+        if w > 0 and n > 0:
+            total += n
+    return int(total)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", required=True)
+    p.add_argument("--bins", required=True, help="root of binary shard dirs")
+    p.add_argument("--out", required=True, help="checkpoint output root")
+    p.add_argument("--phase", type=int, choices=[1, 2, 3], required=True)
+    p.add_argument("--resume", type=str, default=None)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--grad-accum", type=int, default=8)
+    p.add_argument("--max-steps", type=int, default=None)
+    p.add_argument("--epochs", type=float, default=2.0,
+                   help="estimate steps as epochs*phase_tokens/(batch*ga*block)")
+    p.add_argument("--lr", type=float, default=None)
+    p.add_argument("--weight-decay", type=float, default=0.1)
+    p.add_argument("--grad-clip", type=float, default=1.0)
+    p.add_argument("--num-workers", type=int, default=2)
+    p.add_argument("--log-every", type=int, default=20)
+    p.add_argument("--save-every", type=int, default=1000)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--replay-conv", type=float, default=None,
+                   help="override replay ratio of conversational data in phase 2/3")
+    p.add_argument("--replay-tech", type=float, default=None,
+                   help="override replay ratio of technical data in phase 3")
+    p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--dtype", type=str, default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    p.add_argument("--compile", action="store_true")
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    cfg = ModelConfig.from_json(args.config)
+    model = VectraYXNano(cfg).to(args.device)
+    n_params = model.num_params()
+    print(f"[model] {n_params/1e6:.2f}M params · cfg={cfg}")
+
+    mix = make_phase_mix(args.phase, replay_conv=args.replay_conv, replay_tech=args.replay_tech)
+    summary = load_phase_summary(args.bins)
+    phase_tokens = estimate_phase_tokens(args.phase, mix, summary)
+    tokens_per_step = args.batch_size * args.grad_accum * cfg.max_seq_len
+    if args.max_steps is None:
+        args.max_steps = max(1000, int(args.epochs * phase_tokens / tokens_per_step))
+    print(f"[phase {args.phase}] mix={mix}")
+    print(f"[phase {args.phase}] phase_tokens={phase_tokens:,}  tokens/step={tokens_per_step:,}  steps={args.max_steps}")
+
+    lr = args.lr if args.lr is not None else PHASE_LR[args.phase]
+    warmup = max(50, int(PHASE_WARMUP_FRAC[args.phase] * args.max_steps))
+    optimizer = make_optimizer(model, lr=lr, weight_decay=args.weight_decay)
+
+    start_step = 0
+    if args.resume:
+        start_step, _ = load_checkpoint(args.resume, model, optimizer=None, map_location=args.device)
+        print(f"[resume] loaded weights from {args.resume} (step={start_step})")
+        start_step = 0  # fresh optimizer for new phase
+
+    loader = build_dataloader(args, mix, cfg.max_seq_len)
+    data_iter = iter(loader)
+
+    out_dir = Path(args.out) / f"phase{args.phase}"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    log_path = out_dir / "train_log.jsonl"
+
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    use_amp = args.device == "cuda" and dtype != torch.float32
+    scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
+
+    if args.compile:
+        try:
+            model = torch.compile(model)
+        except Exception as e:
+            print(f"[compile] skipped: {e}")
+
+    model.train()
+    t_start = time.time()
+    tokens_seen = 0
+    running_loss = 0.0
+    running_n = 0
+
+    for step in range(start_step, args.max_steps):
+        cur_lr = cosine_with_warmup(step, warmup, args.max_steps, lr)
+        for g in optimizer.param_groups:
+            g["lr"] = cur_lr
+
+        optimizer.zero_grad(set_to_none=True)
+        loss_accum = 0.0
+        for micro in range(args.grad_accum):
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(loader)
+                batch = next(data_iter)
+            xs, ys = batch[0], batch[1]
+            xs = xs.to(args.device, non_blocking=True)
+            ys = ys.to(args.device, non_blocking=True)
+            with torch.amp.autocast("cuda", dtype=dtype, enabled=use_amp):
+                _, loss = model(xs, targets=ys)
+                loss = loss / args.grad_accum
+            if scaler.is_enabled():
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+            loss_accum += loss.item() * args.grad_accum
+
+        if scaler.is_enabled():
+            scaler.unscale_(optimizer)
+        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        if scaler.is_enabled():
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            optimizer.step()
+
+        tokens_seen += tokens_per_step
+        running_loss += loss_accum / args.grad_accum
+        running_n += 1
+
+        if (step + 1) % args.log_every == 0:
+            elapsed = time.time() - t_start
+            tps = tokens_seen / max(1.0, elapsed)
+            avg_loss = running_loss / running_n
+            print(f"[p{args.phase} step {step+1:>6}/{args.max_steps}] "
+                  f"loss={avg_loss:.4f} lr={cur_lr:.2e} gnorm={gnorm:.2f} "
+                  f"tok/s={tps:>7,.0f} elapsed={elapsed/60:.1f}min")
+            log_jsonl(log_path, {
+                "phase": args.phase, "step": step + 1, "loss": avg_loss,
+                "lr": cur_lr, "gnorm": float(gnorm), "tok_per_s": tps,
+                "tokens_seen": tokens_seen,
+            })
+            running_loss = 0.0
+            running_n = 0
+
+        if (step + 1) % args.save_every == 0 or (step + 1) == args.max_steps:
+            ckpt_path = out_dir / "last.pt"
+            save_checkpoint(ckpt_path, model, optimizer, {"step": step + 1}, step + 1,
+                            extra={"phase": args.phase, "mix": mix, "lr": lr})
+            print(f"[save] {ckpt_path}")
+
+    final = out_dir / "last.pt"
+    save_checkpoint(final, model, optimizer, {"step": args.max_steps}, args.max_steps,
+                    extra={"phase": args.phase, "mix": mix, "lr": lr, "done": True})
+    print(f"[done] phase {args.phase} → {final}")
+
+
+if __name__ == "__main__":
+    main()

training/sft_dataset.py

@@ -0,0 +1,105 @@
+"""SFT dataset with proper assistant-only loss masking and safe packing.
+
+Each example is a chat-formatted string with `<|system|> <|user|> <|assistant|> <|end|>`
+turn delimiters.  We tokenize on the fly (corpus is small, ~25M tokens) and build a
+mask=1 only on tokens that are part of an assistant response (everything between
+`<|assistant|>` and the next `<|end|>`).
+
+For pre-training-style packing without cross-example contamination we group multiple
+short examples into a fixed-length window using `cu_seqlens`-style document boundaries
+implemented via per-document attention reset.  Here we keep it simple: pad/truncate
+each example to `block_size`. Throughput is still high (>40k tok/s on L4) for this
+volume.
+"""
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+def _read_jsonl(path):
+    out = []
+    with open(path, "r", encoding="utf-8", errors="replace") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            t = obj.get("text") or ""
+            if t:
+                out.append({"text": t, "source": obj.get("source", Path(path).stem)})
+    return out
+
+
+def build_assistant_mask(token_ids, assistant_id, end_id):
+    """mask[i] = 1 iff token_ids[i] is inside an `<|assistant|> ... <|end|>` span.
+
+    We mark from the token AFTER `<|assistant|>` up to and including `<|end|>` so the
+    model learns to emit the closing delimiter.
+    """
+    mask = np.zeros(len(token_ids), dtype=np.int64)
+    inside = False
+    for i, t in enumerate(token_ids):
+        if t == assistant_id and not inside:
+            inside = True
+            continue  # don't include the assistant tag itself
+        if inside:
+            mask[i] = 1
+            if t == end_id:
+                inside = False
+    return mask
+
+
+class SFTDataset(Dataset):
+    def __init__(self, jsonl_paths, sp, block_size, assistant_token="<|assistant|>",
+                 end_token="<|end|>", pad_id=0, seed=42, mix_weights=None):
+        self.sp = sp
+        self.block_size = block_size
+        self.pad_id = pad_id
+        self.assistant_id = sp.piece_to_id(assistant_token)
+        self.end_id = sp.piece_to_id(end_token)
+        if self.assistant_id < 0 or self.end_id < 0:
+            raise ValueError(f"missing special tokens in tokenizer: "
+                             f"{assistant_token}={self.assistant_id} {end_token}={self.end_id}")
+
+        self.examples = []
+        rng = random.Random(seed)
+        for p in jsonl_paths:
+            recs = _read_jsonl(p)
+            w = (mix_weights or {}).get(Path(p).name, 1.0)
+            if w != 1.0:
+                k = int(len(recs) * w)
+                recs = rng.sample(recs, min(k, len(recs)))
+            self.examples.extend(recs)
+            print(f"  [sft] {p}: {len(recs):,} ex (w={w})")
+        rng.shuffle(self.examples)
+        print(f"[sft] total: {len(self.examples):,} examples")
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        text = self.examples[idx]["text"]
+        ids = self.sp.encode(text, out_type=int)
+        ids = ids[: self.block_size + 1]
+        mask = build_assistant_mask(ids, self.assistant_id, self.end_id)
+
+        if len(ids) < self.block_size + 1:
+            need = self.block_size + 1 - len(ids)
+            ids = ids + [self.pad_id] * need
+            mask = np.concatenate([mask, np.zeros(need, dtype=np.int64)])
+
+        ids = np.asarray(ids, dtype=np.int64)
+        x = torch.from_numpy(ids[:-1])
+        y = torch.from_numpy(ids[1:].copy())
+        m = torch.from_numpy(mask[1:].copy())  # mask aligned with targets
+        # zero out padded targets
+        y[m == 0] = -100
+        return x, y, m

training/transformer.py

@@ -0,0 +1,259 @@
+"""VectraYX-Nano transformer (decoder-only, ~42M params).
+
+Modern small-LLM stack:
+  RMSNorm (pre-norm)  ·  SwiGLU FFN  ·  RoPE  ·  GQA (8q/2kv)
+  QK-Norm  ·  no biases  ·  tied embeddings  ·  z-loss
+"""
+
+import json
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 16384
+    n_layers: int = 8
+    n_heads: int = 8
+    n_kv_heads: int = 2
+    d_model: int = 512
+    d_ffn: int = 2048
+    max_seq_len: int = 1024
+    rope_theta: float = 10000.0
+    rms_eps: float = 1e-6
+    init_std: float = 0.02
+    dropout: float = 0.0
+    tie_embeddings: bool = True
+    qk_norm: bool = True
+    z_loss_coef: float = 1e-4
+
+    @classmethod
+    def from_json(cls, path):
+        cfg = json.loads(open(path).read())["model"]
+        return cls(**{k: cfg[k] for k in cfg if k in cls.__dataclass_fields__})
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return x.to(self.weight.dtype) * self.weight
+
+
+def precompute_rope(head_dim, max_seq_len, theta=10000.0, device=None):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)
+    cos = freqs.cos()
+    sin = freqs.sin()
+    if device is not None:
+        cos = cos.to(device)
+        sin = sin.to(device)
+    return cos, sin
+
+
+def apply_rope(x, cos, sin):
+    # x: (B, H, T, D) with D even.  cos/sin: (T, D/2)
+    T, D = x.shape[-2], x.shape[-1]
+    cos = cos[:T].view(1, 1, T, D // 2)
+    sin = sin[:T].view(1, 1, T, D // 2)
+    x1 = x[..., : D // 2]
+    x2 = x[..., D // 2:]
+    rx1 = x1 * cos - x2 * sin
+    rx2 = x1 * sin + x2 * cos
+    return torch.cat([rx1, rx2], dim=-1)
+
+
+class GQAttention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        assert cfg.d_model % cfg.n_heads == 0
+        assert cfg.n_heads % cfg.n_kv_heads == 0
+        self.n_heads = cfg.n_heads
+        self.n_kv_heads = cfg.n_kv_heads
+        self.head_dim = cfg.d_model // cfg.n_heads
+        self.repeat = self.n_heads // self.n_kv_heads
+
+        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.d_model, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.d_model, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+
+        self.qk_norm = cfg.qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, eps=cfg.rms_eps)
+            self.k_norm = RMSNorm(self.head_dim, eps=cfg.rms_eps)
+
+        self.dropout = cfg.dropout
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        if self.repeat > 1:
+            k = k.repeat_interleave(self.repeat, dim=1)
+            v = v.repeat_interleave(self.repeat, dim=1)
+
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=True,
+        )
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.w_gate = nn.Linear(cfg.d_model, cfg.d_ffn, bias=False)
+        self.w_up = nn.Linear(cfg.d_model, cfg.d_ffn, bias=False)
+        self.w_down = nn.Linear(cfg.d_ffn, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.d_model, eps=cfg.rms_eps)
+        self.attn = GQAttention(cfg)
+        self.ffn_norm = RMSNorm(cfg.d_model, eps=cfg.rms_eps)
+        self.ffn = SwiGLU(cfg)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class VectraYXNano(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
+        self.final_norm = RMSNorm(cfg.d_model, eps=cfg.rms_eps)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+
+        if cfg.tie_embeddings:
+            self.lm_head.weight = self.tok_emb.weight
+
+        head_dim = cfg.d_model // cfg.n_heads
+        cos, sin = precompute_rope(head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+        self.apply(self._init_weights)
+        residual_std = cfg.init_std / math.sqrt(2 * cfg.n_layers)
+        for n, p in self.named_parameters():
+            if n.endswith("wo.weight") or n.endswith("w_down.weight"):
+                nn.init.normal_(p, mean=0.0, std=residual_std)
+
+    def _init_weights(self, m):
+        std = self.cfg.init_std
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=std)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=std)
+
+    def num_params(self, exclude_embedding=False):
+        n = sum(p.numel() for p in self.parameters())
+        if exclude_embedding and self.cfg.tie_embeddings:
+            n -= self.tok_emb.weight.numel()
+        return n
+
+    def forward(self, idx, targets=None, loss_mask=None):
+        B, T = idx.shape
+        assert T <= self.cfg.max_seq_len, f"seq {T} > max {self.cfg.max_seq_len}"
+        x = self.tok_emb(idx)
+        cos = self.rope_cos
+        sin = self.rope_sin
+        for layer in self.layers:
+            x = layer(x, cos, sin)
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        if targets is None:
+            return logits, None
+
+        # cross-entropy + z-loss for stability
+        flat_logits = logits.view(-1, logits.size(-1))
+        flat_tgt = targets.view(-1)
+        ce = F.cross_entropy(flat_logits, flat_tgt, reduction="none", ignore_index=-100)
+        if loss_mask is not None:
+            mask = loss_mask.view(-1).float()
+            denom = mask.sum().clamp_min(1.0)
+            ce_loss = (ce * mask).sum() / denom
+        else:
+            valid = (flat_tgt != -100).float()
+            denom = valid.sum().clamp_min(1.0)
+            ce_loss = (ce * valid).sum() / denom
+
+        if self.cfg.z_loss_coef > 0:
+            lse = torch.logsumexp(flat_logits.float(), dim=-1)
+            if loss_mask is not None:
+                z = ((lse ** 2) * loss_mask.view(-1).float()).sum() / denom
+            else:
+                z = ((lse ** 2) * (flat_tgt != -100).float()).sum() / denom
+            loss = ce_loss + self.cfg.z_loss_coef * z
+        else:
+            loss = ce_loss
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=0.7, top_k=40, top_p=0.9,
+                 eos_id=None, repeat_penalty=1.0):
+        self.eval()
+        for _ in range(max_new_tokens):
+            cond = idx[:, -self.cfg.max_seq_len:]
+            logits, _ = self(cond)
+            logits = logits[:, -1, :].float()
+
+            if repeat_penalty != 1.0:
+                for token in set(idx[0].tolist()):
+                    logits[0, token] = logits[0, token] / repeat_penalty if logits[0, token] > 0 else logits[0, token] * repeat_penalty
+
+            if temperature <= 0:
+                next_id = logits.argmax(-1, keepdim=True)
+            else:
+                logits = logits / temperature
+                if top_k:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = -float("inf")
+                if top_p and top_p < 1.0:
+                    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+                    probs = F.softmax(sorted_logits, dim=-1)
+                    cumprobs = probs.cumsum(-1)
+                    drop = cumprobs > top_p
+                    drop[..., 1:] = drop[..., :-1].clone()
+                    drop[..., 0] = False
+                    sorted_logits[drop] = -float("inf")
+                    logits = torch.full_like(logits, -float("inf")).scatter(-1, sorted_idx, sorted_logits)
+                probs = F.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, 1)
+            idx = torch.cat([idx, next_id], dim=-1)
+            if eos_id is not None and next_id.item() == eos_id:
+                break
+        return idx

training/utils.py

@@ -0,0 +1,81 @@
+"""Training utilities: optimizer setup, LR schedule, checkpointing."""
+
+import json
+import math
+import os
+from pathlib import Path
+
+import torch
+
+
+def cosine_with_warmup(step, warmup, total, max_lr, min_lr_ratio=0.1):
+    if step < warmup:
+        return max_lr * (step + 1) / warmup
+    progress = (step - warmup) / max(1, total - warmup)
+    progress = min(1.0, progress)
+    return min_lr_ratio * max_lr + 0.5 * (max_lr - min_lr_ratio * max_lr) * (1 + math.cos(math.pi * progress))
+
+
+def make_optimizer(model, lr, weight_decay=0.1, betas=(0.9, 0.95), fused=True):
+    """AdamW with weight decay only on 2D weights (no decay on biases / norms / embeddings).
+
+    Per Loshchilov & Hutter; same convention as nanoGPT.
+    """
+    decay, no_decay = [], []
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.dim() >= 2 and "tok_emb" not in n:
+            decay.append(p)
+        else:
+            no_decay.append(p)
+    groups = [
+        {"params": decay, "weight_decay": weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+    extra = {}
+    if fused and torch.cuda.is_available():
+        try:
+            return torch.optim.AdamW(groups, lr=lr, betas=betas, fused=True)
+        except TypeError:
+            pass
+    return torch.optim.AdamW(groups, lr=lr, betas=betas, **extra)
+
+
+def save_checkpoint(path, model, optimizer, scheduler_state, step, extra=None):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict() if optimizer is not None else None,
+        "scheduler": scheduler_state,
+        "step": step,
+        "config": {k: getattr(model.cfg, k) for k in model.cfg.__dataclass_fields__},
+        "extra": extra or {},
+    }
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    torch.save(payload, tmp)
+    os.replace(tmp, path)
+
+
+def load_checkpoint(path, model, optimizer=None, map_location="cpu"):
+    payload = torch.load(path, map_location=map_location, weights_only=False)
+    # Si el checkpoint tiene tie_embeddings=True, usar strict=False
+    # (lm_head comparte pesos con tok_emb y no se guarda por separado)
+    strict = not payload.get("tie_embeddings", False)
+    missing, unexpected = model.load_state_dict(payload["model"], strict=strict)
+    if missing:
+        print(f"[load_checkpoint] missing keys (expected with tie_embeddings): {missing[:3]}")
+    if optimizer is not None and payload.get("optimizer"):
+        optimizer.load_state_dict(payload["optimizer"])
+    return payload.get("step", 0), payload.get("extra", {})
+
+
+def count_tokens(loader_output_iter, n_steps, block_size, batch_size):
+    """Approximate; effective tokens consumed per step."""
+    return n_steps * block_size * batch_size
+
+
+def log_jsonl(path, record):
+    with open(path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(record, ensure_ascii=False) + "\n")