train.py

"""
ZeroShot-1B — single-file LLaMA-style 1.19B trainer + inference.

Architecture trimmed from the original spec to fit microbatch=4 on a 32 GB
5090. Original was FFN=8192, ctx=4096 → ~29 GB activations alone, OOM.
This: FFN=5632, ctx=3072 → ~21 GB activations, ~29 GB peak.

Subcommands:
  python train.py base     --no-compile
  python train.py mid      --checkpoint ckpt_base_final.pt --no-compile
  python train.py finetune --checkpoint ckpt_mid_final.pt  --no-compile
  python train.py finetune --checkpoint ckpt_mid_final.pt  --no-compile --skip_mid
  python train.py generate --checkpoint ckpt_sft_final.pt  --prompt "..." --chat

Auto-resume: re-launching a stage picks up the newest matching ckpt in
--ckpt_dir (model + optimizer + step + RNG + dataset sample counter). Pass
--fresh to ignore. --checkpoint is for cross-stage *weight* init only
(fresh optimizer).

Exit codes:
  0  — clean finish
  2  — recoverable failure (data, OOM, SIGTERM) — emergency ckpt saved, run.sh restarts
  130 — Ctrl-C
"""
import argparse
import math
import os
import queue
import random
import signal
import sys
import threading
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterator

import numpy as np
import tiktoken
import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset


# =============================================================================
# CONFIG
# =============================================================================

@dataclass
class ModelConfig:
    vocab_size: int = 50304       # GPT-2 (50257) padded for tensor cores
    n_layers: int = 24
    n_heads: int = 16
    n_kv_heads: int = 4           # GQA: 4 KV heads, 4 query heads per group
    d_model: int = 2048
    d_ff: int = 5632              # SwiGLU expansion (~2.75x d_model). 8192 OOMs at ctx>=3072 mb=4.
    max_seq_len: int = 3072
    rope_theta: float = 10000.0
    rms_eps: float = 1e-5


@dataclass
class TrainConfig:
    stage: str = "base"
    micro_batch_size: int = 4
    grad_accum: int = 20          # tokens/step = 4 * 20 * 3072 = 245,760 (~0.25M)
    seq_len: int = 3072
    max_steps: int = 25_000
    warmup_steps: int = 2_000
    peak_lr: float = 3e-4
    min_lr: float = 3e-5
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    log_every: int = 10
    ckpt_interval_sec: int = 30 * 60   # time-based; vast instances die at unpredictable steps
    ckpt_keep: int = 1
    ckpt_dir: str = "./checkpoints"
    seed: int = 42
    cost_per_hour: float = 0.365
    wandb_project: str = "zeroshot-1b"
    data_max_retries: int = 5
    data_retry_base_delay: float = 2.0


STAGE_PRESETS = {
    "base": dict(
        max_steps=25_000, warmup_steps=2_000, peak_lr=3e-4, min_lr=3e-5,
        ckpt_prefix="ckpt_base", micro_batch_size=4, grad_accum=20,
    ),
    "mid": dict(
        max_steps=2_500, warmup_steps=200, peak_lr=1e-4, min_lr=1e-5,
        ckpt_prefix="ckpt_mid", micro_batch_size=4, grad_accum=20,
    ),
    "finetune": dict(
        max_steps=2_000, warmup_steps=100, peak_lr=5e-5, min_lr=5e-6,
        ckpt_prefix="ckpt_sft", micro_batch_size=4, grad_accum=8,
    ),
}


# =============================================================================
# MODEL — RMSNorm, RoPE, GQA attention, SwiGLU, transformer
# =============================================================================

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        x32 = x.float()
        norm = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return (norm * self.weight.float()).type_as(x)


def precompute_rope(head_dim, max_seq_len, theta, device="cpu"):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    # x: (B, H, T, D). cos/sin: (T, D/2). GPT-NeoX-style split-halves rotation.
    x1, x2 = x.chunk(2, dim=-1)
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class Attention(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.n_rep = cfg.n_heads // cfg.n_kv_heads
        self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)

    def forward(self, x, cos, sin, kv_cache=None, pos=0):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        cos_t = cos[pos:pos + T]
        sin_t = sin[pos:pos + T]
        q = apply_rope(q, cos_t, sin_t)
        k = apply_rope(k, cos_t, sin_t)

        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k_cache[:, :, pos:pos + T] = k
            v_cache[:, :, pos:pos + T] = v
            k = k_cache[:, :, :pos + T]
            v = v_cache[:, :, :pos + T]

        # SDPA's enable_gqa is still flagged experimental and has open bug reports;
        # explicit repeat is correct on every nightly we'd be using.
        k = k.repeat_interleave(self.n_rep, dim=1)
        v = v.repeat_interleave(self.n_rep, dim=1)

        # is_causal=True uses top-left mask alignment in SDPA, which is correct for
        # training (Q==K len) and prefill at pos=0. Single-token decode wants no
        # mask. We never do chunked prefill (pos>0 with T>1), so this is sufficient.
        is_causal = (kv_cache is None) or (T > 1)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)


class SwiGLU(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.up_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.down_proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.attn = Attention(cfg)
        self.mlp_norm = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.mlp = SwiGLU(cfg)

    def forward(self, x, cos, sin, kv_cache=None, pos=0):
        x = x + self.attn(self.attn_norm(x), cos, sin, kv_cache, pos)
        x = x + self.mlp(self.mlp_norm(x))
        return x


class LLaMA(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.norm = RMSNorm(cfg.d_model, cfg.rms_eps)

        cos, sin = precompute_rope(cfg.d_model // cfg.n_heads, 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)
        for name, p in self.named_parameters():
            if name.endswith("o_proj.weight") or name.endswith("down_proj.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layers))

    @staticmethod
    def _init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            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=0.02)

    def forward(self, idx, targets=None, kv_caches=None, pos=0, ignore_index=-100):
        x = self.tok_emb(idx)
        for i, layer in enumerate(self.layers):
            kv = kv_caches[i] if kv_caches is not None else None
            x = layer(x, self.rope_cos, self.rope_sin, kv, pos)
        x = self.norm(x)
        logits = F.linear(x, self.tok_emb.weight)  # tied lm head
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)).float(),
                targets.view(-1),
                ignore_index=ignore_index,
            )
        return logits, loss

    def num_params(self):
        return sum(p.numel() for p in self.parameters())

    def configure_optimizer(self, lr, weight_decay, betas):
        # 8-bit AdamW: more stable than bf16 Adam states (variance term is fragile),
        # less memory than fp32 states. Embedding gets fp32 override per bnb docs.
        import bitsandbytes as bnb
        decay, no_decay = [], []
        for _, p in self.named_parameters():
            if not p.requires_grad:
                continue
            (decay if p.dim() >= 2 else no_decay).append(p)
        groups = [
            {"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0},
        ]
        opt = bnb.optim.AdamW8bit(groups, lr=lr, betas=betas)
        bnb.optim.GlobalOptimManager.get_instance().register_module_override(
            self.tok_emb, "weight", {"optim_bits": 32}
        )
        return opt


# =============================================================================
# DATA — streaming loaders with retry + background prefetch
# =============================================================================

ENC = tiktoken.get_encoding("gpt2")
EOT = ENC.eot_token  # 50256
IGNORE = -100

USER_TAG = "\n<|user|>\n"
ASSISTANT_TAG = "\n<|assistant|>\n"


class _StreamError(RuntimeError):
    pass


def _retry_iter(open_fn, max_retries, base_delay):
    """Run open_fn() to get a fresh iterator, yield from it; on exception
    recreate via open_fn() with exp backoff. open_fn closes over the right
    skip count (set on stage start; on hard restart, run.sh re-invokes us
    and we recreate from the latest ckpt's samples_consumed)."""
    attempt = 0
    while True:
        try:
            it = open_fn()
            for x in it:
                yield x
                attempt = 0
            return
        except StopIteration:
            return
        except Exception as e:
            attempt += 1
            if attempt > max_retries:
                raise _StreamError(f"max retries ({max_retries}) exceeded: {e}") from e
            delay = base_delay * (2 ** (attempt - 1))
            print(f"[data] {type(e).__name__}: {e}; retry {attempt}/{max_retries} in {delay:.1f}s",
                  flush=True)
            time.sleep(delay)


def _open_fineweb(seed, skip):
    ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT",
                      split="train", streaming=True)
    ds = ds.shuffle(seed=seed, buffer_size=10_000)
    return ds.skip(skip) if skip else ds


def _open_cosmopedia(seed, skip):
    ds = load_dataset("HuggingFaceTB/cosmopedia", name="web_samples_v2",
                      split="train", streaming=True)
    ds = ds.shuffle(seed=seed, buffer_size=10_000)
    return ds.skip(skip) if skip else ds


def _open_ultrachat(seed, skip):
    ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft", streaming=True)
    ds = ds.shuffle(seed=seed, buffer_size=2_000)
    return ds.skip(skip) if skip else ds


def _doc_tokens(ds_iter, text_key):
    for sample in ds_iter:
        text = sample.get(text_key) or ""
        if not text:
            continue
        ids = ENC.encode_ordinary(text)
        ids.append(EOT)
        for t in ids:
            yield t


def _pack(token_iter: Iterator[int], seq_len):
    buf = []
    for tok in token_iter:
        buf.append(tok)
        if len(buf) >= seq_len + 1:
            chunk = buf[: seq_len + 1]
            del buf[: seq_len]
            x = torch.tensor(chunk[:-1], dtype=torch.long)
            y = torch.tensor(chunk[1:], dtype=torch.long)
            yield x, y


def base_loader(seq_len, seed, skip_samples, max_retries, base_delay):
    yield from _pack(
        _doc_tokens(_retry_iter(lambda: _open_fineweb(seed, skip_samples),
                                max_retries, base_delay), "text"),
        seq_len,
    )


def mid_loader(seq_len, seed, skip_samples, max_retries, base_delay, mix_ratio=0.1):
    """90% FineWeb-Edu / 10% Cosmopedia, switched at document boundaries."""
    fw_iter = _doc_tokens(
        _retry_iter(lambda: _open_fineweb(seed, skip_samples), max_retries, base_delay), "text")
    co_iter = _doc_tokens(
        _retry_iter(lambda: _open_cosmopedia(seed + 1, skip_samples // 10), max_retries, base_delay), "text")
    rng = random.Random(seed)

    def mixed():
        while True:
            src = co_iter if rng.random() < mix_ratio else fw_iter
            while True:
                t = next(src)
                yield t
                if t == EOT:
                    break

    yield from _pack(mixed(), seq_len)


def _format_chat(messages):
    ids, mask = [], []
    for msg in messages:
        role = msg.get("role")
        content = (msg.get("content") or "").strip()
        if not content:
            continue
        if role == "user":
            tag = ENC.encode_ordinary(USER_TAG + content)
            ids.extend(tag); mask.extend([0] * len(tag))
        elif role == "assistant":
            head = ENC.encode_ordinary(ASSISTANT_TAG)
            body = ENC.encode_ordinary(content)
            ids.extend(head); mask.extend([0] * len(head))
            ids.extend(body); mask.extend([1] * len(body))
            ids.append(EOT); mask.append(1)
    return ids, mask


def sft_loader(seq_len, seed, skip_samples, max_retries, base_delay):
    for sample in _retry_iter(lambda: _open_ultrachat(seed, skip_samples), max_retries, base_delay):
        ids, mask = _format_chat(sample.get("messages") or [])
        if len(ids) < 8:
            continue
        ids = ids[: seq_len + 1]
        mask = mask[: seq_len + 1]
        pad_n = (seq_len + 1) - len(ids)
        if pad_n > 0:
            ids += [EOT] * pad_n
            mask += [0] * pad_n
        x = torch.tensor(ids[:-1], dtype=torch.long)
        y = torch.tensor(
            [t if m else IGNORE for t, m in zip(ids[1:], mask[1:])],
            dtype=torch.long,
        )
        yield x, y


class _Prefetcher:
    """Background thread fills a bounded queue. Errors propagate on next()."""
    def __init__(self, gen, max_prefetch=8):
        self.q = queue.Queue(maxsize=max_prefetch)
        self._sentinel = object()
        self._error = None
        self.thread = threading.Thread(target=self._run, args=(gen,), daemon=True)
        self.thread.start()

    def _run(self, gen):
        try:
            for x in gen:
                self.q.put(x)
        except BaseException as e:
            self._error = e
        finally:
            self.q.put(self._sentinel)

    def __iter__(self):
        return self

    def __next__(self):
        x = self.q.get()
        if x is self._sentinel:
            if self._error is not None:
                raise self._error
            raise StopIteration
        return x


class StagedLoader:
    """Wraps a sample generator into a batch iterator and counts samples consumed."""
    def __init__(self, gen_fn, train_cfg, stage_seed, skip_samples=0, prefetch=8):
        self.samples_consumed = skip_samples
        self.micro_batch = train_cfg.micro_batch_size
        gen = gen_fn(
            seq_len=train_cfg.seq_len,
            seed=stage_seed,
            skip_samples=skip_samples,
            max_retries=train_cfg.data_max_retries,
            base_delay=train_cfg.data_retry_base_delay,
        )
        self._gen = _Prefetcher(gen, max_prefetch=prefetch)

    def __iter__(self):
        return self

    def __next__(self):
        xs, ys = [], []
        for _ in range(self.micro_batch):
            x, y = next(self._gen)
            xs.append(x); ys.append(y)
            self.samples_consumed += 1
        return torch.stack(xs), torch.stack(ys)


def build_loader(stage, train_cfg, stage_seed, skip_samples=0):
    gen = {"base": base_loader, "mid": mid_loader, "finetune": sft_loader}[stage]
    return StagedLoader(gen, train_cfg, stage_seed, skip_samples)


# =============================================================================
# CHECKPOINT + TRAIN HELPERS
# =============================================================================

def cosine_lr(step, warmup, max_steps, peak, min_lr):
    if step < warmup:
        return peak * (step + 1) / max(1, warmup)
    if step >= max_steps:
        return min_lr
    progress = (step - warmup) / max(1, max_steps - warmup)
    return min_lr + 0.5 * (peak - min_lr) * (1.0 + math.cos(math.pi * progress))


def find_latest_ckpt(ckpt_dir: Path, prefix: str):
    files = list(ckpt_dir.glob(f"{prefix}_step*.pt"))
    return max(files, key=lambda p: p.stat().st_mtime) if files else None


def prune_ckpts(ckpt_dir: Path, prefix: str, keep: int):
    files = sorted(ckpt_dir.glob(f"{prefix}_step*.pt"), key=lambda p: p.stat().st_mtime)
    for f in files[:-keep]:
        try:
            f.unlink()
        except OSError:
            pass


def get_rng_state():
    return {
        "torch": torch.get_rng_state(),
        "torch_cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
        "numpy": np.random.get_state(),
        "python": random.getstate(),
    }


def set_rng_state(state):
    torch.set_rng_state(state["torch"].cpu().byte())
    if state.get("torch_cuda") is not None and torch.cuda.is_available():
        torch.cuda.set_rng_state_all([s.cpu().byte() for s in state["torch_cuda"]])
    np.random.set_state(state["numpy"])
    random.setstate(state["python"])


def save_checkpoint(path, model, optimizer, step, train_cfg, model_cfg, samples_consumed):
    payload = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "step": step,
        "train_cfg": asdict(train_cfg),
        "model_cfg": asdict(model_cfg),
        "samples_consumed": samples_consumed,
        "rng": get_rng_state(),
    }
    tmp = str(path) + ".tmp"
    torch.save(payload, tmp)
    os.replace(tmp, path)


def load_checkpoint(path, model, optimizer=None, map_location="cuda"):
    ckpt = torch.load(path, map_location=map_location, weights_only=False)
    sd = ckpt["model"]
    if any(k.startswith("_orig_mod.") for k in sd): sd = {k.removeprefix("_orig_mod."): v for k, v in sd.items()}
    model.load_state_dict(sd)
    if optimizer is not None and "optimizer" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer"])
    return ckpt


def estimate_vram_gb(model_cfg: ModelConfig, train_cfg: TrainConfig, n_params: int):
    weight_gb = n_params * 2 / 1e9
    grad_gb = n_params * 2 / 1e9
    optim_gb = n_params * 2 / 1e9 + 0.82  # 8-bit AdamW + fp32 emb override
    fixed = weight_gb + grad_gb + optim_gb
    B, T = train_cfg.micro_batch_size, train_cfg.seq_len
    per_layer = B * T * (13 * model_cfg.d_model + 8 * model_cfg.d_ff) * 2 / 1e9
    return fixed, per_layer * model_cfg.n_layers


def print_startup(model, train_cfg, model_cfg, ckpt_prefix, start_step):
    n = model.num_params()
    fixed, acts = estimate_vram_gb(model_cfg, train_cfg, n)
    tps = train_cfg.micro_batch_size * train_cfg.grad_accum * train_cfg.seq_len
    bar = "=" * 64
    print(bar)
    print(f"  ZeroShot-1B — stage: {train_cfg.stage}  (prefix: {ckpt_prefix})")
    print(bar)
    print(f"  params:       {n/1e9:.3f}B  ({n:,})")
    print(f"  dtype:        bfloat16   optim: AdamW8bit (bnb)   ctx: {train_cfg.seq_len}")
    print(f"  arch:         L={model_cfg.n_layers}  H={model_cfg.n_heads}  KV={model_cfg.n_kv_heads}  "
          f"d={model_cfg.d_model}  ffn={model_cfg.d_ff}")
    print(f"  micro_batch:  {train_cfg.micro_batch_size}   grad_accum: {train_cfg.grad_accum}   "
          f"tokens/step: {tps:,}")
    print(f"  schedule:     cosine warmup={train_cfg.warmup_steps}  peak={train_cfg.peak_lr:.1e}  "
          f"min={train_cfg.min_lr:.1e}  steps={train_cfg.max_steps}")
    print(f"  total tokens: {tps*train_cfg.max_steps/1e9:.2f}B   (start step {start_step})")
    print(f"  vram (est):   fixed {fixed:.1f} + acts {acts:.1f} = {fixed+acts:.1f} GB  / 32 GB")
    print(f"  ckpt every:   {train_cfg.ckpt_interval_sec//60} min wall  →  "
          f"{train_cfg.ckpt_dir}/  (keep last {train_cfg.ckpt_keep})")
    print(bar, flush=True)


def fmt_eta(seconds):
    if seconds <= 0 or math.isinf(seconds):
        return "?"
    h = int(seconds // 3600); m = int((seconds % 3600) // 60)
    return f"{h:d}h{m:02d}m"


# =============================================================================
# TRAIN LOOP
# =============================================================================

def train(stage: str, args):
    preset = STAGE_PRESETS[stage]
    train_cfg = TrainConfig(stage=stage, **{k: v for k, v in preset.items() if k != "ckpt_prefix"})
    ckpt_prefix = preset["ckpt_prefix"]

    if args.lr is not None: train_cfg.peak_lr = args.lr
    if args.min_lr is not None: train_cfg.min_lr = args.min_lr
    if args.batch_size is not None: train_cfg.micro_batch_size = args.batch_size
    if args.grad_accum is not None: train_cfg.grad_accum = args.grad_accum
    if args.steps is not None: train_cfg.max_steps = args.steps
    if args.warmup is not None: train_cfg.warmup_steps = args.warmup
    if args.ckpt_dir is not None: train_cfg.ckpt_dir = args.ckpt_dir
    if args.seed is not None: train_cfg.seed = args.seed

    ckpt_dir = Path(train_cfg.ckpt_dir)
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    device = "cuda"
    torch.manual_seed(train_cfg.seed)
    np.random.seed(train_cfg.seed)
    random.seed(train_cfg.seed)
    torch.set_float32_matmul_precision("high")

    model_cfg = ModelConfig()
    model = LLaMA(model_cfg).to(device=device, dtype=torch.bfloat16)
    optimizer = model.configure_optimizer(
        lr=train_cfg.peak_lr,
        weight_decay=train_cfg.weight_decay,
        betas=(train_cfg.beta1, train_cfg.beta2),
    )

    start_step = 0
    samples_consumed = 0
    resume_path = None if args.fresh else find_latest_ckpt(ckpt_dir, ckpt_prefix)

    if resume_path is not None:
        print(f"[resume] {resume_path}")
        ckpt = load_checkpoint(resume_path, model, optimizer, map_location=device)
        start_step = ckpt["step"]
        samples_consumed = ckpt.get("samples_consumed", 0)
        if "rng" in ckpt:
            set_rng_state(ckpt["rng"])
    elif args.checkpoint is not None:
        init_path = Path(args.checkpoint)
        if not init_path.is_absolute() and not init_path.exists():
            init_path = ckpt_dir / args.checkpoint
        print(f"[init-from] {init_path}  (weights only, fresh optimizer)")
        load_checkpoint(init_path, model, optimizer=None, map_location=device)
    elif stage != "base":
        raise SystemExit(f"{stage} requires --checkpoint pointing at the prior stage's ckpt "
                         f"(no auto-resume ckpt found in {ckpt_dir})")

    stage_seed = train_cfg.seed + {"base": 0, "mid": 1000, "finetune": 2000}[stage]
    model = torch.compile(model, mode="default")
    print_startup(model, train_cfg, model_cfg, ckpt_prefix, start_step)

    use_wandb = bool(os.environ.get("WANDB_API_KEY")) and not args.no_wandb
    if use_wandb:
        import wandb
        wandb.init(
            project=train_cfg.wandb_project,
            name=f"{stage}-{int(time.time())}",
            config={**asdict(train_cfg), **asdict(model_cfg)},
            resume="allow",
        )

    loader = build_loader(stage, train_cfg, stage_seed, samples_consumed)
    data_iter = iter(loader)
    tokens_per_step = train_cfg.micro_batch_size * train_cfg.grad_accum * train_cfg.seq_len

    # ---- emergency-ckpt helpers ----
    _state = {"step": start_step}

    def emergency_save(reason):
        path = ckpt_dir / f"{ckpt_prefix}_emergency.pt"
        save_checkpoint(path, model, optimizer, _state["step"], train_cfg, model_cfg,
                        loader.samples_consumed)
        print(f"[emergency] saved {path}  reason: {reason}", flush=True)

    def on_sigterm(signum, frame):
        print("[signal] SIGTERM received", flush=True)
        emergency_save("SIGTERM")
        sys.exit(2)

    signal.signal(signal.SIGTERM, on_sigterm)

    # ---- the loop ----
    model.train()
    run_start = time.time()
    last_ckpt_time = time.time()
    t_window = time.time()
    tokens_window = 0
    loss_window = 0.0
    loss_count = 0

    try:
        for step in range(start_step, train_cfg.max_steps):
            _state["step"] = step
            lr = cosine_lr(step, train_cfg.warmup_steps, train_cfg.max_steps,
                           train_cfg.peak_lr, train_cfg.min_lr)
            for g in optimizer.param_groups:
                g["lr"] = lr

            optimizer.zero_grad(set_to_none=True)
            accum_loss = 0.0
            for _ in range(train_cfg.grad_accum):
                x, y = next(data_iter)
                x = x.to(device, non_blocking=True)
                y = y.to(device, non_blocking=True)
                _, loss = model(x, targets=y)
                (loss / train_cfg.grad_accum).backward()
                accum_loss += loss.item()
            accum_loss /= train_cfg.grad_accum

            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip)
            optimizer.step()

            loss_window += accum_loss
            loss_count += 1
            tokens_window += tokens_per_step

            if (step + 1) % train_cfg.log_every == 0:
                dt = time.time() - t_window
                tps = tokens_window / max(dt, 1e-6)
                avg_loss = loss_window / loss_count
                elapsed = time.time() - run_start
                eta = (train_cfg.max_steps - (step + 1)) * (dt / train_cfg.log_every)
                cost = (elapsed / 3600.0) * train_cfg.cost_per_hour
                print(
                    f"step {step+1:>6d}/{train_cfg.max_steps} | "
                    f"loss {avg_loss:.4f} | lr {lr:.2e} | grad {grad_norm.item():.2f} | "
                    f"{tps/1e3:.1f}k tok/s | "
                    f"elapsed {fmt_eta(elapsed)} | eta {fmt_eta(eta)} | ${cost:.2f}"
                )
                if use_wandb:
                    import wandb
                    wandb.log({
                        "loss": avg_loss, "lr": lr, "grad_norm": grad_norm.item(),
                        "tokens_per_sec": tps, "step": step + 1,
                        "elapsed_sec": elapsed, "cost_usd": cost,
                    })
                t_window = time.time()
                tokens_window = 0
                loss_window = 0.0
                loss_count = 0

            if time.time() - last_ckpt_time >= train_cfg.ckpt_interval_sec:
                path = ckpt_dir / f"{ckpt_prefix}_step{step+1:07d}.pt"
                save_checkpoint(path, model, optimizer, step + 1, train_cfg, model_cfg,
                                loader.samples_consumed)
                prune_ckpts(ckpt_dir, ckpt_prefix, train_cfg.ckpt_keep)
                last_ckpt_time = time.time()
                print(f"[ckpt] {path.name}  saved at step {step+1}, "
                      f"can resume from step {step+1}", flush=True)

    except _StreamError as e:
        print(f"[fatal] dataloader: {e}", flush=True)
        emergency_save(f"_StreamError: {e}")
        sys.exit(2)
    except torch.cuda.OutOfMemoryError as e:
        print(f"[fatal] OOM: {e}", flush=True)
        emergency_save("OOM")
        sys.exit(2)
    except KeyboardInterrupt:
        print("[interrupt] user", flush=True)
        emergency_save("KeyboardInterrupt")
        sys.exit(130)

    final = ckpt_dir / f"{ckpt_prefix}_final.pt"
    save_checkpoint(final, model, optimizer, train_cfg.max_steps, train_cfg, model_cfg,
                    loader.samples_consumed)
    elapsed = time.time() - run_start
    cost = (elapsed / 3600.0) * train_cfg.cost_per_hour
    print(f"[done] {final}  total elapsed {fmt_eta(elapsed)}  approx cost ${cost:.2f}")


# =============================================================================
# INFERENCE — pre-allocated KV cache
# =============================================================================

def alloc_kv_cache(cfg: ModelConfig, batch_size, max_len, device, dtype):
    head_dim = cfg.d_model // cfg.n_heads
    return [
        (
            torch.zeros(batch_size, cfg.n_kv_heads, max_len, head_dim, device=device, dtype=dtype),
            torch.zeros(batch_size, cfg.n_kv_heads, max_len, head_dim, device=device, dtype=dtype),
        )
        for _ in range(cfg.n_layers)
    ]


@torch.no_grad()
def sample(model, prompt_ids, max_new=200, temperature=0.8, top_k=50, device="cuda"):
    model.eval()
    cfg = model.cfg
    T0 = len(prompt_ids)
    max_len = min(T0 + max_new, cfg.max_seq_len)
    kv = alloc_kv_cache(cfg, 1, max_len, device, torch.bfloat16)

    idx = torch.tensor([prompt_ids], dtype=torch.long, device=device)
    logits, _ = model(idx, kv_caches=kv, pos=0)
    next_logits = logits[:, -1, :]

    out = list(prompt_ids)
    for i in range(max_new):
        if T0 + i >= cfg.max_seq_len:
            break
        if temperature <= 0:
            tok = next_logits.argmax(-1, keepdim=True)
        else:
            scaled = next_logits / temperature
            if top_k:
                v, _ = torch.topk(scaled, min(top_k, scaled.size(-1)))
                scaled = scaled.masked_fill(scaled < v[:, [-1]], -float("inf"))
            probs = torch.softmax(scaled.float(), dim=-1)
            tok = torch.multinomial(probs, 1)
        tid = tok.item()
        out.append(tid)
        if tid == EOT:
            break
        logits, _ = model(tok, kv_caches=kv, pos=T0 + i)
        next_logits = logits[:, -1, :]
    return out


def generate_cmd(args):
    device = "cuda"
    ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False)
    cfg = ModelConfig(**ckpt["model_cfg"]) if "model_cfg" in ckpt else ModelConfig()
    model = LLaMA(cfg).to(device=device, dtype=torch.bfloat16)
    sd = ckpt["model"]
    if any(k.startswith("_orig_mod.") for k in sd): sd = {k.removeprefix("_orig_mod."): v for k, v in sd.items()}
    model.load_state_dict(sd)

    prompt = f"\n<|user|>\n{args.prompt}\n<|assistant|>\n" if args.chat else args.prompt
    ids = ENC.encode_ordinary(prompt)

    t0 = time.time()
    out = sample(model, ids, args.max_new, args.temperature, args.top_k, device)
    dt = time.time() - t0
    new_tokens = len(out) - len(ids)
    print(ENC.decode(out))
    print(f"\n[gen] {new_tokens} new tokens in {dt:.2f}s = {new_tokens/dt:.1f} tok/s")


# =============================================================================
# CLI
# =============================================================================

def build_parser():
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd", required=True)

    def add_train_args(sp):
        sp.add_argument("--no-compile", action="store_true",
                        help="Required on Blackwell; compile is never invoked here regardless.")
        sp.add_argument("--checkpoint", type=str, default=None,
                        help="Initialize weights from this file (for cross-stage seeding).")
        sp.add_argument("--fresh", action="store_true",
                        help="Ignore any existing ckpt in --ckpt_dir for this stage; start over.")
        sp.add_argument("--ckpt_dir", type=str, default=None)
        sp.add_argument("--lr", type=float, default=None)
        sp.add_argument("--min_lr", type=float, default=None)
        sp.add_argument("--batch_size", type=int, default=None)
        sp.add_argument("--grad_accum", type=int, default=None)
        sp.add_argument("--steps", type=int, default=None)
        sp.add_argument("--warmup", type=int, default=None)
        sp.add_argument("--seed", type=int, default=None)
        sp.add_argument("--no-wandb", action="store_true")

    add_train_args(sub.add_parser("base"))
    add_train_args(sub.add_parser("mid"))
    ft = sub.add_parser("finetune"); add_train_args(ft)
    ft.add_argument("--skip_mid", action="store_true",
                    help="UX flag: confirms starting SFT directly from a mid checkpoint.")

    g = sub.add_parser("generate")
    g.add_argument("--checkpoint", required=True)
    g.add_argument("--prompt", default="Once upon a time")
    g.add_argument("--max_new", type=int, default=200)
    g.add_argument("--temperature", type=float, default=0.8)
    g.add_argument("--top_k", type=int, default=50)
    g.add_argument("--chat", action="store_true",
                   help="Wrap prompt in <|user|>/<|assistant|> tags (use with SFT checkpoints).")
    return p


def main():
    args = build_parser().parse_args()
    if args.cmd == "generate":
        generate_cmd(args)
    else:
        train(args.cmd, args)


if __name__ == "__main__":
    main()