Upload 11 files

5acp.py

@@ -0,0 +1,971 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# ALiBi: fixed to work with KV cache (q_len ≠ k_len), mask may be None during cached decode.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"], add_special_tokens=False)
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+
+# ───────────────────────── Relative positional bias (ALiBi) ─────────────────────────
+def _alibi_slopes(n_heads: int, device=None):
+    """
+    Return shape (1, h, 1, 1) slopes tensor on device.
+    """
+    device = device or DEV
+    import math as _m
+    def pow2slopes(n):
+        start = 2 ** (-2 ** -(_m.log2(n) - 3))
+        ratio = start
+        return [start * (ratio ** i) for i in range(n)]
+    if float(int(_m.log2(n_heads))) == _m.log2(n_heads):
+        vals = pow2slopes(n_heads)
+    else:
+        closest = 2 ** int(_m.floor(_m.log2(n_heads)))
+        vals = pow2slopes(closest)
+        extra = pow2slopes(2 * closest)
+        vals += extra[0::2][: n_heads - closest]
+    return torch.tensor(vals, device=device).view(1, n_heads, 1, 1)
+
+def alibi_bias_qk(n_heads: int, q_len: int, k_len: int, device=None):
+    """
+    Build ALiBi bias for arbitrary q_len × k_len with causal structure.
+    Returns shape (1, h, q_len, k_len).
+    """
+    device = device or DEV
+    i = torch.arange(q_len, device=device).view(1, 1, q_len, 1)   # queries
+    j = torch.arange(k_len, device=device).view(1, 1, 1, k_len)   # keys
+    dist = (j - i).clamp_min(0)                                   # only penalize future
+    slopes = _alibi_slopes(n_heads, device=device)
+    return -slopes * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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,  # kept for compat; unused
+        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)
+
+        # (B, h, Nq, Nk)
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)
+
+        # Always add ALiBi relative bias if enabled, even when q_len ≠ k_len or mask is None
+        if self.use_relpos:
+            q_len = q.size(2)
+            k_len = k.size(2)
+            att = att + alibi_bias_qk(self.h, q_len, k_len, device=att.device)
+
+        # Apply mask if provided (square causal or SAT masks in non-cached passes)
+        if mask is not None:
+            att = att + mask
+
+        z = (att.softmax(-1) @ v).transpose(1, 2)  # (B,Nq,h,r)
+        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=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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    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 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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = torch.tensor([ids_nat.size(1)], device=DEV)
+                tlen = torch.tensor([tgt_ar.size(1)], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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, nat_h, sat_h, opt, scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "py_state": random.getstate(),
+            "torch_state": rng_state(),
+        },
+    )
+    print("🎉 training complete")
+
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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), idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+
+    # Min-p
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+
+if __name__ == "__main__":
+    main()

5ap1.py

@@ -0,0 +1,966 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+
+# ───────────────────────── 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)  # only penalize future
+    slopes = _alibi_slopes(n_heads)
+    return -slopes * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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)  # (B,Nq,h,r)
+        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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    now_wall = time.time()
+    now_mono = time.monotonic()
+    if resume_wall_time is None:
+        return now_wall, now_mono
+    # How long since the previous save in wall-clock?
+    elapsed_wall = max(0.0, now_wall - resume_wall_time)
+    # Clamp to interval so we don't try to "catch up" multiple times
+    elapsed_clamped = min(float(interval_sec), elapsed_wall)
+    # Pretend we last saved 'elapsed_clamped' ago on the monotonic clock
+    return now_wall, now_mono - elapsed_clamped
+
+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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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, nat_h, sat_h, opt, scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "py_state": random.getstate(),
+            "torch_state": rng_state(),
+        },
+    )
+    print("🎉 training complete")
+
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        # subtract presence for seen tokens; subtract frequency * count
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
+        sel = logits[..., uniq]
+        # if logit > 0: divide by penalty; else multiply by penalty
+        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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    # shape handling
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    # Convert to probabilities for p-based filtering
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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, idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+        # Always keep at least one
+        keep[..., 0] = True
+        # Build mask
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
+        probs = probs * mask
+
+    # Min-p
+    if min_p > 0.0:
+        probs = torch.where(probs >= min_p, probs, torch.zeros_like(probs))
+
+    # If everything zeroed (can happen at extreme settings), fall back to the argmax token
+    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)))
+
+    # Renormalize
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+
+if __name__ == "__main__":
+    main()

5ap1a.py

@@ -0,0 +1,1090 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# NEW: HF token auto-load from ./h.txt or CLI; authenticated streaming; retry/backoff;
+#      optional local snapshot prefetch; fast transfer path.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os
+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
+from huggingface_hub.utils import HfHubHTTPError
+from huggingface_hub import snapshot_download
+
+# ───────────────────────── 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
+
+# Enable fast transfer path for large files unless explicitly disabled
+if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") is None:
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+# Tokenizer ID (can override with env TOKENIZER_ID)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Will be initialized in init_tokenizer() after HF token is set
+tok: Optional[AutoTokenizer] = None
+VOCAB = BLANK = EOS = 0
+
+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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+
+# ───────────────────────── Token / tokenizer setup ─────────────────────────
+def _read_first_line(path: str) -> Optional[str]:
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            line = f.readline().strip()
+            return line if line else None
+    except Exception:
+        return None
+
+def setup_hf_token(cli_token: Optional[str] = None, token_file: Optional[str] = None):
+    """
+    Determine a HF token from: CLI -> env -> CLI file -> ./h.txt, then export HF_TOKEN.
+    Never prints the token.
+    """
+    token = None
+    if cli_token:
+        token = cli_token.strip()
+    elif os.environ.get("HF_TOKEN"):
+        token = os.environ["HF_TOKEN"].strip()
+    elif token_file:
+        token = _read_first_line(token_file)
+    if token is None:
+        token = _read_first_line("h.txt")
+
+    if token:
+        os.environ["HF_TOKEN"] = token  # used by datasets/hub internals
+
+def init_tokenizer():
+    global tok, VOCAB, BLANK, EOS
+    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  # allow new [PAD] if appended
+    BLANK = tok.pad_token_id
+    EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id
+
+
+# ───────────────────────── 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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Data stream ─────────────────────────
+def _resilient_dataset_iter(ds, max_retries: int = 7, base_backoff: float = 2.0):
+    """
+    Iterate a streaming dataset with retry/backoff on transient hub errors.
+    """
+    it = iter(ds)
+    while True:
+        try:
+            yield next(it)
+        except HfHubHTTPError as e:
+            code = getattr(getattr(e, "response", None), "status_code", None)
+            if code in (401, 403, 408, 429, 500, 502, 503, 504) and max_retries > 0:
+                wait = base_backoff * (2 ** (7 - max_retries))
+                print(f"[hub-retry] HTTP {code}; backoff {wait:.1f}s; retries left {max_retries}")
+                time.sleep(wait)
+                max_retries -= 1
+                continue
+            raise
+        except StopIteration:
+            break
+
+def token_stream(
+    ds_name: str,
+    target: int,
+    seed: int = 42,
+    hf_token: Optional[str] = None,
+    prefetch_dir: Optional[str] = None,
+    allow_patterns: Optional[List[str]] = None,
+    shuffle_buf: int = 10_000,
+    max_retries: int = 7,
+):
+    """
+    Stream tokens from a dataset with:
+      • optional authenticated hub access (reduces 403s)
+      • optional local snapshot (prefetch subset to disk)
+      • retry/backoff on transient errors
+    """
+    if tok is None:
+        raise RuntimeError("Tokenizer not initialized. Call init_tokenizer() first.")
+
+    hf_token = (hf_token or os.environ.get("HF_TOKEN")) or None
+    dl_cfg = DownloadConfig(max_retries=10)
+
+    if prefetch_dir:
+        # Pull a slice locally to avoid range-request roulette
+        snapshot_download(
+            repo_id=ds_name,
+            repo_type="dataset",
+            local_dir=prefetch_dir,
+            allow_patterns=allow_patterns or ["train/**"],
+            token=hf_token,
+            max_workers=4,
+        )
+        # Stream from local JSONL.ZST files
+        pattern = os.path.join(prefetch_dir, "train", "**", "*.jsonl.zst")
+        ds = load_dataset(
+            "json",
+            data_files={"train": pattern},
+            split="train",
+            streaming=True,
+        )
+    else:
+        ds = load_dataset(
+            ds_name,
+            split="train",
+            streaming=True,
+            token=hf_token,
+            download_config=dl_cfg,
+        )
+
+    ds = ds.shuffle(buffer_size=shuffle_buf, seed=seed)
+
+    emitted = 0
+    for ex in _resilient_dataset_iter(ds, max_retries=max_retries):
+        txt = ex.get("text") if isinstance(ex, dict) else None
+        if not txt:
+            continue
+        enc = tok.encode(txt)
+        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
+
+
+# ───────────────────────── 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)  # only penalize future
+    slopes = _alibi_slopes(n_heads)
+    return -slopes * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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)  # (B,Nq,h,r)
+        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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    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 train(args):
+    cfg = PRESETS[args.preset].copy()
+
+    # Tokenizer must be ready before model build
+    init_tokenizer()
+
+    # 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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(
+        args.source,
+        target_tokens,
+        seed=42,
+        hf_token=args.hf_token,
+        prefetch_dir=args.prefetch_dir,
+        allow_patterns=args.allow_patterns.split(",") if args.allow_patterns else None,
+        shuffle_buf=args.shuffle_buf,
+        max_retries=args.hf_retries,
+    )
+    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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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, nat_h, sat_h, opt, scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "py_state": random.getstate(),
+            "torch_state": rng_state(),
+        },
+    )
+    print("🎉 training complete")
+
+
+# ────────────────��──────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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, idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+
+    # Min-p
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    if tok is None:
+        raise RuntimeError("Tokenizer not initialized. Call init_tokenizer() first.")
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        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)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    if tok is None:
+        raise RuntimeError("Tokenizer not initialized. Call init_tokenizer() first.")
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    if tok is None:
+        raise RuntimeError("Tokenizer not initialized. Call init_tokenizer() first.")
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    # Hub/dataset robustness
+    tr.add_argument("--hf_token", type=str, default=None, help="HF token for authenticated streaming (overrides env)")
+    tr.add_argument("--hf_token_file", type=str, default=None, help="Path to a file containing the HF token (1st line)")
+    tr.add_argument("--hf_retries", type=int, default=7, help="Retries for transient hub errors")
+    tr.add_argument("--shuffle_buf", type=int, default=10000, help="Streaming shuffle buffer")
+    tr.add_argument("--prefetch_dir", type=str, default=None, help="Optional local snapshot dir (prefetch before training)")
+    tr.add_argument("--allow_patterns", type=str, default=None, help="Comma-separated allow patterns for snapshot_download")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+
+    # Make sure HF token is exported before any hub usage
+    if args.cmd == "train":
+        setup_hf_token(args.hf_token, args.hf_token_file)
+        train(args)
+    else:
+        setup_hf_token(None, None)  # harmless; may help if tokenizer is private
+        init_tokenizer()
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+
+if __name__ == "__main__":
+    main()

Av2.py

@@ -0,0 +1,977 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# ALiBi: fixed to work with KV cache (q_len ≠ k_len), absolute-position offset, correct sign,
+#        and attention scaling by sqrt(projected_rank) instead of sqrt(head_dim).
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"], add_special_tokens=False)
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+
+# ───────────────────────── Relative positional bias (ALiBi) ─────────────────────────
+def _alibi_slopes(n_heads: int, device=None):
+    """
+    Return shape (1, h, 1, 1) slopes tensor on device.
+    """
+    device = device or DEV
+    import math as _m
+    def pow2slopes(n):
+        start = 2 ** (-2 ** -(_m.log2(n) - 3))
+        ratio = start
+        return [start * (ratio ** i) for i in range(n)]
+    if float(int(_m.log2(n_heads))) == _m.log2(n_heads):
+        vals = pow2slopes(n_heads)
+    else:
+        closest = 2 ** int(_m.floor(_m.log2(n_heads)))
+        vals = pow2slopes(closest)
+        extra = pow2slopes(2 * closest)
+        vals += extra[0::2][: n_heads - closest]
+    return torch.tensor(vals, device=device).view(1, n_heads, 1, 1)
+
+def alibi_bias_qk(n_heads: int, q_len: int, k_len: int, q_offset: int = 0, device=None):
+    """
+    Build ALiBi bias for arbitrary q_len × k_len with causal structure.
+    Positions are absolute: queries start at q_offset, keys start at 0.
+    Penalize older past more strongly; most-recent past has smallest penalty.
+    Returns shape (1, h, q_len, k_len).
+    """
+    device = device or DEV
+    # absolute positions
+    i = torch.arange(q_offset, q_offset + q_len, device=device).view(1, 1, q_len, 1)  # queries
+    j = torch.arange(0, k_len, device=device).view(1, 1, 1, k_len)                    # keys
+    # distance into the past from query to key
+    dist = (i - j).clamp_min(0)
+    slopes = _alibi_slopes(n_heads, device=device)
+    return -slopes * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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,  # kept for compat; unused
+        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ):
+        q = self._proj(self.q(x))      # (B, h, Nq, r)
+        k_new = self._proj(self.k(x))  # (B, h, Nk_new, r)
+        v_new = self._proj(self.v(x))  # (B, h, Nk_new, r)
+
+        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)
+
+        # (B, h, Nq, Nk), scale by projected rank r
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(q.size(-1))
+
+        # Add ALiBi relative bias with absolute offset for cache
+        if self.use_relpos:
+            q_len = q.size(2)
+            k_len = k.size(2)
+            q_off = k_len - q_len if use_cache else 0
+            att = att + alibi_bias_qk(self.h, q_len, k_len, q_offset=q_off, device=att.device)
+
+        # Apply mask if provided (square causal or SAT masks in non-cached passes)
+        if mask is not None:
+            att = att + mask
+
+        z = (att.softmax(-1) @ v).transpose(1, 2)  # (B,Nq,h,r)
+        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=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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    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 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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = torch.tensor([ids_nat.size(1)], device=DEV)
+                tlen = torch.tensor([tgt_ar.size(1)], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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, nat_h, sat_h, opt, scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "py_state": random.getstate(),
+            "torch_state": rng_state(),
+        },
+    )
+    print("🎉 training complete")
+
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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), idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+
+    # Min-p
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+
+if __name__ == "__main__":
+    main()

G.py

@@ -0,0 +1,1017 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# NEW: Graceful shutdown: catches SIGINT/SIGTERM, writes an atomic "interrupt.pt", then exits.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os, sys, signal, atexit, threading, traceback
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+# Interrupt state
+_interrupt_flag = threading.Event()
+_interrupt_reason = {"sig": None, "trace": None}
+_last_emergency_save_mono = 0.0
+
+# ───────────────────────── 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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+# ───────────────────────── 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)  # only penalize future
+    slopes = _alibi_slopes(n_heads)
+    return -slopes * dist
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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)  # (B,Nq,h,r)
+        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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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
+
+# ───────────────────────── Interrupt handling ─────────────────────────
+def _mark_interrupt(sig_name: str):
+    if not _interrupt_flag.is_set():
+        _interrupt_reason["sig"] = sig_name
+        try:
+            _interrupt_reason["trace"] = "".join(traceback.format_stack(limit=5))
+        except Exception:
+            _interrupt_reason["trace"] = None
+        _interrupt_flag.set()
+        print(f"\n[interrupt] received {sig_name}; will save an emergency checkpoint and exit...")
+
+def _install_signal_handlers():
+    def _handler(signum, frame):
+        name = {signal.SIGINT: "SIGINT", signal.SIGTERM: "SIGTERM"}.get(signum, f"SIG{signum}")
+        _mark_interrupt(name)
+    try:
+        signal.signal(signal.SIGINT, _handler)
+    except Exception:
+        pass
+    try:
+        signal.signal(signal.SIGTERM, _handler)
+    except Exception:
+        pass
+
+_install_signal_handlers()
+
+# ───────────────────────── 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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    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 _emergency_save_if_needed(args, meta_basics, core, ar_h, nat_h, sat_h, opt, scaler):
+    """
+    If an interrupt was requested, write ckpts_joint/interrupt.pt atomically and exit.
+    Throttled to avoid duplicate writes on repeated signals.
+    """
+    global _last_emergency_save_mono
+    if not _interrupt_flag.is_set():
+        return False
+    now = time.monotonic()
+    if now - _last_emergency_save_mono < 1.0:
+        return True
+    _last_emergency_save_mono = now
+    out_dir = pathlib.Path(args.save_dir)
+    out_path = out_dir / "interrupt.pt"
+    meta = {
+        **meta_basics,
+        "interrupt": {
+            "sig": _interrupt_reason.get("sig"),
+            "trace": _interrupt_reason.get("trace"),
+            "wall_time": time.time(),
+        },
+    }
+    try:
+        save_ckpt(out_path, core, ar_h, nat_h, sat_h, opt, scaler, meta)
+        print("🛑 emergency checkpoint written; exiting due to interrupt.")
+    except Exception as e:
+        print(f"[interrupt-save-failed] {e}")
+    return True
+
+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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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
+
+    # register atexit for best-effort save if Python exits normally (not SIGKILL/power loss)
+    def _atexit_note():
+        if _interrupt_flag.is_set():
+            print("[atexit] process exiting after interrupt; latest emergency checkpoint already attempted.")
+    atexit.register(_atexit_note)
+
+    while seen_tok < total_tokens_needed:
+        # Check for interrupt before assembling batch
+        if _emergency_save_if_needed(
+            args,
+            meta_basics={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                         "py_state": random.getstate(), "torch_state": rng_state()},
+            core=core, ar_h=ar_h, nat_h=nat_h, sat_h=sat_h, opt=opt, scaler=scaler
+        ):
+            return
+
+        # ------- 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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 (only if not interrupted)
+    if not _interrupt_flag.is_set():
+        save_ckpt(
+            pathlib.Path(args.save_dir) / "final.pt",
+            core, ar_h, nat_h, sat_h, opt, scaler,
+            meta={
+                "cfg": cfg,
+                "step": step,
+                "seen_tok": seen_tok,
+                "wall_time": time.time(),
+                "py_state": random.getstate(),
+                "torch_state": rng_state(),
+            },
+        )
+        print("🎉 training complete")
+    else:
+        print("Ended after interrupt; final save skipped (emergency checkpoint already written).")
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        # subtract presence for seen tokens; subtract frequency * count
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
+        sel = logits[..., uniq]
+        # if logit > 0: divide by penalty; else multiply by penalty
+        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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    # shape handling
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    # Convert to probabilities for p-based filtering
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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, idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+        # Always keep at least one
+        keep[..., 0] = True
+        # Build mask
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
+        probs = probs * mask
+
+    # Min-p
+    if min_p > 0.0:
+        probs = torch.where(probs >= min_p, torch.zeros_like(probs) + probs, torch.zeros_like(probs))
+
+    # If everything zeroed (can happen at extreme settings), fall back to the argmax token
+    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)))
+
+    # Renormalize
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+if __name__ == "__main__":
+    main()

ap.py

@@ -0,0 +1,999 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# UPDATE2: seed support for train and infer; seeds Python/NumPy/Torch and dataset shuffle.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+
+# ───────────────────────── Utilities ─────────────────────────
+def set_seed(seed: Optional[int]) -> None:
+    """
+    Best-effort reproducibility. Seeds Python/NumPy/Torch and toggles CuDNN to deterministic.
+    If seed is None, does nothing.
+    """
+    if seed is None:
+        return
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    import numpy as np
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    try:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    except Exception:
+        pass
+
+
+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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+
+# ───────────────────────── 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)  # only penalize future
+    slopes = _alibi_slopes(n_heads)
+    return -slopes * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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)  # (B,Nq,h,r)
+        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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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,
+    nat_h: nn.Module,
+    sat_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"])
+    ar_h.load_state_dict(ck["ar"])
+    nat_h.load_state_dict(ck["nat"])
+    sat_h.load_state_dict(ck["sat"])
+    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.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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    now_wall = time.time()
+    now_mono = time.monotonic()
+    if resume_wall_time is None:
+        return now_wall, now_mono
+    # How long since the previous save in wall-clock?
+    elapsed_wall = max(0.0, now_wall - resume_wall_time)
+    # Clamp to interval so we don't try to "catch up" multiple times
+    elapsed_clamped = min(float(interval_sec), elapsed_wall)
+    # Pretend we last saved 'elapsed_clamped' ago on the monotonic clock
+    return now_wall, now_mono - elapsed_clamped
+
+def train(args):
+    set_seed(args.seed)
+
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, target_tokens, seed=(args.seed if args.seed is not None else 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                        "seed": args.seed,
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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, nat_h, sat_h, opt, scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "py_state": random.getstate(),
+            "torch_state": rng_state(),
+            "seed": args.seed,
+        },
+    )
+    print("🎉 training complete")
+
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        # subtract presence for seen tokens; subtract frequency * count
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
+        sel = logits[..., uniq]
+        # if logit > 0: divide by penalty; else multiply by penalty
+        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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    # shape handling
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    # Convert to probabilities for p-based filtering
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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, idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+        # Always keep at least one
+        keep[..., 0] = True
+        # Build mask
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
+        probs = probs * mask
+
+    # Min-p
+    if min_p > 0.0:
+        probs = torch.where(probs >= min_p, probs, torch.zeros_like(probs))
+
+    # If everything zeroed (can happen at extreme settings), fall back to the argmax token
+    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)))
+
+    # Renormalize
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+    tr.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+    inf.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+
+    # Set seed early for both train and infer
+    if getattr(args, "seed", None) is not None:
+        set_seed(args.seed)
+
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+
+if __name__ == "__main__":
+    main()

ep.py

@@ -0,0 +1,1066 @@
+#!/usr/bin/env python3
+# 5L.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# NEW: Graceful shutdown: catches SIGINT/SIGTERM, writes an atomic "interrupt.pt", then exits.
+# NEW: Checkpoint dir quota: after each save, prune oldest *.pt (not .tmp) if save_dir usage exceeds limit.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os, sys, signal, atexit, threading, traceback
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+TOKENIZER_ID = os.environ.get(
+    "TOKENIZER_ID",
+    "Qwen/Qwen3-235B-A22B-Thinking-2507"
+)
+
+# Some Qwen tokenizers require trust_remote_code
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,   # allow new [PAD] if appended
+    tok.pad_token_id,
+    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),
+}
+
+# Safe default for 1× Tesla P40; override with --block
+DEFAULT_BLOCK = 576
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+# Default interval: 24 hours. Override with --save_every_sec (e.g., 86400).
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+# Interrupt state
+_interrupt_flag = threading.Event()
+_interrupt_reason = {"sig": None, "trace": None}
+_last_emergency_save_mono = 0.0
+
+# ───────────────────────── 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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    # Only enable if explicitly requested AND CUDA is available
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        # ensure EOS between docs
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+# ───────────────────────── 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)  # only penalize future
+    slopes = _alibi_slopes(n_heads)
+    return -slopes * dist
+
+# ───────────────────────── Model components ─────────────────────────
+class LowRankMHA(nn.Module):
+    """
+    Cache-aware MHA with low-rank projections; supports kv caching for decode.
+    """
+    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)  # (B,Nq,h,r)
+        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):
+    """
+    Transformer encoder with optional kv caching (for AR/SAT decode).
+    """
+    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)
+
+class NATHead(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+# ───────────────────────── Checkpoint helpers & quota ─────────────────────────
+def save_ckpt(
+    path: pathlib.Path,
+    core: nn.Module,
+    ar_h: nn.Module,
+    nat_h: nn.Module,
+    sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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 _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.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
+
+# ---- Quota helpers: prune oldest *.pt in a directory when exceeding byte limit
+def _ckpt_dir_usage_bytes(dirpath: pathlib.Path) -> Tuple[int, List[pathlib.Path]]:
+    files = [p for p in dirpath.glob("*.pt") if _is_probably_ckpt(p)]
+    files_sorted = sorted(files, key=lambda p: p.stat().st_mtime)  # oldest first
+    total = 0
+    for p in files_sorted:
+        try:
+            total += p.stat().st_size
+        except Exception:
+            pass
+    return total, files_sorted
+
+def prune_ckpts_to_quota(
+    dirpath: pathlib.Path,
+    quota_bytes: int,
+    delete_only_one: bool = True,
+    protect: Optional[pathlib.Path] = None
+) -> Tuple[List[pathlib.Path], int]:
+    """
+    If total size of *.pt in dirpath exceeds quota_bytes, delete the oldest file(s).
+    If delete_only_one is True, remove just one oldest per invocation.
+    'protect' (if provided) will not be deleted.
+    Returns (deleted_files, new_total_bytes).
+    """
+    deleted: List[pathlib.Path] = []
+    total, files = _ckpt_dir_usage_bytes(dirpath)
+    if protect is not None:
+        try:
+            pr = protect.resolve()
+            files = [f for f in files if f.resolve() != pr]
+        except Exception:
+            pass
+    while total > quota_bytes and files:
+        victim = files.pop(0)  # oldest
+        try:
+            sz = victim.stat().st_size
+        except Exception:
+            sz = 0
+        try:
+            victim.unlink()
+            total -= sz
+            deleted.append(victim)
+            print(f"[quota] removed oldest checkpoint {victim.name} ({sz/ (1024**3):.2f} GiB). Current dir usage ~{total/(1024**3):.2f} GiB.")
+        except Exception as e:
+            print(f"[quota] failed to remove {victim}: {e}")
+        if delete_only_one:
+            break
+    return deleted, total
+
+# ───────────────────────── Interrupt handling ─────────────────────────
+def _mark_interrupt(sig_name: str):
+    if not _interrupt_flag.is_set():
+        _interrupt_reason["sig"] = sig_name
+        try:
+            _interrupt_reason["trace"] = "".join(traceback.format_stack(limit=5))
+        except Exception:
+            _interrupt_reason["trace"] = None
+        _interrupt_flag.set()
+        print(f"\n[interrupt] received {sig_name}; will save an emergency checkpoint and exit...")
+
+def _install_signal_handlers():
+    def _handler(signum, frame):
+        name = {signal.SIGINT: "SIGINT", signal.SIGTERM: "SIGTERM"}.get(signum, f"SIG{signum}")
+        _mark_interrupt(name)
+    try:
+        signal.signal(signal.SIGINT, _handler)
+    except Exception:
+        pass
+    try:
+        signal.signal(signal.SIGTERM, _handler)
+    except Exception:
+        pass
+
+_install_signal_handlers()
+
+# ───────────────────────── 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]:
+    """
+    Returns (last_save_wall, last_save_mono).
+    We use wall time for metadata, monotonic for interval checks.
+    If resuming and the last save was long ago, schedule next save accordingly.
+    """
+    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 _emergency_save_if_needed(args, meta_basics, core, ar_h, nat_h, sat_h, opt, scaler):
+    """
+    If an interrupt was requested, write ckpts_joint/interrupt.pt atomically and exit.
+    Throttled to avoid duplicate writes on repeated signals.
+    """
+    global _last_emergency_save_mono
+    if not _interrupt_flag.is_set():
+        return False
+    now = time.monotonic()
+    if now - _last_emergency_save_mono < 1.0:
+        return True
+    _last_emergency_save_mono = now
+    out_dir = pathlib.Path(args.save_dir)
+    out_path = out_dir / "interrupt.pt"
+    meta = {
+        **meta_basics,
+        "interrupt": {
+            "sig": _interrupt_reason.get("sig"),
+            "trace": _interrupt_reason.get("trace"),
+            "wall_time": time.time(),
+        },
+    }
+    try:
+        save_ckpt(out_path, core, ar_h, nat_h, sat_h, opt, scaler, meta)
+        # enforce quota after emergency save
+        if args.max_save_dir_gb and args.max_save_dir_gb > 0:
+            quota_bytes = int(args.max_save_dir_gb * (1024**3))
+            prune_ckpts_to_quota(out_dir, quota_bytes, delete_only_one=(not args.prune_until_under), protect=out_path)
+        print("🛑 emergency checkpoint written; exiting due to interrupt.")
+    except Exception as e:
+        print(f"[interrupt-save-failed] {e}")
+    return True
+
+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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_h, opt, scaler
+        )
+        print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
+    # Initialize save timers
+    last_save_wall, last_save_mono = _init_save_timers(last_save_wall, args.save_every_sec)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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
+
+    # register atexit for best-effort save if Python exits normally (not SIGKILL/power loss)
+    def _atexit_note():
+        if _interrupt_flag.is_set():
+            print("[atexit] process exiting after interrupt; latest emergency checkpoint already attempted.")
+    atexit.register(_atexit_note)
+
+    while seen_tok < total_tokens_needed:
+        # Check for interrupt before assembling batch
+        if _emergency_save_if_needed(
+            args,
+            meta_basics={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                         "py_state": random.getstate(), "torch_state": rng_state()},
+            core=core, ar_h=ar_h, nat_h=nat_h, sat_h=sat_h, opt=opt, scaler=scaler
+        ):
+            return
+
+        # ------- 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()                           # (1, N)
+        ids_nat = torch.repeat_interleave(ids, 2, 1)   # (1, 2N) for NAT only
+
+        try:
+            with amp(args.amp):
+                # AR path
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+
+                # NAT path (uses doubled sequence)
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)  # (T,B,V)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+
+                # SAT path
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+
+                loss = loss_ar + loss_nat + loss_sat
+
+            # optimisation
+            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"
+                ck_path = pathlib.Path(args.save_dir) / ck_name
+                save_ckpt(
+                    ck_path,
+                    core, ar_h, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                # enforce quota after periodic save
+                if args.max_save_dir_gb and args.max_save_dir_gb > 0:
+                    quota_bytes = int(args.max_save_dir_gb * (1024**3))
+                    prune_ckpts_to_quota(ck_path.parent, quota_bytes, delete_only_one=(not args.prune_until_under), protect=ck_path)
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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 (only if not interrupted)
+    if not _interrupt_flag.is_set():
+        final_path = pathlib.Path(args.save_dir) / "final.pt"
+        save_ckpt(
+            final_path,
+            core, ar_h, nat_h, sat_h, opt, scaler,
+            meta={
+                "cfg": cfg,
+                "step": step,
+                "seen_tok": seen_tok,
+                "wall_time": time.time(),
+                "py_state": random.getstate(),
+                "torch_state": rng_state(),
+            },
+        )
+        # enforce quota after final save
+        if args.max_save_dir_gb and args.max_save_dir_gb > 0:
+            quota_bytes = int(args.max_save_dir_gb * (1024**3))
+            prune_ckpts_to_quota(final_path.parent, quota_bytes, delete_only_one=(not args.prune_until_under), protect=final_path)
+        print("🎉 training complete")
+    else:
+        print("Ended after interrupt; final save skipped (emergency checkpoint already written).")
+
+# ───────────────────────── Sampling utils ─────────────────────────
+def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
+    """
+    Block tokens that would complete any previously seen n-gram.
+    ids: (1, t)
+    logits: (..., V) where ... may be (1,) or (stride,)
+    """
+    if n <= 0 or ids.size(1) < n - 1:
+        return logits
+    prefix = ids[0, - (n - 1):].tolist()
+    # Build set of next tokens forbidden after this prefix.
+    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
+):
+    """
+    logits: (..., V) where ... may be (1,) or (stride,)
+    ids: (1, t) history
+    """
+    if ids.numel() == 0:
+        return logits
+    if last_n > 0:
+        hist = ids[0, -last_n:].to(torch.long)
+    else:
+        hist = ids[0].to(torch.long)
+
+    if hist.numel() == 0:
+        return logits
+
+    uniq, counts = torch.unique(hist, return_counts=True)
+
+    # presence/frequency penalties (OpenAI-like)
+    if presence_penalty != 0.0 or frequency_penalty != 0.0:
+        # subtract presence for seen tokens; subtract frequency * count
+        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
+        logits[..., uniq] = logits[..., uniq] - adjust
+
+    # repetition penalty (CTRL/GPT-NeoX style)
+    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
+        sel = logits[..., uniq]
+        # if logit > 0: divide by penalty; else multiply by penalty
+        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:
+    """
+    Works on 1D or 2D logits (..., V). Applies temperature, then filtering.
+    Returns normalized probabilities ready for sampling.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    # shape handling
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+
+    B, V = logits.size(0), logits.size(-1)
+
+    # Convert to probabilities for p-based filtering
+    probs = logits.softmax(-1)
+
+    # Top-k
+    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, idx, 1.0)
+        probs = probs * mask
+
+    # Top-p (nucleus)
+    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
+        # Always keep at least one
+        keep[..., 0] = True
+        # Build mask
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
+        probs = probs * mask
+
+    # Min-p
+    if min_p > 0.0:
+        probs = torch.where(probs >= min_p, torch.zeros_like(probs) + probs, torch.zeros_like(probs))
+
+    # If everything zeroed (can happen at extreme settings), fall back to the argmax token
+    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)))
+
+    # Renormalize
+    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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+
+    start = time.time()
+    for _ in range(max_new):
+        logits = ar_h(h_full)[:, -1]  # (1, V)
+
+        # penalties
+        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)
+
+        # step with kv cache
+        x = ids[:, -1:]
+        h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - start:.2f}s]")
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])  # (1, SAT_BLOCK, V)
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+
+        # Sequentially sample within the stride so penalties apply cumulatively
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]  # (1, V)
+
+            # penalties
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(
+                row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty
+            )
+
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)  # (1,1)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)  # (1,1)
+
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    # Checkpoint dir quota controls (default 28 GiB per request)
+    tr.add_argument("--max_save_dir_gb", type=float, default=28.0,
+                    help="If > 0, after each save prune oldest *.pt in save_dir when total *.pt size exceeds this many GiB.")
+    tr.add_argument("--prune_until_under", action="store_true",
+                    help="If set, keep deleting oldest checkpoints until usage <= limit. By default only one oldest is removed per save.")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+if __name__ == "__main__":
+    main()

ep1.py

@@ -0,0 +1,861 @@
+#!/usr/bin/env python3
+# ep.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# NEW: Graceful shutdown: catches SIGINT/SIGTERM, writes an atomic "interrupt.pt", then exits.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os, sys, signal, atexit, threading, traceback
+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
+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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,
+    tok.pad_token_id,
+    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
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+# Interrupt state
+_interrupt_flag = threading.Event()
+_interrupt_reason = {"sig": None, "trace": None}
+_last_emergency_save_mono = 0.0
+
+# ───────────────────────── 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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    try:
+        # When PyTorch flips default, this still works. If you want, pass weights_only=True later.
+        return torch.load(path, map_location=map_location)
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+# ───────────────────────── 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, mask=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, mask, kv=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, mask, 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)
+
+class NATHead(nn.Module):
+    def __init__(self, d): super().__init__(); self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(path: pathlib.Path, core: nn.Module, ar_h: nn.Module, nat_h: nn.Module, sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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 != "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, nat_h: nn.Module, sat_h: nn.Module,
+              opt: torch.optim.Optimizer, scaler: GradScaler):
+    """
+    Load a full training state from a checkpoint file or directory.
+    Returns (step, seen_tok, wall_time)
+    """
+    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
+    if "core" in ck: core.load_state_dict(ck["core"])
+    if "ar" in ck:   ar_h.load_state_dict(ck["ar"])
+    if "nat" in ck:  nat_h.load_state_dict(ck["nat"])
+    if "sat" in ck:  sat_h.load_state_dict(ck["sat"])
+    # opt/scaler can be missing if you saved partials; load best-effort
+    try:
+        if "opt" in ck: opt.load_state_dict(ck["opt"])
+    except Exception as e:
+        print(f"[resume] optimizer load skipped: {e}")
+    try:
+        if "scaler" in ck: scaler.load_state_dict(ck["scaler"])
+    except Exception as e:
+        print(f"[resume] scaler load skipped: {e}")
+    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.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
+
+# ───────────────────────── Interrupt handling ─────────────────────────
+def _mark_interrupt(sig_name: str):
+    if not _interrupt_flag.is_set():
+        _interrupt_reason["sig"] = sig_name
+        try:
+            _interrupt_reason["trace"] = "".join(traceback.format_stack(limit=5))
+        except Exception:
+            _interrupt_reason["trace"] = None
+        _interrupt_flag.set()
+        print(f"\n[interrupt] received {sig_name}; will save an emergency checkpoint and exit...")
+
+def _install_signal_handlers():
+    def _handler(signum, frame):
+        name = {signal.SIGINT: "SIGINT", signal.SIGTERM: "SIGTERM"}.get(signum, f"SIG{signum}")
+        _mark_interrupt(name)
+    try: signal.signal(signal.SIGINT, _handler)
+    except Exception: pass
+    try: signal.signal(signal.SIGTERM, _handler)
+    except Exception: pass
+
+_install_signal_handlers()
+
+# ───────────────────────── 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 _emergency_save_if_needed(args, meta_basics, core, ar_h, nat_h, sat_h, opt, scaler):
+    global _last_emergency_save_mono
+    if not _interrupt_flag.is_set():
+        return False
+    now = time.monotonic()
+    if now - _last_emergency_save_mono < 1.0:
+        return True
+    _last_emergency_save_mono = now
+    out_dir = pathlib.Path(args.save_dir)
+    out_path = out_dir / "interrupt.pt"
+    meta = {**meta_basics, "interrupt": {"sig": _interrupt_reason.get("sig"), "trace": _interrupt_reason.get("trace"), "wall_time": time.time()}}
+    try:
+        save_ckpt(out_path, core, ar_h, nat_h, sat_h, opt, scaler, meta)
+        print("🛑 emergency checkpoint written; exiting due to interrupt.")
+    except Exception as e:
+        print(f"[interrupt-save-failed] {e}")
+    return True
+
+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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_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)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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
+
+    def _atexit_note():
+        if _interrupt_flag.is_set():
+            print("[atexit] process exiting after interrupt; latest emergency checkpoint already attempted.")
+    atexit.register(_atexit_note)
+
+    while seen_tok < total_tokens_needed:
+        if _emergency_save_if_needed(
+            args,
+            meta_basics={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                         "py_state": random.getstate(), "torch_state": rng_state()},
+            core=core, ar_h=ar_h, nat_h=nat_h, sat_h=sat_h, opt=opt, scaler=scaler
+        ):
+            return
+
+        try:
+            while len(buf) < BLOCK:
+                buf.append(next(stream))
+        except StopIteration:
+            break
+        ids = torch.tensor(buf[:BLOCK], device=DEV).unsqueeze(0)
+        buf = buf[BLOCK:]
+
+        tgt_ar = ids.clone()
+        ids_nat = torch.repeat_interleave(ids, 2, 1)
+
+        try:
+            with amp(args.amp):
+                # AR
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+                # NAT
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+                # SAT
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+                loss = loss_ar + loss_nat + loss_sat
+
+            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
+
+        step += 1
+        seen_tok += BLOCK
+        pbar.update(BLOCK)
+        pbar.set_postfix(loss=f"{loss.item():.3f}", block=BLOCK)
+
+        # time-based checkpoint cadence
+        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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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()
+
+    if not _interrupt_flag.is_set():
+        save_ckpt(
+            pathlib.Path(args.save_dir) / "final.pt",
+            core, ar_h, nat_h, sat_h, opt, scaler,
+            meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                  "py_state": random.getstate(), "torch_state": rng_state()}
+        )
+        print("🎉 training complete")
+    else:
+        print("Ended after interrupt; final save skipped (emergency checkpoint already written).")
+
+# ───────────────────────── 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)
+    B, V = logits.size(0), logits.size(-1)
+    probs = logits.softmax(-1)
+    if top_k and top_k < V:
+        _, idx = torch.topk(probs, top_k, dim=-1)
+        mask = torch.full_like(probs, 0.0)
+        mask.scatter_(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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
+    h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
+    t0 = time.time()
+    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)
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{max_new} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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):
+    ids = torch.tensor([tok.encode(prompt)], device=DEV)
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(row_logits, ids, penalty_last_n, repetition_penalty, presence_penalty, frequency_penalty)
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+    print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams):
+    ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    print(tok.decode(out, skip_special_tokens=True))
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
+
+if __name__ == "__main__":
+    main()

ep2.py

@@ -0,0 +1,924 @@
+#!/usr/bin/env python3
+# ep.py — joint AR+NAT+SAT trainer/decoder (Qwen3 tokenizer)
+# Robust fresh-start, ignores *.pt.tmp, AMP dtype auto, OOM backoff, progressive block growth.
+# Added: repetition/presence/frequency penalties, top-k/top-p/min-p, greedy, no-repeat-ngrams.
+# Fixes: SAT multinomial shape; checkpoint loads on CPU; cfg fallback if ckpt missing cfg.
+# UPDATE: time-based checkpointing only (monotonic), no step-based saving. Resume respects interval.
+# NEW: Graceful shutdown: catches SIGINT/SIGTERM, writes an atomic "interrupt.pt", then exits.
+# NEW: Prompt coloring in output; default bright gray, override with --prompt_color (name or ANSI code), or 'none' to disable.
+
+from __future__ import annotations
+import argparse, json, math, pathlib, random, time, os, sys, signal, atexit, threading, traceback
+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
+from transformers import AutoTokenizer, logging as hf_log
+from tqdm.auto import tqdm
+
+# ───────────────────────── ANSI color helpers ─────────────────────────
+ANSI_COLORS = {
+    "black": "30", "red": "31", "green": "32", "yellow": "33",
+    "blue": "34", "magenta": "35", "cyan": "36", "white": "37",
+    "gray": "90", "bright_black": "90", "bright_gray": "90"
+}
+def _ansi(code: str) -> str:
+    return f"\x1b[{code}m"
+
+def _resolve_prompt_color(s: Optional[str]) -> Optional[str]:
+    if s is None:
+        return "90"  # default bright gray
+    s = s.strip().lower()
+    if s in ("none", "off", "no", "false"):
+        return None
+    return ANSI_COLORS.get(s, s)  # allow raw numeric like "31"
+
+def _print_with_prompt_color(prompt_text: str, gen_text: str, prompt_color: Optional[str]):
+    code = _resolve_prompt_color(prompt_color)
+    if code is None:
+        sys.stdout.write(prompt_text + gen_text + "\n")
+        return
+    sys.stdout.write(_ansi(code))
+    sys.stdout.write(prompt_text)
+    sys.stdout.write(_ansi("0"))  # reset
+    sys.stdout.write(gen_text + "\n")
+
+# ───────────────────────── 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
+
+# Use the Qwen3 tokenizer (can override with env TOKENIZER_ID if needed)
+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, BLANK, EOS = (
+    max(tok.get_vocab().values()) + 1,
+    tok.pad_token_id,
+    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
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_joint")
+
+# Interrupt state
+_interrupt_flag = threading.Event()
+_interrupt_reason = {"sig": None, "trace": None}
+_last_emergency_save_mono = 0.0
+
+# ───────────────────────── 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:
+    """
+    Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
+    If not usable, return 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"):
+    """
+    Always load on CPU to avoid CUDA fragmentation/OOM during torch.load.
+    """
+    try:
+        return torch.load(path, map_location=map_location)
+    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 _auto_amp_dtype():
+    if DEV.type == "cuda":
+        try:
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        except Exception:
+            return torch.float16
+    return torch.float32
+
+def amp(enabled: bool):
+    return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
+
+# ───────────────────────── Data stream ─────────────────────────
+def token_stream(ds_name: str, target: int, seed: int = 42):
+    ds = load_dataset(ds_name, split="train", streaming=True)
+    ds = ds.shuffle(buffer_size=10_000, seed=seed)
+    emitted = 0
+    for ex in ds:
+        enc = tok.encode(ex["text"])
+        if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+            enc = enc + [EOS]
+        for t in enc:
+            yield t
+            emitted += 1
+            if emitted >= target:
+                return
+
+# ───────────────────────── 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, mask=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, mask, kv=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, mask, 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)
+
+class NATHead(nn.Module):
+    def __init__(self, d): super().__init__(); self.proj = nn.Linear(d, VOCAB)
+    def forward(self, h): return self.proj(h)
+
+class SATHead(nn.Module):
+    def __init__(self, d, mode="var"):
+        super().__init__()
+        self.proj = nn.Linear(d, VOCAB)
+        self.mode = mode
+        self.gate = nn.Linear(d, 2) if mode == "var" else None
+    def forward(self, h_last):
+        logits = self.proj(h_last)
+        gate = self.gate(h_last[:, 0]) if self.gate is not None else None
+        return logits, gate
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
+    return torch.triu(m, 1)
+
+def sat_mask(n, block=SAT_BLOCK):
+    idx = torch.arange(n, device=DEV)
+    grp = idx.unsqueeze(0) // block
+    allow = (grp.T == grp) | (grp.T > grp)
+    return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(path: pathlib.Path, core: nn.Module, ar_h: nn.Module, nat_h: nn.Module, sat_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(),
+        "nat": nat_h.state_dict(),
+        "sat": sat_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 != "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, nat_h: nn.Module, sat_h: nn.Module,
+              opt: torch.optim.Optimizer, scaler: GradScaler):
+    """
+    Load a full training state from a checkpoint file or directory.
+    Returns (step, seen_tok, wall_time)
+    """
+    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
+    if "core" in ck: core.load_state_dict(ck["core"])
+    if "ar" in ck:   ar_h.load_state_dict(ck["ar"])
+    if "nat" in ck:  nat_h.load_state_dict(ck["nat"])
+    if "sat" in ck:  sat_h.load_state_dict(ck["sat"])
+    # opt/scaler can be missing if you saved partials; load best-effort
+    try:
+        if "opt" in ck: opt.load_state_dict(ck["opt"])
+    except Exception as e:
+        print(f"[resume] optimizer load skipped: {e}")
+    try:
+        if "scaler" in ck: scaler.load_state_dict(ck["scaler"])
+    except Exception as e:
+        print(f"[resume] scaler load skipped: {e}")
+    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
+
+# ───────────────────────── Interrupt handling ─────────────────────────
+def _mark_interrupt(sig_name: str):
+    if not _interrupt_flag.is_set():
+        _interrupt_reason["sig"] = sig_name
+        try:
+            _interrupt_reason["trace"] = "".join(traceback.format_stack(limit=5))
+        except Exception:
+            _interrupt_reason["trace"] = None
+        _interrupt_flag.set()
+        print(f"\n[interrupt] received {sig_name}; will save an emergency checkpoint and exit...")
+
+def _install_signal_handlers():
+    def _handler(signum, frame):
+        name = {signal.SIGINT: "SIGINT", signal.SIGTERM: "SIGTERM"}.get(signum, f"SIG{signum}")
+        _mark_interrupt(name)
+    try: signal.signal(signal.SIGINT, _handler)
+    except Exception: pass
+    try: signal.signal(signal.SIGTERM, _handler)
+    except Exception: pass
+
+_install_signal_handlers()
+
+# ───────────────────────── 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 _emergency_save_if_needed(args, meta_basics, core, ar_h, nat_h, sat_h, opt, scaler):
+    global _last_emergency_save_mono
+    if not _interrupt_flag.is_set():
+        return False
+    now = time.monotonic()
+    if now - _last_emergency_save_mono < 1.0:
+        return True
+    _last_emergency_save_mono = now
+    out_dir = pathlib.Path(args.save_dir)
+    out_path = out_dir / "interrupt.pt"
+    meta = {**meta_basics, "interrupt": {"sig": _interrupt_reason.get("sig"), "trace": _interrupt_reason.get("trace"), "wall_time": time.time()}}
+    try:
+        save_ckpt(out_path, core, ar_h, nat_h, sat_h, opt, scaler, meta)
+        print("🛑 emergency checkpoint written; exiting due to interrupt.")
+    except Exception as e:
+        print(f"[interrupt-save-failed] {e}")
+    return True
+
+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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"], mode="var").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")
+            loaded += _safe_load_any(src, nat_h, key="nat")
+            loaded += _safe_load_any(src, sat_h, key="sat")
+            if loaded:
+                print(f"Warm-start: loaded {loaded} matching tensors from {src}")
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": core.parameters(), "lr": LR_CORE},
+            {"params": ar_h.parameters(), "lr": LR_HEAD},
+            {"params": nat_h.parameters(), "lr": LR_HEAD},
+            {"params": sat_h.parameters(), "lr": LR_HEAD},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
+    ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
+    ce_gate = nn.CrossEntropyLoss()
+
+    # ---------- 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, nat_h, sat_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)
+
+    # Target tokens
+    if args.target_tokens:
+        target_tokens = args.target_tokens
+    else:
+        param_count = sum(p.numel() for p in core.parameters())
+        target_tokens = int(25 * 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")
+
+    stream = token_stream(args.source, 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
+
+    def _atexit_note():
+        if _interrupt_flag.is_set():
+            print("[atexit] process exiting after interrupt; latest emergency checkpoint already attempted.")
+    atexit.register(_atexit_note)
+
+    while seen_tok < total_tokens_needed:
+        if _emergency_save_if_needed(
+            args,
+            meta_basics={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                         "py_state": random.getstate(), "torch_state": rng_state()},
+            core=core, ar_h=ar_h, nat_h=nat_h, sat_h=sat_h, opt=opt, scaler=scaler
+        ):
+            return
+
+        try:
+            while len(buf) < BLOCK:
+                buf.append(next(stream))
+        except StopIteration:
+            break
+        ids = torch.tensor(buf[:BLOCK], device=DEV).unsqueeze(0)
+        buf = buf[BLOCK:]
+
+        tgt_ar = ids.clone()
+        ids_nat = torch.repeat_interleave(ids, 2, 1)
+
+        try:
+            with amp(args.amp):
+                # AR
+                h_ar = core(ids, causal_mask(ids.size(1)))
+                logits_ar = ar_h(h_ar)[:, :-1]
+                loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
+                # NAT
+                h_nat = core(ids_nat, None)
+                log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)
+                ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
+                loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
+                # SAT
+                h_sat = core(ids, sat_mask(ids.size(1)))
+                logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
+                tgt_sat = ids[:, 1:SAT_BLOCK+1]
+                loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
+                if gate is not None:
+                    loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+                loss = loss_ar + loss_nat + loss_sat
+
+            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
+
+        step += 1
+        seen_tok += BLOCK
+        pbar.update(BLOCK)
+        pbar.set_postfix(loss=f"{loss.item():.3f}", block=BLOCK)
+
+        # time-based checkpoint cadence
+        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, nat_h, sat_h, opt, scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "py_state": random.getstate(),
+                        "torch_state": rng_state(),
+                    },
+                )
+                last_save_mono = now_mono
+                last_save_wall = time.time()
+
+        # 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()
+
+    if not _interrupt_flag.is_set():
+        save_ckpt(
+            pathlib.Path(args.save_dir) / "final.pt",
+            core, ar_h, nat_h, sat_h, opt, scaler,
+            meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(),
+                  "py_state": random.getstate(), "torch_state": rng_state()}
+        )
+        print("🎉 training complete")
+    else:
+        print("Ended after interrupt; final save skipped (emergency checkpoint already written).")
+
+# ───────────────────────── 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)
+    B, V = logits.size(0), logits.size(-1)
+    probs = logits.softmax(-1)
+    if top_k and top_k < V:
+        _, idx = torch.topk(probs, top_k, dim=-1)
+        mask = torch.full_like(probs, 0.0)
+        mask.scatter_(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, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
+    sat_h = SATHead(cfg["d"]).to(DEV)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    nat_h.load_state_dict(sd["nat"])
+    sat_h.load_state_dict(sd["sat"])
+    return core, ar_h, nat_h, sat_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,
+              prompt_color: Optional[str] = "90"):
+    # tokenize and remember prompt 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()
+    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)
+
+    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)
+    _print_with_prompt_color(prompt_text, gen_text, prompt_color)
+    print(f"[{len(full_ids) - prompt_len} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def sat_decode(core, sat_h, prompt, max_new, T, var,
+               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,
+               prompt_color: Optional[str] = "90"):
+    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)
+
+    added, t0 = 0, time.time()
+    while added < max_new:
+        h = core(ids, sat_mask(ids.size(1)))
+        logits_all, gate = sat_h(h[:, -SAT_BLOCK:])
+        stride = 2 if (not var or gate is None) else (gate.softmax(-1).multinomial(1) + 1).item()
+        stride = int(stride)
+        for pos in range(stride):
+            row_logits = logits_all[:, pos, :]
+            row_logits = _apply_no_repeat_ngram(row_logits, ids, no_repeat_ngram_size)
+            row_logits = _apply_rep_presence_frequency(row_logits, ids, penalty_last_n,
+                                                       repetition_penalty, presence_penalty, frequency_penalty)
+            if greedy:
+                nxt = row_logits.argmax(-1, keepdim=True)
+            else:
+                probs = _filter_top_k_top_p_min_p(row_logits.squeeze(0), top_k, top_p, min_p, T)
+                nxt = probs.multinomial(1)
+            ids = torch.cat([ids, nxt], 1)
+            added += 1
+            if added >= max_new:
+                break
+
+    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)
+    _print_with_prompt_color(prompt_text, gen_text, prompt_color)
+    print(f"[{added} tok in {time.time() - t0:.2f}s]")
+
+@torch.no_grad()
+def nat_decode(core, nat_h, prompt, max_new, passes, streams,
+               prompt_color: Optional[str] = "90"):
+    prompt_ids = tok.encode(prompt)
+    ids = torch.tensor([prompt_ids + [BLANK] * (max_new * 2)], device=DEV)
+    t0 = time.time()
+    for _ in range(passes):
+        h = core(ids, None)
+        logits = nat_h(h)
+        logits[..., BLANK] = -1e9
+        cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
+        best = (cand != BLANK).float().mean(-1).argmax(0)
+        ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
+    out = [t for t in ids[0].tolist() if t != BLANK]
+    gen_text = tok.decode(out, skip_special_tokens=True)
+    prompt_text = tok.decode(prompt_ids, skip_special_tokens=True)
+    _print_with_prompt_color(prompt_text, gen_text, prompt_color)
+    print(f"[{len(out)} output tokens 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")
+    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_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")
+
+    # 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")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--mode", choices=["ar", "nat", "sat"], 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)
+
+    # New 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)
+
+    inf.add_argument("--var", action="store_true")
+    inf.add_argument("--passes", type=int, default=1)
+    inf.add_argument("--streams", type=int, default=5)
+
+    # Prompt color flag (name or raw ANSI code; use 'none' to disable)
+    inf.add_argument("--prompt_color", type=str, default="90",
+                     help="ANSI color name/code for the prompt (e.g., gray, cyan, 90). Use 'none' to disable.")
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    else:
+        core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
+        if args.mode == "ar":
+            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,
+                      prompt_color=args.prompt_color)
+        elif args.mode == "sat":
+            sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
+                       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,
+                       prompt_color=args.prompt_color)
+        else:
+            nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams,
+                       prompt_color=args.prompt_color)
+
+if __name__ == "__main__":
+    main()

step08250364.pt

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