5chp.py

#!/usr/bin/env python3
# 5apg.py — AR-only trainer/decoder (Qwen3 tokenizer)
# Fresh-start safe, AMP dtype auto, OOM backoff, progressive block growth.
# Sampling: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
# Checkpoints: time-based and step-based (monotonic). Resume respects interval.
# FP8: --fp8-only [--fp8-fallback] attempts float8_e4m3fn autocast, otherwise bf16/FP16.
# Chinchilla-style target token calc uses ALL enabled params (core + AR head).
# Robust streaming: retries, dataset fallbacks, dataset:config, and local JSONL support.
# Chat SFT: --chat uses tokenizer.apply_chat_template on records with {role, content} lists;
#           freezing options for core; LR overrides; new run via --warmstart_from (no --resume).

from __future__ import annotations
import argparse, json, math, pathlib, random, time, os, sys
from contextlib import nullcontext
from typing import Dict, Any, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset, DownloadConfig
from transformers import AutoTokenizer, logging as hf_log
from tqdm.auto import tqdm

# ───────────────────────── Globals ─────────────────────────
hf_log.set_verbosity_error()
DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# Tokenizer
TOKENIZER_ID = os.environ.get("TOKENIZER_ID", "Qwen/Qwen3-235B-A22B-Thinking-2507")
tok = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True)
if tok.pad_token is None:
    tok.add_special_tokens({"pad_token": "[PAD]"})
VOCAB = max(tok.get_vocab().values()) + 1
BLANK = tok.pad_token_id
EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id

PRESETS: Dict[str, Dict[str, int]] = {
    "small":   dict(d=512, layers=8,  heads=16, rank=64),
    "smallx2": dict(d=512, layers=16, heads=16, rank=64),
    "base":    dict(d=768, layers=12, heads=24, rank=96),
}

DEFAULT_BLOCK = 576
LR_CORE, LR_HEAD = 5e-5, 2e-4
DEFAULT_SAVE_SEC = 24 * 3600
CKDIR = pathlib.Path("ckpts_joint")

# ───────────────────────── Utilities ─────────────────────────
def rng_state():
    if DEV.type == "cuda":
        try:
            return torch.cuda.get_rng_state(DEV)
        except TypeError:
            return torch.cuda.get_rng_state()
    return torch.get_rng_state()

def _is_probably_ckpt(path: pathlib.Path) -> bool:
    try:
        return path.is_file() and path.suffix == ".pt" and not path.name.endswith(".pt.tmp") and path.stat().st_size > (1<<20)
    except Exception:
        return False

def _resolve_ckpt(path: pathlib.Path) -> pathlib.Path | None:
    try:
        if path.is_dir():
            cands = sorted([p for p in path.glob("*.pt") if _is_probably_ckpt(p)],
                           key=lambda p: p.stat().st_mtime, reverse=True)
            return cands[0] if cands else None
        if path.suffix == ".tmp":
            solid = path.with_suffix("")
            return solid if _is_probably_ckpt(solid) else _resolve_ckpt(path.parent)
        return path if _is_probably_ckpt(path) else _resolve_ckpt(path.parent)
    except Exception:
        return None

def _try_load(path: pathlib.Path, map_location="cpu"):
    try:
        return torch.load(path, map_location="cpu")
    except Exception as e:
        print(f"[ckpt-skip] {path} not usable: {e}")
        return None

# ───────────────────────── AMP helper ─────────────────────────
try:
    from torch.amp import autocast as _ac, GradScaler
except ImportError:
    from torch.cuda.amp import autocast as _ac, GradScaler

def _supports_fp8() -> bool:
    return hasattr(torch, "float8_e4m3fn")

def _auto_amp_dtype(prefer_fp8: bool = False):
    if DEV.type != "cuda":
        return torch.float32
    if prefer_fp8 and _supports_fp8():
        return torch.float8_e4m3fn
    try:
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    except Exception:
        return torch.float16

def amp(enabled: bool, prefer_fp8: bool = False):
    if not (enabled and DEV.type == "cuda"):
        return nullcontext()
    return _ac(device_type="cuda", dtype=_auto_amp_dtype(prefer_fp8=prefer_fp8))

# ───────────────────────── Chat helpers ─────────────────────────
def _coerce_role(r: str) -> str:
    r = (r or "").lower()
    if r in {"user", "human", "customer", "questioner"}:
        return "user"
    if r in {"assistant", "gpt", "bot", "agent", "answerer"}:
        return "assistant"
    if r in {"system", "context", "instruction"}:
        return "system"
    return r or "user"

def _render_chat_text_from_ex(ex: dict, messages_key: str, add_generation_prompt: bool) -> Optional[str]:
    msgs = ex.get(messages_key)
    # common alternates
    if msgs is None:
        for alt in ("conversations", "dialog", "turns"):
            if isinstance(ex.get(alt), list):
                msgs = ex[alt]
                break
    if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict):
        try:
            norm = []
            for m in msgs:
                role = _coerce_role(m.get("role", ""))
                content = m.get("content", m.get("text", ""))
                if not isinstance(content, str):
                    continue
                norm.append({"role": role, "content": content})
            if not norm:
                return None
            return tok.apply_chat_template(norm, tokenize=False, add_generation_prompt=add_generation_prompt)
        except Exception:
            return None
    # prompt/response or instruction/output style
    for a, b in (("prompt", "response"), ("instruction", "output"), ("question", "answer")):
        if isinstance(ex.get(a), str) and isinstance(ex.get(b), str):
            return f"User: {ex[a]}\nAssistant: {ex[b]}"
    return None

# ───────────────────────── Robust streaming data ─────────────────────────
def _open_stream_one(ds_name: str, seed: int):
    """

    Supports:

      - 'dataset' or 'dataset:config' (e.g., 'allenai/c4:en')

      - 'json:/path/file.jsonl' (local JSONL)

    """
    if ":" in ds_name:
        base, config = ds_name.split(":", 1)
    else:
        base, config = ds_name, None

    dc = DownloadConfig(max_retries=5, use_etag=True, resume_download=True)
    if base == "json":
        if not config:
            raise ValueError("Use 'json:/path/to/file.jsonl' or glob like 'json:/data/*.jsonl'")
        data_files = {"train": config}
        ds = load_dataset("json", data_files=data_files, split="train", streaming=True, download_config=dc)
    else:
        if config:
            ds = load_dataset(base, config, split="train", streaming=True, download_config=dc)
        else:
            ds = load_dataset(base, split="train", streaming=True, download_config=dc)
    ds = ds.shuffle(buffer_size=10_000, seed=seed)
    return iter(ds)

def token_stream(args, target: int, seed: int = 42, max_retries: int = 999):
    """

    Comma-separated dataset fallbacks, resilient to HF 5xx, with chat/text handling.

    Example: --source "json:/data/oasst.jsonl,allenai/c4:en"

    """
    ds_names = args.source
    sources = [s.strip() for s in ds_names.split(",") if s.strip()]
    if not sources:
        sources = ["cerebras/SlimPajama-627B"]

    src_idx = 0
    emitted = 0
    it = None
    attempts = 0
    backoff_base = 2.0

    while emitted < target:
        try:
            if it is None:
                it = _open_stream_one(sources[src_idx], seed)
            ex = next(it)
            text = None
            if isinstance(ex, dict):
                if args.chat:
                    text = _render_chat_text_from_ex(ex, args.chat_messages_key, args.sft_add_generation_prompt)
                if text is None:
                    if args.dataset_field_text and isinstance(ex.get(args.dataset_field_text), str):
                        text = ex[args.dataset_field_text]
                    elif isinstance(ex.get("text"), str):
                        text = ex["text"]
            if not isinstance(text, str):
                attempts = 0
                continue
            enc = tok.encode(text)
            if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
                enc.append(EOS)
            for t in enc:
                yield t
                emitted += 1
                if emitted >= target:
                    return
            attempts = 0  # progress resets backoff
        except StopIteration:
            it = None
            src_idx = (src_idx + 1) % len(sources)
        except Exception as e:
            attempts += 1
            sleep_s = min(60.0, backoff_base ** min(attempts, 6))
            print(f"[stream-retry] source={sources[src_idx]} attempts={attempts} sleep={sleep_s:.1f}s reason={type(e).__name__}", flush=True)
            time.sleep(sleep_s)
            it = None
            if attempts % 5 == 0 and len(sources) > 1:
                src_idx = (src_idx + 1) % len(sources)
            if attempts > max_retries:
                raise

# ───────────────────────── Relative positional bias (ALiBi) ─────────────────────────
def _alibi_slopes(n_heads: int):
    import math
    def pow2slopes(n):
        start = 2 ** (-2 ** -(math.log2(n) - 3))
        ratio = start
        return [start * (ratio ** i) for i in range(n)]
    if math.log2(n_heads).is_integer():
        vals = pow2slopes(n_heads)
    else:
        closest = 2 ** math.floor(math.log2(n_heads))
        vals = pow2slopes(closest)
        extra = pow2slopes(2 * closest)
        vals += extra[0::2][: n_heads - closest]
    return torch.tensor(vals, device=DEV).view(1, n_heads, 1, 1)

def alibi_bias(n_heads: int, n_tokens: int):
    i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1)
    j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens)
    dist = (j - i).clamp_min(0)
    slopes = _alibi_slopes(n_heads)
    return -slopes * dist

# ───────────────────────── Model components ─────────────────────────
class LowRankMHA(nn.Module):
    def __init__(self, d: int, h: int, r: int, use_relpos: bool = True):
        super().__init__()
        assert d % h == 0, "d must be divisible by number of heads"
        self.h, self.dk = h, d // h
        self.use_relpos = use_relpos
        self.q = nn.Linear(d, d, bias=False)
        self.k = nn.Linear(d, d, bias=False)
        self.v = nn.Linear(d, d, bias=False)
        self.U = nn.Parameter(torch.randn(self.dk, r))
        nn.init.orthogonal_(self.U)
        self.proj = nn.Linear(h * r, d, bias=False)
        self.drop = nn.Dropout(0.1)

    def _proj(self, x):
        B, N, _ = x.shape
        return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,

                rel_bias_tokens: Optional[int] = None,

                kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,

                use_cache: bool = False):
        q = self._proj(self.q(x))
        k_new = self._proj(self.k(x))
        v_new = self._proj(self.v(x))

        if kv_cache is None:
            k, v = k_new, v_new
        else:
            k, v = kv_cache
            if use_cache:
                k = torch.cat([k, k_new], dim=2)
                v = torch.cat([v, v_new], dim=2)

        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)

        if q.size(2) == k.size(2):
            if self.use_relpos and rel_bias_tokens is not None:
                att = att + alibi_bias(self.h, rel_bias_tokens)
            if mask is not None:
                att = att + mask

        z = (att.softmax(-1) @ v).transpose(1, 2)
        z = z.reshape(x.size(0), x.size(1), -1)
        out = self.drop(self.proj(z))
        return (out, (k, v)) if use_cache else out

class Block(nn.Module):
    def __init__(self, d: int, h: int, r: int):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.mha = LowRankMHA(d, h, r, use_relpos=True)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d))

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor],

                kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,

                use_cache: bool = False):
        n = x.size(1)
        if use_cache:
            y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=n if mask is not None else None, kv_cache=kv, use_cache=True)
            x = x + y
            x = x + self.ff(self.ln2(x))
            return x, new_kv
        else:
            x = x + self.mha(self.ln1(x), mask, rel_bias_tokens=n)
            return x + self.ff(self.ln2(x))

class Encoder(nn.Module):
    def __init__(self, cfg: Dict[str, int]):
        super().__init__()
        d, l, h, r = cfg["d"], cfg["layers"], cfg["heads"], cfg["rank"]
        self.emb = nn.Embedding(VOCAB, d)
        self.blocks = nn.ModuleList([Block(d, h, r) for _ in range(l)])
        self.ln = nn.LayerNorm(d)

    def forward(self, ids: torch.Tensor, mask: Optional[torch.Tensor],

                kv_caches: Optional[List[Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None,

                use_cache: bool = False):
        x = self.emb(ids)
        if not use_cache:
            for blk in self.blocks:
                x = blk(x, mask)
            return self.ln(x)
        new_kvs: List[Tuple[torch.Tensor, torch.Tensor]] = []
        for i, blk in enumerate(self.blocks):
            kv = kv_caches[i] if (kv_caches is not None) else None
            x, kv_out = blk(x, mask, kv, use_cache=True)
            new_kvs.append(kv_out)
        return self.ln(x), new_kvs

class ARHead(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.proj = nn.Linear(d, VOCAB)
    def forward(self, h): return self.proj(h)

# ───────────────────────── Masks ─────────────────────────
def causal_mask(n):
    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
    return torch.triu(m, 1)

# ───────────────────────── Checkpoint helpers ─────────────────────────
def save_ckpt(path: pathlib.Path, core: nn.Module, ar_h: nn.Module,

              opt: torch.optim.Optimizer, scaler: GradScaler, meta: Dict[str, Any]):
    path.parent.mkdir(exist_ok=True, parents=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    state = {
        "core": core.state_dict(),
        "ar": ar_h.state_dict(),
        "opt": opt.state_dict(),
        "scaler": scaler.state_dict(),
        "cfg": meta.get("cfg"),
        "tokenizer_id": TOKENIZER_ID,
        **{k: v for k, v in meta.items() if k not in {"cfg"}},
    }
    torch.save(state, tmp, _use_new_zipfile_serialization=False)
    tmp.replace(path)
    (path.parent / "latest.json").write_text(json.dumps({"path": str(path), "step": meta["step"]}))
    print(f"\n✓ saved checkpoint {path.name}")

def load_ckpt(path: pathlib.Path, core: nn.Module, ar_h: nn.Module,

              opt: torch.optim.Optimizer, scaler: GradScaler):
    p = _resolve_ckpt(path) or path
    ck = _try_load(p, map_location="cpu")
    if ck is None:
        raise FileNotFoundError(f"No valid checkpoint at {p}")
    core.load_state_dict(ck["core"])
    if "ar" in ck:
        ar_h.load_state_dict(ck["ar"])
    opt.load_state_dict(ck["opt"])
    scaler.load_state_dict(ck["scaler"])
    return ck.get("step", 0), ck.get("seen_tok", 0), ck.get("wall_time", time.time())

def _safe_load_any(path: pathlib.Path, tgt: nn.Module, key: str | None = None, rename: str | None = None):
    p = _resolve_ckpt(path) or path
    if not p or not p.exists(): return 0
    ck = _try_load(p, map_location="cpu")
    if ck is None: return 0
    sd = ck.get(key, ck) if key else ck
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    if rename:
        sd = {k.replace(rename, "proj."): v for k, v in sd.items() if rename in k}
    tgt_sd = tgt.state_dict()
    filt = {k: v for k, v in sd.items() if k in tgt_sd and v.shape == tgt_sd[k].shape}
    if filt:
        tgt.load_state_dict(filt, strict=False)
    return len(filt)

def infer_cfg_from_ckpt(path: pathlib.Path):
    p = _resolve_ckpt(path) or path
    if not p.exists(): return None
    sd = _try_load(p, map_location="cpu")
    if sd is None: return None
    if isinstance(sd, dict) and "cfg" in sd and isinstance(sd["cfg"], dict):
        return dict(sd["cfg"])
    core = sd.get("core")
    if core is None: return None
    emb_w = core.get("emb.weight")
    if emb_w is None: return None
    d = emb_w.shape[1]
    layer_ids = []
    for k in core.keys():
        if k.startswith("blocks."):
            parts = k.split(".")
            if len(parts) > 2 and parts[1].isdigit():
                layer_ids.append(int(parts[1]))
    layers = (max(layer_ids) + 1) if layer_ids else None
    U = core.get("blocks.0.mha.U")
    heads = rank = None
    if U is not None:
        dk, r = U.shape
        rank = r
        heads = d // dk if dk > 0 else None
    out = {"d": d}
    if layers is not None: out["layers"] = layers
    if heads is not None:  out["heads"] = heads
    if rank is not None:   out["rank"] = rank
    return out

# ───────────────────────── Train loop ─────────────────────────
def _parse_grow_plan(s: str) -> List[int]:
    steps = []
    for part in s.split(","):
        part = part.strip()
        if part:
            v = int(part)
            if v >= 128:
                steps.append(v)
    return sorted(set(steps))

def _init_save_timers(resume_wall_time: float | None, interval_sec: int) -> Tuple[float, float]:
    now_wall = time.time()
    now_mono = time.monotonic()
    if resume_wall_time is None:
        return now_wall, now_mono
    elapsed_wall = max(0.0, now_wall - resume_wall_time)
    elapsed_clamped = min(float(interval_sec), elapsed_wall)
    return now_wall, now_mono - elapsed_clamped

def _count_enabled_params(*modules: Optional[nn.Module]) -> int:
    total = 0
    for m in modules:
        if m is not None:
            total += sum(p.numel() for p in m.parameters())
    return total

def train(args):
    cfg = PRESETS[args.preset].copy()

    # Previous topology probe (unless --fresh)
    if not args.fresh:
        src_probe = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
        prev_cfg = infer_cfg_from_ckpt(src_probe)
    else:
        prev_cfg = None

    if prev_cfg:
        cfg["d"] = prev_cfg.get("d", cfg["d"])
        if prev_cfg.get("heads"): cfg["heads"] = prev_cfg["heads"]
        if args.rank is None and prev_cfg.get("rank"): cfg["rank"] = prev_cfg["rank"]
        if prev_cfg.get("layers"): cfg["layers"] = prev_cfg["layers"]
        if args.x2 and prev_cfg.get("layers"): cfg["layers"] = max(cfg["layers"], prev_cfg["layers"] * 2)
    if args.rank: cfg["rank"] = args.rank
    if args.x2 and not prev_cfg: cfg["layers"] *= 2

    BLOCK = args.block or DEFAULT_BLOCK

    core = Encoder(cfg).to(DEV)
    ar_h = ARHead(cfg["d"]).to(DEV)

    # Warm start unless --fresh
    loaded = 0
    if not args.fresh:
        src = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
        src = _resolve_ckpt(src)
        if src:
            loaded += _safe_load_any(src, core, key="core")
            loaded += _safe_load_any(src, ar_h, key="ar")
            if loaded:
                print(f"Warm-start: loaded {loaded} matching tensors from {src}")

    # Optional: freeze core; selectively unfreeze LayerNorm and/or embeddings for SFT
    if args.freeze_core:
        for p in core.parameters(): p.requires_grad = False
        if args.unfreeze_ln:
            for blk in core.blocks:
                for p in blk.ln1.parameters(): p.requires_grad = True
                for p in blk.ln2.parameters(): p.requires_grad = True
            for p in core.ln.parameters(): p.requires_grad = True
        if args.train_emb:
            for p in core.emb.parameters(): p.requires_grad = True

    # Optimizer (respect requires_grad)
    opt = torch.optim.AdamW([
        {"params": [p for p in core.parameters() if p.requires_grad], "lr": args.lr_core},
        {"params": ar_h.parameters(), "lr": args.lr_head},
    ])
    scaler = GradScaler(enabled=((args.amp or args.fp8_only) and DEV.type == "cuda"))
    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)

    # ---------- resume bookkeeping ----------
    start_step, seen_tok = 0, 0
    last_save_wall = None
    if args.resume and not args.fresh:
        start_step, seen_tok, last_save_wall = load_ckpt(pathlib.Path(args.resume), core, ar_h, opt, scaler)
        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)

    # Chinchilla-style target tokens: ALL enabled params (core + ar head)
    if args.target_tokens:
        target_tokens = args.target_tokens
    else:
        enabled_param_count = _count_enabled_params(core, ar_h)
        target_tokens = int(25 * enabled_param_count)

    new_tokens_needed = target_tokens - seen_tok
    if new_tokens_needed <= 0:
        print("Target already reached – nothing to train.")
        return
    new_steps = new_tokens_needed // BLOCK
    if args.steps:
        new_steps = min(new_steps, args.steps)
        new_tokens_needed = new_steps * BLOCK

    total_tokens_needed = seen_tok + new_tokens_needed
    print(f"[auto-steps] {new_steps:,} training steps (@ {BLOCK} tokens/step)")

    # Progressive growth plan
    grow_plan = _parse_grow_plan(args.grow_plan) if args.auto_grow else []
    if args.auto_grow:
        if BLOCK not in grow_plan:
            grow_plan = sorted(set(grow_plan + [BLOCK]))
        print(f"[auto-grow] plan: {grow_plan} every {args.grow_every_steps} steps")

    # FP8 guard
    if args.fp8_only and not _supports_fp8() and not args.fp8_fallback:
        raise RuntimeError("FP8 not supported by your torch build/hardware. Use --fp8-fallback to continue with bf16.")

    stream = token_stream(args, target_tokens, seed=42)
    buf: list[int] = []
    pbar = tqdm(total=total_tokens_needed, initial=seen_tok, unit="tok")
    step = start_step
    steps_since_last_grow = 0

    while seen_tok < total_tokens_needed:
        # ------- assemble one batch -------
        try:
            while len(buf) < BLOCK:
                buf.append(next(stream))
        except StopIteration:
            break
        ids = torch.tensor(buf[:BLOCK], device=DEV).unsqueeze(0)  # (B=1, N)
        buf = buf[BLOCK:]

        tgt_ar = ids.clone()

        try:
            with amp(args.amp or args.fp8_only, prefer_fp8=args.fp8_only and (_supports_fp8() or args.fp8_fallback)):
                h_ar = core(ids, causal_mask(ids.size(1)))
                logits_ar = ar_h(h_ar)[:, :-1]
                loss = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))

            scaler.scale(loss).backward()
            scaler.unscale_(opt)
            nn.utils.clip_grad_norm_(core.parameters(), 1.0)
            scaler.step(opt)
            scaler.update()
            opt.zero_grad(set_to_none=True)

        except RuntimeError as e:
            msg = str(e).lower()
            if "out of memory" in msg or "cuda error" in msg:
                new_block = max(128, BLOCK // 2)
                if new_block < BLOCK:
                    print(f"\n[OOM] reducing block from {BLOCK} -> {new_block}")
                    BLOCK = new_block
                    if DEV.type == "cuda":
                        torch.cuda.empty_cache()
                    buf = ids[0].tolist() + buf
                    steps_since_last_grow = 0
                    continue
            raise

        # progress
        step += 1
        seen_tok += BLOCK
        pbar.update(BLOCK)
        pbar.set_postfix(loss=f"{loss.item():.3f}", block=BLOCK)

        # time-based checkpoint cadence only (monotonic)
        if args.save_every_sec > 0:
            now_mono = time.monotonic()
            if now_mono - last_save_mono >= args.save_every_sec:
                ck_name = f"step{step:08d}.pt"
                save_ckpt(
                    pathlib.Path(args.save_dir) / ck_name,
                    core, ar_h, opt, scaler,
                    meta={
                        "cfg": cfg,
                        "step": step,
                        "seen_tok": seen_tok,
                        "wall_time": time.time(),
                        "py_state": random.getstate(),
                        "torch_state": rng_state(),
                        "fp8_only": args.fp8_only,
                    },
                )
                last_save_mono = now_mono

        # optional step-based checkpoint cadence
        if args.save_every_steps > 0 and step > 0 and (step % args.save_every_steps == 0):
            ck_name = f"step{step:08d}.pt"
            save_ckpt(
                pathlib.Path(args.save_dir) / ck_name,
                core, ar_h, opt, scaler,
                meta={
                    "cfg": cfg,
                    "step": step,
                    "seen_tok": seen_tok,
                    "wall_time": time.time(),
                    "py_state": random.getstate(),
                    "torch_state": rng_state(),
                    "fp8_only": args.fp8_only,
                },
            )

        # progressive growth
        if args.auto_grow:
            steps_since_last_grow += 1
            if steps_since_last_grow >= args.grow_every_steps:
                steps_since_last_grow = 0
                try:
                    idx = grow_plan.index(BLOCK)
                    if idx + 1 < len(grow_plan):
                        candidate = grow_plan[idx + 1]
                        print(f"[auto-grow] attempting BLOCK {BLOCK} -> {candidate}")
                        BLOCK = candidate
                        if DEV.type == "cuda":
                            torch.cuda.empty_cache()
                    else:
                        print("[auto-grow] at max planned block; no further growth.")
                except ValueError:
                    grow_plan = sorted(set(grow_plan + [BLOCK]))
                    idx = grow_plan.index(BLOCK)
                    if idx + 1 < len(grow_plan):
                        candidate = grow_plan[idx + 1]
                        print(f"[auto-grow] moving to planned BLOCK {candidate}")
                        BLOCK = candidate
                        if DEV.type == "cuda":
                            torch.cuda.empty_cache()

    pbar.close()

    # final save
    save_ckpt(
        pathlib.Path(args.save_dir) / "final.pt",
        core, ar_h, opt, scaler,
        meta={
            "cfg": cfg,
            "step": step,
            "seen_tok": seen_tok,
            "wall_time": time.time(),
            "py_state": random.getstate(),
            "torch_state": rng_state(),
            "fp8_only": args.fp8_only,
        },
    )
    print("🎉 training complete")

# ───────────────────────── Sampling utils ─────────────────────────
def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
    if n <= 0 or ids.size(1) < n - 1:
        return logits
    prefix = ids[0, - (n - 1):].tolist()
    banned = []
    tokens = ids[0].tolist()
    for i in range(len(tokens) - n + 1):
        if tokens[i:i + n - 1] == prefix:
            banned.append(tokens[i + n - 1])
    if banned:
        banned_idx = torch.tensor(banned, device=logits.device, dtype=torch.long)
        logits[..., banned_idx] = float("-inf")
    return logits

def _apply_rep_presence_frequency(

    logits: torch.Tensor, ids: torch.Tensor, last_n: int,

    repetition_penalty: float, presence_penalty: float, frequency_penalty: float

):
    if ids.numel() == 0:
        return logits
    hist = ids[0, -last_n:].to(torch.long) if last_n > 0 else ids[0].to(torch.long)
    if hist.numel() == 0:
        return logits
    uniq, counts = torch.unique(hist, return_counts=True)
    if presence_penalty != 0.0 or frequency_penalty != 0.0:
        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
        logits[..., uniq] = logits[..., uniq] - adjust
    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
        sel = logits[..., uniq]
        sel = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty)
        logits[..., uniq] = sel
    return logits

def _filter_top_k_top_p_min_p(

    logits: torch.Tensor, top_k: int, top_p: float, min_p: float, temperature: float

) -> torch.Tensor:
    logits = logits / max(temperature, 1e-8)
    if logits.dim() == 1:
        logits = logits.unsqueeze(0)
    probs = logits.softmax(-1)

    V = probs.size(-1)
    if top_k and top_k < V:
        vals, idx = torch.topk(probs, top_k, dim=-1)
        mask = torch.full_like(probs, 0.0)
        mask.scatter_((1 if probs.dim() == 2 else -1), idx, 1.0)
        probs = probs * mask

    if top_p < 1.0:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
        cumsum = torch.cumsum(sorted_probs, dim=-1)
        keep = cumsum <= top_p
        keep[..., 0] = True
        mask = torch.zeros_like(probs)
        mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
        probs = probs * mask

    if min_p > 0.0:
        probs = torch.where(probs >= min_p, probs, torch.zeros_like(probs))

    sums = probs.sum(-1, keepdim=True)
    empty = (sums == 0)
    if empty.any():
        fallback_idx = logits.argmax(-1, keepdim=True)
        probs = torch.where(empty, torch.zeros_like(probs), probs)
        probs.scatter_(-1, fallback_idx, torch.where(empty, torch.ones_like(sums), torch.zeros_like(sums)))

    probs = probs / probs.sum(-1, keepdim=True)
    return probs

# ───────────────────────── Inference helpers ─────────────────────────
def load_joint(ckpt: str, preset: str):
    path = _resolve_ckpt(pathlib.Path(ckpt)) or pathlib.Path(ckpt)
    sd = _try_load(path, map_location="cpu")
    if sd is None:
        raise FileNotFoundError(f"No valid checkpoint at {path}")
    cfg = sd["cfg"] if "cfg" in sd and isinstance(sd["cfg"], dict) else (infer_cfg_from_ckpt(path) or PRESETS[preset])
    core = Encoder(cfg).to(DEV)
    ar_h = ARHead(cfg["d"]).to(DEV)
    core.load_state_dict(sd["core"])
    if "ar" in sd:
        ar_h.load_state_dict(sd["ar"])
    return core, ar_h

@torch.no_grad()
def ar_decode(core, ar_h, prompt: str, max_new: int, T: float,

              greedy: bool, top_k: int, top_p: float, min_p: float,

              repetition_penalty: float, presence_penalty: float,

              frequency_penalty: float, penalty_last_n: int,

              no_repeat_ngram_size: int,

              use_fp8: bool, fp8_fallback: bool):
    # Tokenize prompt and remember its length
    prompt_ids = tok.encode(prompt)
    if len(prompt_ids) == 0:
        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
        prompt_len = 0
    else:
        ids = torch.tensor([prompt_ids], device=DEV)
        prompt_len = ids.size(1)

    t0 = time.time()
    with amp(use_fp8 or False, prefer_fp8=use_fp8 and (_supports_fp8() or fp8_fallback)):
        h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
        for _ in range(max_new):
            logits = ar_h(h_full)[:, -1]
            logits = _apply_no_repeat_ngram(logits, ids, no_repeat_ngram_size)
            logits = _apply_rep_presence_frequency(
                logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
            )
            if greedy:
                nxt = logits.argmax(-1, keepdim=True)
            else:
                probs = _filter_top_k_top_p_min_p(logits.squeeze(0), top_k, top_p, min_p, T)
                nxt = probs.multinomial(1)
            ids = torch.cat([ids, nxt.unsqueeze(0) if nxt.dim()==1 else nxt], 1)
            x = ids[:, -1:]
            h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)

    # Decode prompt vs generation separately
    full_ids = ids[0].tolist()
    prompt_text = tok.decode(full_ids[:prompt_len], skip_special_tokens=True)
    gen_text    = tok.decode(full_ids[prompt_len:], skip_special_tokens=True)

    if sys.stdout.isatty():
        sys.stdout.write("\x1b[90m")
        sys.stdout.write(prompt_text)
        sys.stdout.write("\x1b[0m")
        sys.stdout.write(gen_text + "\n")
    else:
        sys.stdout.write(prompt_text + gen_text + "\n")

    print(f"[{len(full_ids) - prompt_len} tok in {time.time() - t0:.2f}s]")

# ───────────────────────── CLI ─────────────────────────
def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    tr = sub.add_parser("train")
    tr.add_argument("--preset", choices=PRESETS, default="small")
    tr.add_argument("--rank", type=int)
    tr.add_argument("--block", type=int, default=DEFAULT_BLOCK)
    tr.add_argument("--source", default="cerebras/SlimPajama-627B",
                    help="Comma-separated datasets (optionally dataset:config), or json:/path.jsonl")
    tr.add_argument("--target_tokens", type=int)
    tr.add_argument("--steps", type=int)
    tr.add_argument("--amp", action="store_true")
    tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC)
    tr.add_argument("--save_every_steps", type=int, default=0, help="Also checkpoint every N steps (0 = disabled)")
    tr.add_argument("--save_dir", default=str(CKDIR))
    tr.add_argument("--resume", type=str)
    tr.add_argument("--x2", action="store_true", help="~2x params by doubling layers")
    tr.add_argument("--warmstart_from", type=str, default=None, help="Path to previous final.pt for shape-safe warm start")
    tr.add_argument("--fresh", action="store_true", help="Start from scratch: do not probe or load any checkpoints")
    # FP8 control
    tr.add_argument("--fp8-only", action="store_true", dest="fp8_only", help="Attempt FP8 autocast (float8_e4m3fn) for compute")
    tr.add_argument("--fp8-fallback", action="store_true", dest="fp8_fallback", help="If FP8 unsupported, fall back to bf16 instead of erroring")
    # Progressive block growth
    tr.add_argument("--auto_grow", action="store_true", help="Automatically grow block size over time")
    tr.add_argument("--grow_plan", type=str, default="576,640,768,896,1024", help="Comma list of block sizes to try in order")
    tr.add_argument("--grow_every_steps", type=int, default=50000, help="Steps between growth attempts")
    # Chat / dataset fields
    tr.add_argument("--chat", action="store_true", help="Treat rows as chat and render via tokenizer chat template")
    tr.add_argument("--chat_messages_key", type=str, default="messages", help="Field name with list[{role,content}]")
    tr.add_argument("--dataset_field_text", type=str, default="text", help="Field to read when not in --chat mode")
    tr.add_argument("--sft_add_generation_prompt", action="store_true", help="Pass add_generation_prompt=True to chat template")
    # Freezing / LR overrides
    tr.add_argument("--freeze_core", action="store_true", help="Freeze encoder (core) weights for SFT")
    tr.add_argument("--unfreeze_ln", action="store_true", help="When freezing core, still train LayerNorms")
    tr.add_argument("--train_emb", action="store_true", help="When freezing core, also train token embeddings")
    tr.add_argument("--lr_core", type=float, default=LR_CORE, help="LR for core (trainable subset)")
    tr.add_argument("--lr_head", type=float, default=LR_HEAD, help="LR for AR head")

    inf = sub.add_parser("infer")
    inf.add_argument("--mode", choices=["ar"], required=True)
    inf.add_argument("--ckpt", required=True)
    inf.add_argument("--preset", default="small")
    inf.add_argument("--prompt", required=True)
    inf.add_argument("--max_new", type=int, default=120)
    inf.add_argument("--temperature", type=float, default=1.0)

    # Decode controls
    inf.add_argument("--greedy", action="store_true", help="Greedy decode (overrides sampling)")
    inf.add_argument("--top_k", type=int, default=0)
    inf.add_argument("--top_p", type=float, default=1.0)
    inf.add_argument("--min_p", type=float, default=0.0)
    inf.add_argument("--repetition_penalty", type=float, default=1.0)
    inf.add_argument("--presence_penalty", type=float, default=0.0)
    inf.add_argument("--frequency_penalty", type=float, default=0.0)
    inf.add_argument("--penalty_last_n", type=int, default=64)
    inf.add_argument("--no_repeat_ngram_size", type=int, default=0)

    # Inference FP8
    inf.add_argument("--fp8-only", action="store_true", dest="fp8_only", help="Attempt FP8 autocast during decode")
    inf.add_argument("--fp8-fallback", action="store_true", default=False, dest="fp8_fallback", help=argparse.SUPPRESS)

    args = ap.parse_args()
    if args.cmd == "train":
        if args.fp8_only:
            print("[init] FP8-only requested. If FP8 kernels are missing, using --fp8-fallback will continue with bf16.")
        train(args)
    else:
        core, ar_h = load_joint(args.ckpt, args.preset)
        ar_decode(core, ar_h, args.prompt, args.max_new, args.temperature,
                  args.greedy, args.top_k, args.top_p, args.min_p,
                  args.repetition_penalty, args.presence_penalty,
                  args.frequency_penalty, args.penalty_last_n,
                  args.no_repeat_ngram_size,
                  use_fp8=args.fp8_only, fp8_fallback=args.fp8_fallback if hasattr(args, "fp8_fallback") else False)

if __name__ == "__main__":
    main()