sft_train.py

"""
LUNA 100M — SFT Fine-Tuning Script
====================================
Fine-tunes the pretrained LUNA-100M on instruction-following (SFT) data.

Features:
  - Loads pretrained checkpoint (latest.pt from pretraining)
  - JSON-based SFT dataset (instruction/input/output format)
  - Prompt masking: loss computed only on the output portion
  - Checkpoint eval: runs identity + knowledge prompts after each save
  - Cosine LR with warmup
  - Auto hardware detection (same as train.py)

Usage:
    python sft_train.py                                    # uses sft_config.yaml
    python sft_train.py --config sft_config.yaml           # explicit config
    python sft_train.py --train_json /data/train.json      # override data path
"""

import os
import gc
import sys
import math
import time
import json
import argparse
import yaml
import psutil
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import autocast, GradScaler
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


# ─── Model (identical to train.py) ───────────────────────────────────────────

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=1024):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_seq_len).float()
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos())
        self.register_buffer("sin_cached", emb.sin())

    def forward(self, seq_len):
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary(x, cos, sin):
    c = cos.unsqueeze(0).unsqueeze(0)
    s = sin.unsqueeze(0).unsqueeze(0)
    return x * c + rotate_half(x) * s


class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.rot_dim = int(self.head_dim * rotary_pct)
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
        self.rotary = RotaryEmbedding(self.rot_dim, block_size)

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        cos, sin = self.rotary(T)
        q = torch.cat([apply_rotary(q[..., :self.rot_dim], cos, sin), q[..., self.rot_dim:]], dim=-1)
        k = torch.cat([apply_rotary(k[..., :self.rot_dim], cos, sin), k[..., self.rot_dim:]], dim=-1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))


class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
        self.gelu = nn.GELU()
        self.proj = nn.Linear(4 * n_embd, n_embd, bias=True)

    def forward(self, x):
        return self.proj(self.gelu(self.fc(x)))


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class LUNAModel(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_embd, n_head):
        super().__init__()
        self.block_size = block_size
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight  # tied

    def _init_weights(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            m.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                m.bias.data.zero_()

    def forward(self, idx, targets=None, loss_mask=None, return_logits=True):
        x = self.wte(idx)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_targets = targets[:, 1:].contiguous()
            if loss_mask is not None:
                shift_mask = loss_mask[:, 1:].contiguous()
                # Only compute loss on output tokens
                flat_logits = shift_logits.view(-1, shift_logits.size(-1))
                flat_targets = shift_targets.view(-1)
                flat_mask = shift_mask.view(-1).float()
                per_token_loss = F.cross_entropy(flat_logits, flat_targets, reduction='none')
                loss = (per_token_loss * flat_mask).sum() / flat_mask.sum().clamp(min=1)
            else:
                loss = F.cross_entropy(
                    shift_logits.view(-1, shift_logits.size(-1)),
                    shift_targets.view(-1)
                )
            if not return_logits:
                logits = None
        return logits, loss

    @property
    def num_params(self):
        return sum(p.numel() for p in self.parameters()) - self.wte.weight.numel()


# ─── SFT Dataset ──────────────────────────────────────────────────────────────

class SFTDataset(torch.utils.data.Dataset):
    """
    Loads JSON SFT data (instruction/input/output) and tokenizes with prompt masking.
    Format per entry: {"instruction": "...", "input": "...", "output": "..."}

    Prompt template (Alpaca-style):
        ### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n{output}<|endoftext|>

    Loss mask: 0 for prompt tokens, 1 for response tokens (including EOS).
    """

    def __init__(self, json_path, tokenizer, max_len=1024):
        with open(json_path, "r", encoding="utf-8") as f:
            self.data = json.load(f)
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.eos_id = tokenizer.eos_token_id or 0

    def __len__(self):
        return len(self.data)

    def _format_prompt(self, entry):
        inst = entry.get("instruction", "").strip()
        inp = entry.get("input", "").strip()
        out = entry.get("output", "").strip()

        if inst and inp:
            prompt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n"
        elif inst:
            prompt = f"### Instruction:\n{inst}\n\n### Response:\n"
        else:
            # input-only format
            prompt = f"### Input:\n{inp}\n\n### Response:\n"

        return prompt, out

    def __getitem__(self, idx):
        entry = self.data[idx]
        prompt, response = self._format_prompt(entry)

        prompt_ids = self.tokenizer.encode(prompt)
        response_ids = self.tokenizer.encode(response) + [self.eos_id]

        total_ids = prompt_ids + response_ids

        # Truncate to max_len
        if len(total_ids) > self.max_len:
            total_ids = total_ids[:self.max_len]
            # Ensure EOS at end
            total_ids[-1] = self.eos_id
            # Recalculate prompt boundary
            prompt_len = min(len(prompt_ids), self.max_len)
        else:
            prompt_len = len(prompt_ids)

        # Build loss mask: 0 for prompt, 1 for response
        loss_mask = [0] * prompt_len + [1] * (len(total_ids) - prompt_len)

        # Pad to max_len
        pad_len = self.max_len - len(total_ids)
        total_ids = total_ids + [self.eos_id] * pad_len
        loss_mask = loss_mask + [0] * pad_len  # don't compute loss on padding

        input_ids = torch.tensor(total_ids, dtype=torch.long)
        loss_mask = torch.tensor(loss_mask, dtype=torch.long)

        return input_ids, loss_mask


# ─── Generation (for eval) ───────────────────────────────────────────────────

@torch.no_grad()
def generate(model, input_ids, max_new=150, temperature=0.7,
             top_p=0.9, top_k=40, device="cpu"):
    model.eval()
    ids = input_ids.clone().to(device)
    for _ in range(max_new):
        ctx = ids[:, -model.block_size:]
        logits, _ = model(ctx)
        logits = logits[:, -1, :] / max(temperature, 1e-8)
        if top_k > 0:
            vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < vals[:, -1:]] = -float("inf")
        probs = torch.softmax(logits, dim=-1)
        if top_p < 1.0:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cum = torch.cumsum(sorted_probs, dim=-1)
            mask = cum - sorted_probs > top_p
            sorted_probs[mask] = 0.0
            sorted_probs /= sorted_probs.sum()
            next_token = sorted_idx[0, torch.multinomial(sorted_probs[0], 1)]
        else:
            next_token = torch.multinomial(probs[0], 1)
        ids = torch.cat([ids, next_token.view(1, 1)], dim=1)
        if next_token.item() == 0:  # EOS
            break
    model.train()
    return ids[0, input_ids.size(1):]


def run_eval_prompts(model, tokenizer, prompts, device, step, out_dir):
    """Run eval prompts and print + log results."""
    model.eval()
    results = []
    sep = "─" * 60

    print(f"\n{sep}")
    print(f"  EVAL @ step {step}")
    print(sep)

    for prompt_text in prompts:
        # Format as instruction
        formatted = f"### Instruction:\n{prompt_text}\n\n### Response:\n"
        ids = tokenizer.encode(formatted, return_tensors="pt").to(device)
        out_ids = generate(model, ids, max_new=150, temperature=0.7, device=device)
        response = tokenizer.decode(out_ids.tolist(), skip_special_tokens=True).strip()

        print(f"  Q: {prompt_text}")
        print(f"  A: {response[:200]}")
        print()
        results.append({"prompt": prompt_text, "response": response[:500]})

    print(sep)

    # Save eval log
    eval_dir = Path(out_dir) / "evals"
    eval_dir.mkdir(parents=True, exist_ok=True)
    with open(eval_dir / f"eval_step_{step:06d}.json", "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    model.train()
    return results


# ─── Hardware Detection (same as train.py) ────────────────────────────────────

def probe_hardware():
    info = {
        "cpu_cores": os.cpu_count() or 4,
        "ram_gb": psutil.virtual_memory().total / 1024**3,
    }
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        info.update({
            "device": "cuda",
            "gpu_name": props.name,
            "vram_gb": props.total_memory / 1024**3,
            "sm_major": props.major,
        })
        if props.major >= 8:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            info["precision"] = "bf16"
            info["dtype"] = torch.bfloat16
        else:
            info["precision"] = "fp16"
            info["dtype"] = torch.float16
    else:
        info.update({
            "device": "cpu", "gpu_name": "CPU", "vram_gb": 0,
            "sm_major": 0, "precision": "fp32", "dtype": torch.float32,
        })
    return info


def probe_max_batch(model, device, dtype, seq_len, vocab_size, grad_accum_sim=4):
    """Binary search for max micro_batch. Safety: x0.70."""
    tmp_opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    lo, hi, best = 1, 512, 1
    while lo <= hi:
        mid = (lo + hi) // 2
        try:
            torch.cuda.empty_cache(); gc.collect()
            tmp_opt.zero_grad(set_to_none=True)
            for _ in range(grad_accum_sim):
                x = torch.randint(0, vocab_size, (mid, seq_len), device=device)
                mask = torch.ones_like(x)
                with autocast(device_type="cuda", dtype=dtype):
                    _, loss = model(x, x, loss_mask=mask, return_logits=False)
                    loss = loss / grad_accum_sim
                loss.backward()
                del x, mask, loss
            tmp_opt.step()
            tmp_opt.zero_grad(set_to_none=True)
            best = mid; lo = mid + 1
            torch.cuda.empty_cache()
        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
            if "out of memory" in str(e).lower() or isinstance(e, torch.cuda.OutOfMemoryError):
                try: del x, mask, loss
                except: pass
                torch.cuda.empty_cache()
                tmp_opt.zero_grad(set_to_none=True)
                hi = mid - 1
            else:
                raise
    del tmp_opt; torch.cuda.empty_cache(); gc.collect()
    safe = max(1, int(best * 0.70))
    print(f"  Probe: max_batch={best}, using {safe} (70% safety)")
    return safe


# ─── LR Schedule ──────────────────────────────────────────────────────────────

def cosine_lr(step, warmup, total, lr_max, lr_min):
    if step < warmup:
        return lr_max * (step + 1) / warmup
    p = (step - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (1 + math.cos(math.pi * p)) * (lr_max - lr_min)


# ─── Config ───────────────────────────────────────────────────────────────────

def load_sft_config(config_path):
    with open(config_path, encoding="utf-8") as f:
        raw = yaml.safe_load(f)

    cfg = {
        "auto_config":     raw.get("auto_config", True),
        "hf_model_repo":   raw.get("hf_model_repo", "ASTERIZER/LUNA-100M"),
        "hf_model_file":   raw.get("hf_model_file", "latest.pt"),
        "hf_dataset_repo": raw.get("hf_dataset_repo", "ASTERIZER/Luna_Dataset"),
        "pretrained_ckpt": raw.get("pretrained_ckpt", "Base/out/pretrain/luna_100m/latest.pt"),
        "train_json":      raw.get("train_json", "Base/Datasets/sft_clean/train.json"),
        "val_json":        raw.get("val_json", "Base/Datasets/sft_clean/val.json"),
        "out_dir":         raw.get("out_dir", "Base/out/sft/luna_100m_sft"),
        "tokenizer_dir":   raw.get("tokenizer_dir", "Base/checkpoints/EleutherAI/pythia-160m"),
        # model
        "vocab_size":  raw["model"]["vocab_size"],
        "seq_len":     raw["model"]["seq_len"],
        "n_layer":     raw["model"]["n_layer"],
        "n_embd":      raw["model"]["n_embd"],
        "n_head":      raw["model"]["n_head"],
        # train
        "epochs":          raw["train"]["epochs"],
        "max_tokens":      raw["train"].get("max_tokens", 0),
        "lr_warmup_steps": raw["train"]["lr_warmup_steps"],
        "save_interval":   raw["train"]["save_interval"],
        "log_interval":    raw["train"]["log_interval"],
        "eval_interval":   raw["train"]["eval_interval"],
        "max_norm":        raw["train"]["max_norm"],
        # optimizer
        "lr":           raw["optimizer"]["lr"],
        "min_lr":       raw["optimizer"]["min_lr"],
        "weight_decay": raw["optimizer"]["weight_decay"],
        "betas":        tuple(raw["optimizer"]["betas"]),
        "eps":          raw["optimizer"]["eps"],
        # batch
        "global_batch": raw["batch"]["global_batch"],
        "micro_batch":  raw["batch"]["micro_batch"],
        "grad_accum":   raw["batch"]["grad_accum"],
        # dataloader
        "num_workers": raw["dataloader"]["num_workers"],
        "pin_memory":  raw["dataloader"]["pin_memory"],
        # hardware
        "precision": raw["hardware"]["precision"],
        # eval prompts
        "eval_prompts": raw.get("eval_prompts", []),
    }
    return cfg


# ─── Training ─────────────────────────────────────────────────────────────────

SEP = "=" * 72

def sft_train(cfg):
    hw = probe_hardware()
    device = torch.device(hw["device"])

    if device.type == "cuda":
        torch.cuda.empty_cache(); gc.collect()

    # Precision
    if cfg["auto_config"]:
        dtype = hw.get("dtype", torch.float32)
        cfg["precision"] = hw["precision"]
    else:
        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16,
                 "fp32": torch.float32}.get(cfg["precision"], torch.float32)

    print(SEP)
    print("  LUNA 100M - SFT Fine-Tuning")
    print(SEP)
    print(f"  GPU          : {hw['gpu_name']}  ({hw['vram_gb']:.1f} GB)")
    print(f"  RAM          : {hw['ram_gb']:.1f} GB   CPU: {hw['cpu_cores']} cores")
    print(f"  Precision    : {cfg['precision']}   dtype={dtype}")
    print(f"  Pretrained   : {cfg['pretrained_ckpt']}")

    # ── Tokenizer ─────────────────────────────────────────────────────────────
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(cfg["tokenizer_dir"])
    print(f"  Tokenizer    : {cfg['tokenizer_dir']} (vocab={tokenizer.vocab_size})")

    # ── Model ─────────────────────────────────────────────────────────────────
    print(f"\n  Building LUNA-100M...")
    model = LUNAModel(
        vocab_size=cfg["vocab_size"],
        block_size=cfg["seq_len"],
        n_layer=cfg["n_layer"],
        n_embd=cfg["n_embd"],
        n_head=cfg["n_head"],
    ).to(device)
    print(f"  Parameters: {model.num_params:,} (unique)")

    # ── Load pretrained weights ───────────────────────────────────────────────
    ckpt_path = Path(cfg["pretrained_ckpt"])
    if not ckpt_path.exists() and cfg.get("hf_model_repo"):
        # Auto-download from HuggingFace model repo
        print(f"\n  Pretrained checkpoint not found locally.")
        print(f"  Downloading from HuggingFace: {cfg['hf_model_repo']} ({cfg['hf_model_file']})")
        from huggingface_hub import hf_hub_download
        ckpt_path.parent.mkdir(parents=True, exist_ok=True)
        downloaded = hf_hub_download(
            repo_id=cfg["hf_model_repo"],
            filename=cfg["hf_model_file"],
            local_dir=str(ckpt_path.parent),
            token=os.environ.get("HF_TOKEN"),
        )
        downloaded_path = Path(downloaded)
        if not ckpt_path.exists() and downloaded_path.exists():
            ckpt_path = downloaded_path
        print(f"  Downloaded to: {ckpt_path}")

    if ckpt_path.exists():
        print(f"\n  Loading pretrained checkpoint: {ckpt_path}")
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
        state = ckpt["model"] if "model" in ckpt else ckpt
        model.load_state_dict(state, strict=True)
        pretrain_step = ckpt.get("step", "?")
        pretrain_tokens = ckpt.get("tokens_seen", 0)
        print(f"  Pretrained @ step {pretrain_step}, tokens seen: {pretrain_tokens:,}")
        # Do NOT load optimizer state — we start fresh for SFT
    else:
        print(f"\n  WARNING: No pretrained checkpoint at {ckpt_path}")
        print(f"  Training from scratch (not recommended for SFT)!")

    # ── Dataset (auto-download from HF if missing) ─────────────────────────────
    train_path = Path(cfg["train_json"])
    val_path = Path(cfg["val_json"]) if cfg["val_json"] else None

    if not train_path.exists() and cfg.get("hf_dataset_repo"):
        print(f"\n  SFT dataset not found locally.")
        print(f"  Downloading from HuggingFace: {cfg['hf_dataset_repo']}")
        from huggingface_hub import hf_hub_download
        train_path.parent.mkdir(parents=True, exist_ok=True)
        hf_hub_download(
            repo_id=cfg["hf_dataset_repo"],
            repo_type="dataset",
            filename="train.json",
            local_dir=str(train_path.parent),
            token=os.environ.get("HF_TOKEN"),
        )
        print(f"  Downloaded train.json")
        if val_path:
            hf_hub_download(
                repo_id=cfg["hf_dataset_repo"],
                repo_type="dataset",
                filename="val.json",
                local_dir=str(val_path.parent),
                token=os.environ.get("HF_TOKEN"),
            )
            print(f"  Downloaded val.json")

    print(f"\n  Train data: {cfg['train_json']}")
    train_dataset = SFTDataset(cfg["train_json"], tokenizer, max_len=cfg["seq_len"])
    print(f"  Train entries: {len(train_dataset):,}")

    val_dataset = None
    if cfg["val_json"] and Path(cfg["val_json"]).exists():
        val_dataset = SFTDataset(cfg["val_json"], tokenizer, max_len=cfg["seq_len"])
        print(f"  Val entries: {len(val_dataset):,}")

    # ── Batch sizing ──────────────────────────────────────────────────────────
    if cfg["auto_config"] and device.type == "cuda":
        print(f"\n  Probing max micro_batch_size...")
        max_mbs = probe_max_batch(model, device, dtype, cfg["seq_len"], cfg["vocab_size"])
        model.load_state_dict(state, strict=True)  # reinit after probe
        torch.cuda.empty_cache(); gc.collect()
        grad_accum = max(1, math.ceil(cfg["global_batch"] / max_mbs))
        effective_batch = max_mbs * grad_accum
    else:
        max_mbs = cfg["micro_batch"]
        grad_accum = cfg["grad_accum"]
        effective_batch = max_mbs * grad_accum

    print(f"  micro_batch={max_mbs}, grad_accum={grad_accum}, effective={effective_batch}")

    # ── DataLoader ────────────────────────────────────────────────────────────
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=max_mbs,
        shuffle=True,
        num_workers=cfg["num_workers"],
        pin_memory=cfg["pin_memory"],
        drop_last=True,
        prefetch_factor=4 if cfg["num_workers"] > 0 else None,
        persistent_workers=cfg["num_workers"] > 0,
    )

    val_loader = None
    if val_dataset:
        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=max_mbs, shuffle=False,
            num_workers=min(2, cfg["num_workers"]),
            pin_memory=cfg["pin_memory"], drop_last=False,
        )

    # ── Optimizer ─────────────────────────────────────────────────────────────
    try:
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=cfg["lr"],
            weight_decay=cfg["weight_decay"],
            betas=cfg["betas"], eps=cfg["eps"], fused=True,
        )
    except TypeError:
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=cfg["lr"],
            weight_decay=cfg["weight_decay"],
            betas=cfg["betas"], eps=cfg["eps"],
        )

    use_scaler = dtype == torch.float16
    scaler = GradScaler(enabled=use_scaler)

    # ── Schedule ──────────────────────────────────────────────────────────────
    steps_per_epoch = len(train_loader) // grad_accum
    total_steps = steps_per_epoch * cfg["epochs"]
    warmup_steps = min(cfg["lr_warmup_steps"], total_steps // 5)

    out_dir = Path(cfg["out_dir"])
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n  Epochs         : {cfg['epochs']}")
    print(f"  Steps/epoch    : {steps_per_epoch:,}")
    print(f"  Total steps    : {total_steps:,}")
    print(f"  Warmup steps   : {warmup_steps}")
    print(f"  LR             : {cfg['lr']:.2e} -> {cfg['min_lr']:.2e}")
    print(f"  Save every     : {cfg['save_interval']} steps")
    print(f"  Eval every     : {cfg['eval_interval']} steps")
    print(f"  Eval prompts   : {len(cfg['eval_prompts'])}")
    print(f"  Out dir        : {out_dir}")
    print(SEP)

    # ── Resume SFT ────────────────────────────────────────────────────────────
    start_step = 0
    sft_ckpt_path = out_dir / "latest.pt"
    if sft_ckpt_path.exists():
        print(f"\n  Resuming SFT from {sft_ckpt_path}...")
        sft_ckpt = torch.load(sft_ckpt_path, map_location=device, weights_only=True)
        model.load_state_dict(sft_ckpt["model"])
        optimizer.load_state_dict(sft_ckpt["optimizer"])
        start_step = sft_ckpt["step"]
        print(f"  Resumed at SFT step {start_step}")

    # ── Eval at start ─────────────────────────────────────────────────────────
    if cfg["eval_prompts"] and start_step == 0:
        print("\n  Running initial eval (before SFT)...")
        run_eval_prompts(model, tokenizer, cfg["eval_prompts"], device, 0, out_dir)

    # ── Training loop ─────────────────────────────────────────────────────────
    model.train()
    run_t0 = time.perf_counter()
    step = start_step
    best_val_loss = float("inf")

    print(f"\n  Starting SFT training (step {start_step} -> {total_steps})...")

    for epoch in range(cfg["epochs"]):
        data_iter = iter(train_loader)
        micro_step = 0

        for batch_idx, (input_ids, loss_mask) in enumerate(data_iter):
            # Skip already-done steps on resume
            current_global_step = epoch * steps_per_epoch + (micro_step // grad_accum)
            if current_global_step < start_step and (micro_step % grad_accum == grad_accum - 1):
                micro_step += 1
                continue
            if current_global_step >= total_steps:
                break

            input_ids = input_ids.to(device, non_blocking=True)
            loss_mask = loss_mask.to(device, non_blocking=True)

            t0 = time.perf_counter()

            # Accumulation step
            with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
                _, loss = model(input_ids, targets=input_ids, loss_mask=loss_mask, return_logits=False)
                loss = loss / grad_accum

            scaler.scale(loss).backward()
            micro_step += 1

            # Optimizer step after grad_accum micro-batches
            if micro_step % grad_accum == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["max_norm"])

                # LR schedule
                lr_now = cosine_lr(step, warmup_steps, total_steps, cfg["lr"], cfg["min_lr"])
                for pg in optimizer.param_groups:
                    pg["lr"] = lr_now

                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)

                if device.type == "cuda":
                    torch.cuda.synchronize()

                dt = time.perf_counter() - t0
                step += 1

                # ── Log ───────────────────────────────────────────────────────
                if step % cfg["log_interval"] == 0 or step <= 3:
                    tokens_step = effective_batch * cfg["seq_len"]
                    tps = tokens_step / dt if dt > 0 else 0
                    vram = torch.cuda.max_memory_allocated() / 1024**3 if device.type == "cuda" else 0
                    eta_h = (total_steps - step) * dt / 3600
                    print(f"  step {step:6d}/{total_steps} | epoch {epoch+1}/{cfg['epochs']} | "
                          f"loss {loss.item()*grad_accum:.4f} | lr {lr_now:.2e} | "
                          f"{tps:,.0f} tok/s | VRAM {vram:.1f}GB | ETA {eta_h:.1f}h")

                # ── Save checkpoint ───────────────────────────────────────────
                if step % cfg["save_interval"] == 0 or step == total_steps:
                    raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
                    step_dir = out_dir / f"step-{step:06d}"
                    step_dir.mkdir(parents=True, exist_ok=True)
                    torch.save(raw_model.state_dict(), step_dir / "model.pth")
                    torch.save({
                        "step": step,
                        "model": raw_model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "epoch": epoch,
                        "sft_loss": loss.item() * grad_accum,
                    }, out_dir / "latest.pt")
                    print(f"  Saved -> {step_dir}")

                # ── Eval ──────────────────────────────────────────────────────
                if step % cfg["eval_interval"] == 0 or step == total_steps:
                    # Validation loss
                    if val_loader:
                        model.eval()
                        val_loss_sum = 0.0
                        val_count = 0
                        with torch.no_grad():
                            for val_ids, val_mask in val_loader:
                                val_ids = val_ids.to(device, non_blocking=True)
                                val_mask = val_mask.to(device, non_blocking=True)
                                with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
                                    _, vl = model(val_ids, targets=val_ids, loss_mask=val_mask, return_logits=False)
                                val_loss_sum += vl.item()
                                val_count += 1
                                if val_count >= 50:  # cap eval to 50 batches
                                    break
                        avg_val = val_loss_sum / max(val_count, 1)
                        print(f"  Val loss: {avg_val:.4f}")
                        if avg_val < best_val_loss:
                            best_val_loss = avg_val
                            raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
                            torch.save(raw_model.state_dict(), out_dir / "best_model.pth")
                            print(f"  New best! Saved best_model.pth")
                        model.train()

                    # Run eval prompts
                    if cfg["eval_prompts"]:
                        run_eval_prompts(model, tokenizer, cfg["eval_prompts"],
                                         device, step, out_dir)

    # ── Final ─────────────────────────────────────────────────────────────────
    final_dir = out_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
    torch.save(raw_model.state_dict(), final_dir / "model.pth")
    torch.save({
        "step": step,
        "model": raw_model.state_dict(),
        "sft_complete": True,
    }, out_dir / "latest.pt")

    # Copy tokenizer
    import shutil
    tok_src = Path(cfg["tokenizer_dir"])
    if tok_src.exists():
        shutil.copytree(tok_src, final_dir / "tokenizer", dirs_exist_ok=True)

    total_h = (time.perf_counter() - run_t0) / 3600
    print(SEP)
    print(f"  SFT Complete! {total_h:.2f}h -> {final_dir}")
    print(f"  Best val loss: {best_val_loss:.4f}")
    print(SEP)


# ─── Entry ────────────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="LUNA 100M — SFT Fine-Tuning")
    p.add_argument("--config",     default="sft_config.yaml")
    p.add_argument("--pretrained_ckpt", default=None)
    p.add_argument("--train_json", default=None)
    p.add_argument("--val_json",   default=None)
    p.add_argument("--out_dir",    default=None)
    p.add_argument("--epochs",     type=int, default=None)
    p.add_argument("--lr",         type=float, default=None)
    p.add_argument("--micro_batch",type=int, default=None)
    p.add_argument("--global_batch",type=int, default=None)
    p.add_argument("--save_interval", type=int, default=None)
    p.add_argument("--eval_interval", type=int, default=None)
    p.add_argument("--auto_config", type=lambda x: x.lower() in ("1","true","yes"),
                   default=None)
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()
    cfg = load_sft_config(args.config)
    # CLI overrides
    for key, val in vars(args).items():
        if key != "config" and val is not None:
            cfg[key] = val
    sft_train(cfg)