Add n_tenstorrent_port.py — Tenstorrent N300s training port

n_tenstorrent_port.py

@@ -0,0 +1,1754 @@
+#!/usr/bin/env python3
+"""
+n_tenstorrent_port.py
+
+Training-first port of the user's joint AR+SAT trainer to support:
+- Tenstorrent via TT-XLA / PJRT (`--backend tt`)
+- NVIDIA CUDA (`--backend cuda`)
+- CPU fallback (`--backend cpu`)
+
+Design goals:
+- Keep checkpoint format PyTorch-native and cross-device loadable.
+- Prioritize stable training on TT over aggressive graph tricks.
+- Preserve NVIDIA-trained checkpoint compatibility for inference.
+- Stay as close as practical to the original single-file workflow.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import pathlib
+import time
+from contextlib import nullcontext
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from datasets import DownloadConfig, load_dataset
+from transformers import AutoTokenizer, logging as hf_log
+
+STATUS_FILE = "/workspace/status.json"
+
+
+# ───────────────────────── Status helpers ─────────────────────────
+def write_status(step, seen_tok, loss, batch, block, tok_per_sec, phase):
+    try:
+        with open(STATUS_FILE, "w") as f:
+            json.dump(
+                {
+                    "step": step,
+                    "seen_tok": seen_tok,
+                    "loss": float(loss) if loss is not None else None,
+                    "batch": batch,
+                    "block": block,
+                    "tok_per_sec": tok_per_sec,
+                    "phase": phase,
+                    "updated": time.time(),
+                    "target_tok": 35737600000,
+                },
+                f,
+            )
+    except Exception:
+        pass
+
+
+def show_status():
+    try:
+        with open(STATUS_FILE) as f:
+            s = json.load(f)
+        age = time.time() - s.get("updated", 0)
+        target = s.get("target_tok") or 35737600000
+        remaining = target - s.get("seen_tok", 0)
+        eta_sec = remaining / max(s.get("tok_per_sec", 1), 1)
+        eta_days = eta_sec / 86400
+        print(
+            f"Step: {s.get('step', '?'):,} | Tokens: {s.get('seen_tok', 0)/1e9:.2f}B / {target/1e9:.1f}B | Loss: {s.get('loss', 0):.4f}"
+        )
+        print(
+            f"Speed: {s.get('tok_per_sec', 0):.0f} tok/s | B={s.get('batch')} L={s.get('block')} | ETA: {eta_days:.1f} days | {age:.0f}s ago"
+        )
+    except FileNotFoundError:
+        print("No status file. Training not running?")
+    except Exception as e:
+        print(f"Error: {e}")
+
+
+# ───────────────────────── Safe progress ─────────────────────────
+class SafeProgress:
+    def __init__(self, total, initial=0, unit="tok"):
+        self.total = total
+        self.n = initial
+        self.unit = unit
+        self.last_print = initial
+        self.postfix = {}
+        self.start_time = time.time()
+
+    def update(self, n=1):
+        self.n += n
+        if self.n - self.last_print >= 1_000_000:
+            self._print()
+            self.last_print = self.n
+
+    def set_postfix(self, **kwargs):
+        self.postfix = kwargs
+
+    def _print(self):
+        elapsed = time.time() - self.start_time
+        rate = self.n / elapsed if elapsed > 0 else 0
+        pct = 100 * self.n / self.total if self.total > 0 else 0
+        pf = " ".join(f"{k}={v}" for k, v in self.postfix.items())
+        print(f"[{pct:.1f}%] {self.n:,}/{self.total:,} {self.unit} | {rate:.0f} tok/s | {pf}")
+
+    def close(self):
+        self._print()
+        print("Done.")
+
+
+# ───────────────────────── ANSI colors ─────────────────────────
+class Colors:
+    RESET = "\033[0m"
+    BOLD = "\033[1m"
+    PROMPT = "\033[36m"
+    GEN = "\033[0m"
+    INFO = "\033[90m"
+    WARN = "\033[93m"
+
+
+hf_log.set_verbosity_error()
+
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+try:
+    torch.set_float32_matmul_precision("high")
+except Exception:
+    pass
+
+
+# ───────────────────────── Runtime backend ─────────────────────────
+@dataclass
+class BackendRuntime:
+    backend: str
+    device: torch.device
+    is_cuda: bool = False
+    is_tt: bool = False
+    is_xla: bool = False
+    dtype: torch.dtype = torch.float32
+    xm: Any = None
+    xr: Any = None
+    xs: Any = None
+    mesh: Any = None
+    spmd: bool = False
+    compile_options: Optional[Dict[str, str]] = None
+    num_devices: int = 1
+
+    def sync(self, wait: bool = False) -> None:
+        if self.is_cuda:
+            torch.cuda.synchronize(self.device)
+            return
+        if self.is_tt:
+            try:
+                import torch_xla
+
+                torch_xla.sync(wait=wait)
+                return
+            except Exception:
+                pass
+            if self.xm is not None:
+                try:
+                    self.xm.mark_step()
+                except Exception:
+                    pass
+
+    def optimizer_step(self, optimizer: torch.optim.Optimizer) -> None:
+        if self.is_tt and self.xm is not None:
+            try:
+                self.xm.optimizer_step(optimizer, barrier=True)
+            except TypeError:
+                self.xm.optimizer_step(optimizer)
+        else:
+            optimizer.step()
+
+    def maybe_mark_batch_sharding(self, *tensors: torch.Tensor) -> None:
+        if not (self.is_tt and self.spmd and self.xs is not None and self.mesh is not None):
+            return
+        for tensor in tensors:
+            if tensor is None:
+                continue
+            try:
+                if tensor.ndim == 1:
+                    self.xs.mark_sharding(tensor, self.mesh, ("batch",))
+                elif tensor.ndim >= 2:
+                    spec = ["batch"] + [None] * (tensor.ndim - 1)
+                    self.xs.mark_sharding(tensor, self.mesh, tuple(spec))
+            except Exception:
+                # Sharding is best-effort and still fairly sharp-edged.
+                pass
+
+
+RUNTIME = BackendRuntime(backend="cpu", device=torch.device("cpu"))
+DEV = RUNTIME.device
+
+
+def setup_runtime(args) -> BackendRuntime:
+    global RUNTIME, DEV
+
+    if getattr(args, "backend", "auto") == "tt" and (
+        getattr(args, "tt_bfp8", False) or getattr(args, "tt_weight_bfp8", False)
+    ) and getattr(args, "tt_dtype", "bf16") != "bf16":
+        print("[tt-xla] forcing --tt_dtype bf16 because bfp8 conversion requires a bf16 model input dtype")
+        args.tt_dtype = "bf16"
+
+    requested = getattr(args, "backend", "auto")
+    if requested == "auto":
+        if os.environ.get("PJRT_DEVICE", "").upper() == "TT":
+            requested = "tt"
+        elif torch.cuda.is_available():
+            requested = "cuda"
+        else:
+            requested = "cpu"
+
+    if requested == "cuda":
+        runtime = BackendRuntime(
+            backend="cuda",
+            device=torch.device("cuda"),
+            is_cuda=True,
+            dtype=torch.float32,
+        )
+        RUNTIME = runtime
+        DEV = runtime.device
+        return runtime
+
+    if requested == "tt":
+        os.environ.setdefault("PJRT_DEVICE", "TT")
+        os.environ.setdefault("XLA_STABLEHLO_COMPILE", "1")
+        if getattr(args, "tt_spmd", False):
+            os.environ.setdefault("XLA_ALWAYS_ALLREDUCE", "1")
+            os.environ.setdefault("CONVERT_SHLO_TO_SHARDY", "1")
+        if getattr(args, "tt_trace", False):
+            os.environ.setdefault(
+                "TT_RUNTIME_TRACE_REGION_SIZE",
+                str(getattr(args, "tt_trace_region_size", 10_000_000)),
+            )
+
+        import numpy as np  # local import to avoid dependency unless needed
+        import torch_xla
+        import torch_xla.core.xla_model as xm
+        import torch_xla.runtime as xr
+
+        xr.set_device_type("TT")
+        compile_options = {
+            "optimization_level": str(getattr(args, "tt_optimization_level", 1)),
+        }
+        if getattr(args, "tt_bfp8", False):
+            compile_options["enable_bfp8_conversion"] = "true"
+        if getattr(args, "tt_weight_bfp8", False):
+            compile_options["experimental_enable_weight_bfp8_conversion"] = "true"
+        if getattr(args, "tt_trace", False):
+            compile_options["enable_trace"] = "true"
+        torch_xla.set_custom_compile_options(compile_options)
+
+        xs = None
+        mesh = None
+        num_devices = 1
+        if getattr(args, "tt_spmd", False):
+            try:
+                import torch_xla.distributed.spmd as xs
+                from torch_xla.distributed.spmd import Mesh
+
+                xr.use_spmd()
+                num_devices = xr.global_runtime_device_count()
+                mesh = Mesh(
+                    device_ids=np.arange(num_devices),
+                    mesh_shape=(1, num_devices),
+                    axis_names=("batch", "model"),
+                )
+            except Exception as e:
+                print(f"[tt-spmd] disabled due to setup failure: {e}")
+                xs = None
+                mesh = None
+                num_devices = 1
+
+        runtime = BackendRuntime(
+            backend="tt",
+            device=xm.xla_device(),
+            is_tt=True,
+            is_xla=True,
+            dtype=torch.bfloat16 if getattr(args, "tt_dtype", "bf16") == "bf16" else torch.float32,
+            xm=xm,
+            xr=xr,
+            xs=xs,
+            mesh=mesh,
+            spmd=bool(mesh is not None),
+            compile_options=compile_options,
+            num_devices=num_devices,
+        )
+        RUNTIME = runtime
+        DEV = runtime.device
+        return runtime
+
+    runtime = BackendRuntime(backend="cpu", device=torch.device("cpu"), dtype=torch.float32)
+    RUNTIME = runtime
+    DEV = runtime.device
+    return runtime
+
+
+# ───────────────────────── Tokenizer / vocab ─────────────────────────
+TOKENIZER_ID = os.environ.get("TOKENIZER_ID", "deepseek-ai/DeepSeek-V3.2")
+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
+EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id
+PAD_ID = tok.pad_token_id if tok.pad_token_id is not None else (EOS if EOS is not None else 0)
+
+
+# ───────────────────────── Presets / defaults ─────────────────────────
+PRESETS: Dict[str, Dict[str, int]] = {
+    "femto_1x": dict(d=16, layers=1, heads=1, rank=16),
+    "femto_12x": dict(d=16, layers=1, heads=1, rank=192),
+    "femto_24x": dict(d=16, layers=1, heads=1, rank=384),
+    "pico_1x": dict(d=32, layers=1, heads=2, rank=16),
+    "pico_3x": dict(d=32, layers=1, heads=2, rank=48),
+    "pico_6x": dict(d=32, layers=1, heads=2, rank=96),
+    "pico_12x": dict(d=32, layers=1, heads=2, rank=192),
+    "pico_24x": dict(d=32, layers=1, heads=2, rank=384),
+    "pico_48x": dict(d=32, layers=1, heads=2, rank=768),
+    "nano_1x": dict(d=64, layers=2, heads=4, rank=16),
+    "nano_3x": dict(d=64, layers=2, heads=4, rank=48),
+    "nano_6x": dict(d=64, layers=2, heads=4, rank=96),
+    "nano_12x": dict(d=64, layers=2, heads=4, rank=192),
+    "nano_24x": dict(d=64, layers=2, heads=4, rank=384),
+    "nano_48x": dict(d=64, layers=2, heads=4, rank=768),
+    "nano_96x": dict(d=64, layers=2, heads=4, rank=1536),
+    "micro_3x": dict(d=128, layers=4, heads=8, rank=48),
+    "micro_6x": dict(d=128, layers=4, heads=8, rank=96),
+    "micro_12x": dict(d=128, layers=4, heads=8, rank=192),
+    "micro_24x": dict(d=128, layers=4, heads=8, rank=384),
+    "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),
+    "base18": dict(d=768, layers=18, heads=24, rank=96),
+    "large": dict(d=1024, layers=24, heads=16, rank=128),
+}
+
+DEFAULT_BLOCK = 1122
+DEFAULT_BATCH = 1
+SAT_BLOCK = 2
+LR_CORE, LR_HEAD = 5e-5, 2e-4
+EMIT_LAMBDA = 0.1
+DEFAULT_SAVE_SEC = 24 * 3600
+CKDIR = pathlib.Path("ckpts_expansion")
+
+DEFAULT_PRETRAIN_SOURCES = (
+    "OpenTransformer/goddess-crawl,OpenTransformer/agillm-crawl-data,"
+    "OpenTransformer/web-crawl-2026,OpenTransformer/web-crawl-clean-v2,"
+    "OpenTransformer/scraped-web-data,OpenTransformer/turbo-crawl,"
+    "OpenTransformer/sft-data-clean,OpenTransformer/web-crawl-v1"
+)
+DEFAULT_AFTER_SFT_SOURCES = "mlabonne/opc-sft-stage2-chat,HuggingFaceH4/ultrachat_200k"
+DEFAULT_AFTER_SFT_BLOCK = 1122
+
+
+# ───────────────────────── Utilities ─────────────────────────
+def get_uk_time() -> str:
+    utc_now = datetime.now(timezone.utc)
+    year = utc_now.year
+    march_last = datetime(year, 3, 31, 1, 0, tzinfo=timezone.utc)
+    while march_last.weekday() != 6:
+        march_last = march_last.replace(day=march_last.day - 1)
+    oct_last = datetime(year, 10, 31, 1, 0, tzinfo=timezone.utc)
+    while oct_last.weekday() != 6:
+        oct_last = oct_last.replace(day=oct_last.day - 1)
+    if march_last <= utc_now < oct_last:
+        uk_offset = 1
+        tz_name = "BST"
+    else:
+        uk_offset = 0
+        tz_name = "GMT"
+    uk_time = utc_now + timedelta(hours=uk_offset)
+    return uk_time.strftime(f"%Y-%m-%d %H:%M:%S {tz_name}")
+
+
+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) -> Optional[pathlib.Path]:
+    try:
+        if path.is_dir():
+            cands = sorted(
+                [p for p in path.glob("*.pt") if _is_probably_ckpt(p)],
+                key=lambda p: p.stat().st_mtime,
+                reverse=True,
+            )
+            return cands[0] if cands else None
+        if path.suffix == ".tmp":
+            solid = path.with_suffix("")
+            return solid if _is_probably_ckpt(solid) else _resolve_ckpt(path.parent)
+        return path if _is_probably_ckpt(path) else _resolve_ckpt(path.parent)
+    except Exception:
+        return None
+
+
+def _try_load(path: pathlib.Path, map_location="cpu"):
+    try:
+        return torch.load(path, map_location=map_location)
+    except Exception as e:
+        print(f"[ckpt-skip] {path} not usable: {e}")
+        return None
+
+
+def _strip_compiled_prefix(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
+
+
+def _tree_to_cpu(obj: Any) -> Any:
+    if torch.is_tensor(obj):
+        return obj.detach().cpu()
+    if isinstance(obj, dict):
+        return {k: _tree_to_cpu(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_tree_to_cpu(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_tree_to_cpu(v) for v in obj)
+    return obj
+
+
+def optimizer_to(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
+    for state in optimizer.state.values():
+        if not isinstance(state, dict):
+            continue
+        for k, v in list(state.items()):
+            if torch.is_tensor(v):
+                state[k] = v.to(device)
+
+
+def _prune_checkpoints(save_dir: pathlib.Path, phase_name: str, max_ckpts: Optional[int]):
+    if max_ckpts is None or max_ckpts <= 0:
+        return
+    try:
+        for tmp in save_dir.glob("*.pt.tmp"):
+            try:
+                tmp.unlink()
+                print(f"  [prune] cleaned stale tmp {tmp.name}")
+            except Exception:
+                pass
+        pattern = f"{phase_name}_step*.pt"
+        ckpts = sorted(
+            [p for p in save_dir.glob(pattern) if _is_probably_ckpt(p)],
+            key=lambda p: p.stat().st_mtime,
+        )
+        excess = len(ckpts) - max_ckpts
+        if excess > 0:
+            for p in ckpts[:excess]:
+                try:
+                    p.unlink()
+                    print(f"  [prune] deleted old {p.name}")
+                except Exception:
+                    pass
+    except Exception as e:
+        print(f"[ckpt-prune] error: {e}")
+
+
+def print_expansion_info(cfg: dict, tie_weights: bool = False):
+    d_k = cfg["d"] // cfg["heads"]
+    rank = cfg["rank"]
+    ratio = rank / d_k
+    regime = "COMPRESSION" if ratio < 1 else ("IDENTITY" if ratio == 1 else "EXPANSION")
+    tie_str = "YES" if tie_weights else "NO"
+    print("┌─────────────────────────────────────────┐")
+    print("│ TUNEABLE ATTENTION CONFIG               │")
+    print("├─────────────────────────────────────────┤")
+    print(f"│ d_model: {cfg['d']:4d}  heads: {cfg['heads']:2d}  d_k: {d_k:3d}     │")
+    print(f"│ layers: {cfg['layers']:4d}  tie_weights: {tie_str:3s}          │")
+    print(f"│ rank: {rank:4d}  ratio: {ratio:.1f}x  [{regime:11s}] │")
+    print("└─────────────────────────────────────────┘")
+
+
+def _parse_grow_plan(s: str) -> List[int]:
+    return sorted(set(int(x.strip()) for x in s.split(",") if x.strip() and int(x.strip()) >= 128))
+
+
+def _count_enabled_params(*modules) -> int:
+    seen_data_ptrs = set()
+    total = 0
+    for m in modules:
+        if m is None:
+            continue
+        for p in m.parameters():
+            if p.data_ptr() not in seen_data_ptrs:
+                seen_data_ptrs.add(p.data_ptr())
+                total += p.numel()
+    return total
+
+
+def _phase_freeze(core: nn.Module, *, freeze_core: bool, unfreeze_ln: bool, train_emb: bool):
+    for p in core.parameters():
+        p.requires_grad = not freeze_core
+    if freeze_core:
+        if unfreeze_ln:
+            for blk in core.blocks:
+                for p in blk.ln1.parameters():
+                    p.requires_grad = True
+                for p in blk.ln2.parameters():
+                    p.requires_grad = True
+            for p in core.ln.parameters():
+                p.requires_grad = True
+        if train_emb:
+            for p in core.emb.parameters():
+                p.requires_grad = True
+
+
+def retie_weights(core: nn.Module, ar_h: nn.Module, tie_weights: bool) -> None:
+    if tie_weights:
+        ar_h.proj.weight = core.emb.weight
+
+
+# ───────────────────────── AMP helper ─────────────────────────
+try:
+    from torch.amp import GradScaler, autocast as _ac
+except ImportError:
+    from torch.cuda.amp import GradScaler, autocast as _ac
+
+
+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):
+    if not (enabled and DEV.type == "cuda"):
+        return nullcontext()
+    try:
+        return _ac(device_type="cuda", dtype=_auto_amp_dtype())
+    except TypeError:
+        return _ac(dtype=_auto_amp_dtype())
+
+
+# ───────────────────────── Chat & data stream ─────────────────────────
+def _coerce_role(r: str) -> str:
+    r = (r or "").lower()
+    if r in {"user", "human", "customer"}:
+        return "user"
+    if r in {"assistant", "gpt", "bot"}:
+        return "assistant"
+    if r in {"system", "context"}:
+        return "system"
+    return r or "user"
+
+
+def _render_chat_text_from_ex(ex: dict, messages_key: str, add_generation_prompt: bool) -> Optional[str]:
+    msgs = ex.get(messages_key)
+    if msgs is None:
+        for alt in ("conversations", "dialog", "turns"):
+            if isinstance(ex.get(alt), list):
+                msgs = ex[alt]
+                break
+    if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict):
+        try:
+            norm = []
+            for m in msgs:
+                role = _coerce_role(m.get("role", ""))
+                content = m.get("content", m.get("text", ""))
+                if not isinstance(content, str):
+                    continue
+                norm.append({"role": role, "content": content})
+            if not norm:
+                return None
+            return tok.apply_chat_template(norm, tokenize=False, add_generation_prompt=add_generation_prompt)
+        except Exception:
+            return None
+    for a, b in (("prompt", "response"), ("instruction", "output"), ("question", "answer")):
+        if isinstance(ex.get(a), str) and isinstance(ex.get(b), str):
+            return f"User: {ex[a]}\nAssistant: {ex[b]}"
+    return None
+
+
+def _open_stream_one(ds_name: str, seed: int, streaming: bool = True):
+    dc = DownloadConfig(max_retries=5, use_etag=True, resume_download=True)
+    if ":" in ds_name:
+        base, config = ds_name.split(":", 1)
+    else:
+        base, config = ds_name, None
+    if not streaming:
+        print(f"[download] Downloading {ds_name} (non-streaming)...")
+    if base == "json":
+        data_files = {"train": config}
+        ds = load_dataset("json", data_files=data_files, split="train", streaming=streaming, download_config=dc)
+    else:
+        ds = (
+            load_dataset(base, config, split="train", streaming=streaming, download_config=dc)
+            if config
+            else load_dataset(base, split="train", streaming=streaming, download_config=dc)
+        )
+    if streaming:
+        return iter(ds.shuffle(buffer_size=1000, seed=seed))
+    print(f"[download] Got {len(ds):,} examples. Shuffling...")
+    ds = ds.shuffle(seed=seed)
+    return iter(ds)
+
+
+_HOT_CFG_PATH = pathlib.Path("/workspace/hot_config.json")
+_hot_cache = {"mtime": 0, "data": {}}
+
+
+def get_hot_datasets(default):
+    try:
+        if _HOT_CFG_PATH.exists():
+            mt = _HOT_CFG_PATH.stat().st_mtime
+            if mt > _hot_cache["mtime"]:
+                _hot_cache["data"] = json.loads(_HOT_CFG_PATH.read_text())
+                _hot_cache["mtime"] = mt
+            cfg = _hot_cache["data"]
+            if "datasets" in cfg:
+                ds = cfg["datasets"]
+                if isinstance(ds, list):
+                    ds = ",".join(ds)
+                print(f"[HOT] Using: {ds[:60]}...")
+                return ds
+    except Exception as e:
+        print(f"[HOT] Error: {e}")
+    return default
+
+
+def token_stream(
+    ds_names: str,
+    target: int,
+    seed: int = 42,
+    chat: bool = False,
+    chat_messages_key: str = "messages",
+    sft_add_generation_prompt: bool = False,
+    dataset_field_text: str = "text",
+    streaming: bool = True,
+):
+    ds_names = get_hot_datasets(ds_names)
+    sources = [s.strip() for s in ds_names.split(",") if s.strip()]
+    if not sources:
+        return
+    src_idx = 0
+    emitted = 0
+    it = None
+    attempts = 0
+    backoff_base = 2.0
+    while emitted < target:
+        try:
+            if it is None:
+                it = _open_stream_one(sources[src_idx], seed, streaming=streaming)
+            ex = next(it)
+            text = None
+            if isinstance(ex, dict):
+                if chat:
+                    text = _render_chat_text_from_ex(ex, chat_messages_key, sft_add_generation_prompt)
+                if text is None:
+                    if dataset_field_text and isinstance(ex.get(dataset_field_text), str):
+                        text = ex[dataset_field_text]
+                    elif isinstance(ex.get("text"), str):
+                        text = ex["text"]
+            if not isinstance(text, str):
+                attempts = 0
+                continue
+            enc = tok.encode(text)
+            if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
+                enc = enc + [EOS]
+            for t in enc:
+                yield t
+                emitted += 1
+                if emitted >= target:
+                    return
+            attempts = 0
+        except StopIteration:
+            it = None
+            src_idx = (src_idx + 1) % len(sources)
+        except Exception as e:
+            attempts += 1
+            sleep_s = min(60.0, backoff_base ** min(attempts, 6))
+            print(f"[stream-retry] {sources[src_idx]} error: {type(e).__name__}, sleeping {sleep_s:.1f}s")
+            time.sleep(sleep_s)
+            it = None
+            if attempts % 2 == 0 and len(sources) > 1:
+                src_idx = (src_idx + 1) % len(sources)
+
+
+# ───────────────────────── ALiBi ─────────────────────────
+@torch._dynamo.disable
+def _alibi_slopes(n_heads: int):
+    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)
+
+
+@torch._dynamo.disable
+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)
+    return -_alibi_slopes(n_heads) * dist
+
+
+# ───────────────────────── Model components ─────────────────────────
+class TuneableAttentionMHA(nn.Module):
+    def __init__(self, d: int, h: int, r: int, use_relpos: bool = True):
+        super().__init__()
+        assert d % h == 0
+        self.h, self.dk, self.r = h, d // h, r
+        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 * self.dk, d, bias=False)
+        self.drop = nn.Dropout(0.1)
+
+    def _proj_qk(self, x):
+        B, N, _ = x.shape
+        return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U)
+
+    def _reshape_v(self, x):
+        B, N, _ = x.shape
+        return x.view(B, N, self.h, self.dk).transpose(1, 2)
+
+    def forward(self, x, mask=None, rel_bias_tokens=None, kv_cache=None, use_cache=False):
+        q = self._proj_qk(self.q(x))
+        k_new = self._proj_qk(self.k(x))
+        v_new = self._reshape_v(self.v(x))
+        if kv_cache is None:
+            k, v = k_new, v_new
+        else:
+            k_cached, v_cached = kv_cache
+            if use_cache:
+                k = torch.cat([k_cached, k_new], dim=2)
+                v = torch.cat([v_cached, v_new], dim=2)
+            else:
+                k, v = k_new, v_new
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)
+        if self.use_relpos and rel_bias_tokens is not None:
+            att = att + alibi_bias(self.h, rel_bias_tokens).to(att.dtype)[:, :, -q.size(2) :, :]
+        if mask is not None:
+            att = att + mask.to(att.dtype)
+        z = (att.softmax(-1) @ v).transpose(1, 2).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 = TuneableAttentionMHA(d, h, r)
+        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=False, total_seq_len=None):
+        if use_cache:
+            y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=total_seq_len, kv_cache=kv, use_cache=True)
+            x = x + y + self.ff(self.ln2(x + y))
+            return x, new_kv
+        n = x.size(1)
+        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, tie_weights: bool = False):
+        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)
+        self.tie_weights = tie_weights
+
+    def forward(self, ids, mask, kv_caches=None, use_cache=False, total_seq_len=None):
+        x = self.emb(ids)
+        if not use_cache:
+            for blk in self.blocks:
+                x = blk(x, mask)
+            return self.ln(x)
+        new_kvs = []
+        for i, blk in enumerate(self.blocks):
+            kv = kv_caches[i] if kv_caches else None
+            x, kv_out = blk(x, mask, kv, use_cache=True, total_seq_len=total_seq_len)
+            new_kvs.append(kv_out)
+        return self.ln(x), new_kvs
+
+
+class ARHead(nn.Module):
+    def __init__(self, d, tie_weights: bool = False, embedding_weight: nn.Parameter = None):
+        super().__init__()
+        self.tie_weights = tie_weights
+        if tie_weights and embedding_weight is not None:
+            self.proj = nn.Linear(d, VOCAB, bias=False)
+            self.proj.weight = embedding_weight
+        else:
+            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.gate = nn.Linear(d, 2) if mode == "var" else None
+
+    def forward(self, h_last):
+        return self.proj(h_last), (self.gate(h_last[:, 0]) if self.gate else None)
+
+
+# ───────────────────────── Masks ─────────────────────────
+def causal_mask(n):
+    return torch.triu(torch.full((1, 1, n, n), float("-inf"), device=DEV), 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)
+
+
+def sat_mask_cached(new_len: int, cached_len: int, block=SAT_BLOCK):
+    total_len = cached_len + new_len
+    return torch.zeros((1, 1, new_len, total_len), device=DEV)
+
+
+def causal_padded_mask(total_len: int, valid_len: int):
+    mask = causal_mask(total_len)
+    if valid_len < total_len:
+        mask[:, :, :, valid_len:] = float("-inf")
+        mask[:, :, valid_len:, :] = float("-inf")
+    return mask
+
+
+def sat_padded_mask(total_len: int, valid_len: int):
+    mask = sat_mask(total_len)
+    if valid_len < total_len:
+        mask[:, :, :, valid_len:] = float("-inf")
+        mask[:, :, valid_len:, :] = float("-inf")
+    return mask
+
+
+# ───────────────────────── Checkpoint helpers ─────────────────────────
+def save_ckpt(path: pathlib.Path, core, ar_h, sat_h, opt, scaler, meta):
+    if RUNTIME.is_tt:
+        RUNTIME.sync(wait=True)
+    path.parent.mkdir(exist_ok=True, parents=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    state = {
+        "core": _tree_to_cpu(_strip_compiled_prefix(core.state_dict())),
+        "ar": _tree_to_cpu(_strip_compiled_prefix(ar_h.state_dict())),
+        "sat": _tree_to_cpu(_strip_compiled_prefix(sat_h.state_dict())),
+        "opt": _tree_to_cpu(opt.state_dict()),
+        "scaler": _tree_to_cpu(scaler.state_dict()),
+        "cfg": meta.get("cfg"),
+        "tokenizer_id": TOKENIZER_ID,
+        "tie_weights": meta.get("tie_weights", False),
+        **{k: v for k, v in meta.items() if k not in ("cfg", "tie_weights")},
+    }
+    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"],
+                "block_size": meta.get("block_size"),
+                "batch_size": meta.get("batch_size"),
+                "seen_tok": meta.get("seen_tok"),
+            }
+        )
+    )
+    print(f"\n✓ saved checkpoint {path.name}")
+
+
+
+def load_ckpt(path, core, ar_h, sat_h, opt, scaler):
+    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(_strip_compiled_prefix(ck["core"]))
+    ar_h.load_state_dict(_strip_compiled_prefix(ck["ar"]))
+    sat_h.load_state_dict(_strip_compiled_prefix(ck["sat"]))
+    try:
+        opt.load_state_dict(ck["opt"])
+        optimizer_to(opt, DEV)
+    except Exception as e:
+        print(f"[resume] optimizer state skipped: {e}")
+    if ck.get("scaler"):
+        try:
+            scaler.load_state_dict(ck["scaler"])
+        except Exception:
+            pass
+    return ck.get("step", 0), ck.get("seen_tok", 0), ck.get("wall_time", time.time()), ck.get("block_size")
+
+
+
+def _safe_load_any(path: pathlib.Path, tgt: nn.Module, key: str | None = None) -> int:
+    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"]
+    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 "cfg" in sd:
+        return dict(sd["cfg"])
+    return None
+
+
+# ───────────────────────── Training logic ─────────────────────────
+def _loss_float(x: torch.Tensor) -> float:
+    try:
+        return float(x.detach().float().cpu().item())
+    except Exception:
+        return float(x.item())
+
+
+
+def _forward_train_losses(args, core, ar_h, sat_h, ids, ce_tok, ce_gate):
+    h_ar = core(ids, causal_mask(ids.size(1)))
+    logits_ar = ar_h(h_ar)[:, :-1]
+    loss_ar = ce_tok(logits_ar.float().reshape(-1, VOCAB), ids[:, 1:].reshape(-1))
+    if args.ar_only:
+        return loss_ar
+    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.float().reshape(-1, VOCAB), tgt_sat.reshape(-1))
+    if gate is not None:
+        loss_sat += EMIT_LAMBDA * ce_gate(gate.float(), torch.ones(ids.size(0), device=DEV, dtype=torch.long))
+    return loss_ar + loss_sat
+
+
+
+def _run_optimizer_step(args, opt, scaler, loss, trainable_params: Iterable[torch.nn.Parameter]):
+    trainable_params = list(trainable_params)
+    if args.amp and DEV.type == "cuda":
+        scaler.scale(loss).backward()
+        scaler.unscale_(opt)
+        if trainable_params:
+            nn.utils.clip_grad_norm_(trainable_params, 1.0)
+        scaler.step(opt)
+        scaler.update()
+        return
+
+    loss.backward()
+    if trainable_params:
+        nn.utils.clip_grad_norm_(trainable_params, 1.0)
+    RUNTIME.optimizer_step(opt)
+    if RUNTIME.is_tt:
+        RUNTIME.sync(wait=True)
+
+
+
+def _maybe_handle_oom(e: RuntimeError) -> bool:
+    msg = str(e).lower()
+    return (
+        "out of memory" in msg
+        or "cuda out of memory" in msg
+        or "resource exhausted" in msg
+        or "failed to allocate" in msg
+    )
+
+
+
+def _train_phase(
+    args,
+    phase_name: str,
+    core,
+    ar_h,
+    sat_h,
+    opt,
+    scaler,
+    start_step,
+    seen_tok,
+    resume_wall_time,
+    cfg,
+    source,
+    steps,
+    block_size,
+    batch_size,
+    chat_cfg: dict,
+    max_ckpts: Optional[int],
+    target_tokens_override: Optional[int] = None,
+    tie_weights: bool = False,
+    streaming: bool = True,
+):
+    BLOCK = block_size
+    BATCH = batch_size
+    if target_tokens_override is not None:
+        target_tokens = target_tokens_override
+    else:
+        ratio = 51.2 if args.chilla_max_double else 25
+        param_count = _count_enabled_params(core, ar_h, sat_h)
+        target_tokens = int(ratio * param_count)
+
+    if steps:
+        phase_target_tokens = steps * BLOCK * BATCH
+        total_tokens_needed = seen_tok + phase_target_tokens
+    else:
+        total_tokens_needed = target_tokens
+        if total_tokens_needed <= seen_tok:
+            print(f"[{phase_name}] target {total_tokens_needed} already reached.")
+            return start_step, seen_tok, resume_wall_time
+
+    stream = token_stream(
+        source,
+        total_tokens_needed,
+        seed=42,
+        chat=chat_cfg.get("chat", False),
+        chat_messages_key=chat_cfg.get("key", "messages"),
+        sft_add_generation_prompt=chat_cfg.get("gen_prompt", False),
+        dataset_field_text=chat_cfg.get("text_field", "text"),
+        streaming=streaming,
+    )
+
+    ce_tok = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)
+    ce_gate = nn.CrossEntropyLoss()
+    pbar = SafeProgress(total=total_tokens_needed, initial=seen_tok, unit="tok")
+    grow_plan = _parse_grow_plan(args.grow_plan) if args.auto_grow else []
+    buf: List[int] = []
+    batch_accum: List[List[int]] = []
+    step = start_step
+    steps_since_last_grow = 0
+    oom_retries = 0
+    max_oom_retries = 2
+
+    now_wall = time.time()
+    last_save_mono = time.monotonic() - (now_wall - (resume_wall_time or now_wall))
+    print(f"[{phase_name}] Starting. Goal: {total_tokens_needed:,} tokens. Batch={BATCH}, Block={BLOCK}")
+    print(f"[{phase_name}] BACKEND={RUNTIME.backend} AR_ONLY={args.ar_only} TIE_WEIGHTS={tie_weights} STREAMING={streaming}")
+    if RUNTIME.is_tt:
+        print(
+            f"[{phase_name}] TT dtype={str(RUNTIME.dtype).replace('torch.', '')} opt_level={args.tt_optimization_level} spmd={RUNTIME.spmd} devices={RUNTIME.num_devices}"
+        )
+
+    step_start_time = time.monotonic()
+    tok_per_sec_avg = 0.0
+    trainable_params = [p for p in list(core.parameters()) + list(ar_h.parameters()) + list(sat_h.parameters()) if p.requires_grad]
+
+    while seen_tok < total_tokens_needed:
+        try:
+            while len(buf) < BLOCK:
+                buf.append(next(stream))
+        except StopIteration:
+            break
+
+        seq = buf[:BLOCK]
+        buf = buf[BLOCK:]
+        batch_accum.append(seq)
+        if len(batch_accum) < BATCH:
+            continue
+
+        ids = torch.tensor(batch_accum, device=DEV, dtype=torch.long)
+        batch_accum = []
+        if RUNTIME.is_tt:
+            RUNTIME.maybe_mark_batch_sharding(ids)
+
+        try:
+            opt.zero_grad(set_to_none=True)
+            with amp(args.amp):
+                loss = _forward_train_losses(args, core, ar_h, sat_h, ids, ce_tok, ce_gate)
+            _run_optimizer_step(args, opt, scaler, loss, trainable_params)
+            retie_weights(core, ar_h, tie_weights)
+        except RuntimeError as e:
+            if _maybe_handle_oom(e):
+                batch_accum = []
+                opt.zero_grad(set_to_none=True)
+                if DEV.type == "cuda":
+                    torch.cuda.empty_cache()
+                    torch.cuda.synchronize()
+                oom_retries += 1
+                if oom_retries <= max_oom_retries:
+                    print(f"\n[{phase_name} OOM] Retry {oom_retries}/{max_oom_retries} at Batch={BATCH}, clearing caches...")
+                    time.sleep(4)
+                    continue
+                oom_retries = 0
+                if BATCH > 1:
+                    print(f"\n[{phase_name} OOM] Reducing Batch: {BATCH} -> {BATCH - 1}")
+                    BATCH -= 1
+                    time.sleep(4)
+                else:
+                    if grow_plan:
+                        smaller = [b for b in grow_plan if b < BLOCK]
+                        new_block = smaller[-1] if smaller else max(128, BLOCK // 2)
+                    else:
+                        new_block = max(128, BLOCK // 2)
+                    print(f"\n[{phase_name} OOM] Reducing Block: {BLOCK} -> {new_block}")
+                    BLOCK = new_block
+                    time.sleep(4)
+                steps_since_last_grow = 0
+                continue
+            raise
+
+        step += 1
+        oom_retries = 0
+        toks_processed = BLOCK * BATCH
+        seen_tok += toks_processed
+        pbar.update(toks_processed)
+        loss_value = _loss_float(loss)
+        pbar.set_postfix(loss=f"{loss_value:.3f}", B=BATCH, L=BLOCK)
+
+        step_elapsed = time.monotonic() - step_start_time
+        tok_per_sec_now = toks_processed / step_elapsed if step_elapsed > 0 else 0.0
+        tok_per_sec_avg = 0.9 * tok_per_sec_avg + 0.1 * tok_per_sec_now if tok_per_sec_avg > 0 else tok_per_sec_now
+        step_start_time = time.monotonic()
+        write_status(step, seen_tok, loss_value, BATCH, BLOCK, tok_per_sec_avg, phase_name)
+
+        if args.save_every_sec > 0:
+            now_mono = time.monotonic()
+            if now_mono - last_save_mono >= args.save_every_sec:
+                ck_name = f"{phase_name}_step{step:08d}.pt"
+                save_ckpt(
+                    pathlib.Path(args.save_dir) / ck_name,
+                    core,
+                    ar_h,
+                    sat_h,
+                    opt,
+                    scaler,
+                    meta={
+                        "cfg": cfg,
+                        "step": step,
+                        "seen_tok": seen_tok,
+                        "wall_time": time.time(),
+                        "tie_weights": tie_weights,
+                        "block_size": BLOCK,
+                        "batch_size": BATCH,
+                    },
+                )
+                _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts)
+                last_save_mono = now_mono
+
+        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):
+                        BLOCK = grow_plan[idx + 1]
+                        print(f"[{phase_name} Grow] Block -> {BLOCK}")
+                        if DEV.type == "cuda":
+                            torch.cuda.empty_cache()
+                except ValueError:
+                    grow_plan = sorted(set(grow_plan + [BLOCK]))
+
+    pbar.close()
+    save_ckpt(
+        pathlib.Path(args.save_dir) / f"{phase_name}_final.pt",
+        core,
+        ar_h,
+        sat_h,
+        opt,
+        scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "tie_weights": tie_weights,
+            "block_size": BLOCK,
+            "batch_size": BATCH,
+        },
+    )
+    return step, seen_tok, time.time()
+
+
+# ───────────────────────── Main orchestrator ─────────────────────────
+def _build_models(cfg, tie_weights: bool):
+    core = Encoder(cfg, tie_weights=tie_weights)
+    ar_h = ARHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None)
+    sat_h = SATHead(cfg["d"], mode="var")
+    retie_weights(core, ar_h, tie_weights)
+    return core, ar_h, sat_h
+
+
+
+def _maybe_cast_models_for_runtime(core, ar_h, sat_h):
+    if RUNTIME.is_tt and RUNTIME.dtype == torch.bfloat16:
+        core = core.to(dtype=torch.bfloat16)
+        ar_h = ar_h.to(dtype=torch.bfloat16)
+        sat_h = sat_h.to(dtype=torch.bfloat16)
+        retie_weights(core, ar_h, True if getattr(core, "tie_weights", False) or getattr(ar_h, "tie_weights", False) else False)
+    return core, ar_h, sat_h
+
+
+
+def _move_models_to_device(core, ar_h, sat_h, tie_weights: bool):
+    core = core.to(DEV)
+    ar_h = ar_h.to(DEV)
+    sat_h = sat_h.to(DEV)
+    retie_weights(core, ar_h, tie_weights)
+    return core, ar_h, sat_h
+
+
+
+def _maybe_compile_models(args, core, ar_h, sat_h, tie_weights: bool):
+    if not args.compile:
+        return core, ar_h, sat_h
+    if RUNTIME.is_tt:
+        print("[tt-xla] Skipping torch.compile for training stability; TT-XLA lazy compilation is still active.")
+        return core, ar_h, sat_h
+    if hasattr(torch, "compile"):
+        print("[torch.compile] Compiling model...")
+        core = torch.compile(core, mode="reduce-overhead")
+        ar_h = torch.compile(ar_h, mode="reduce-overhead")
+        sat_h = torch.compile(sat_h, mode="reduce-overhead")
+        retie_weights(core, ar_h, tie_weights)
+        print("[torch.compile] Done.")
+    return core, ar_h, sat_h
+
+
+
+def train(args):
+    setup_runtime(args)
+    cfg = PRESETS[args.preset].copy()
+    tie_weights = args.tie_weights
+    print_expansion_info(cfg, tie_weights)
+
+    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.update({k: v for k, v in prev_cfg.items() if k in cfg})
+        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
+
+    print(f"Config: {cfg}")
+    core, ar_h, sat_h = _build_models(cfg, tie_weights=tie_weights)
+
+    total_params = _count_enabled_params(core, ar_h, sat_h)
+    print(f"Total parameters: {total_params:,}")
+    if tie_weights:
+        print(f"{Colors.WARN}[weight-tying] Embedding and LM head share weights{Colors.RESET}")
+
+    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")
+            _safe_load_any(src, ar_h, key="ar")
+            _safe_load_any(src, sat_h, key="sat")
+            retie_weights(core, ar_h, tie_weights)
+            if loaded:
+                print(f"Warm-start loaded from {src}")
+
+    core, ar_h, sat_h = _maybe_cast_models_for_runtime(core, ar_h, sat_h)
+    core, ar_h, sat_h = _move_models_to_device(core, ar_h, sat_h, tie_weights)
+
+    _phase_freeze(core, freeze_core=args.freeze_core, unfreeze_ln=args.unfreeze_ln, train_emb=args.train_emb)
+
+    opt = torch.optim.AdamW(
+        [
+            {"params": [p for p in core.parameters() if p.requires_grad], "lr": args.lr_core},
+            {"params": ar_h.parameters(), "lr": args.lr_head},
+            {"params": sat_h.parameters(), "lr": args.lr_head},
+        ]
+    )
+    scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
+
+    start_step, seen_tok, last_wall, resumed_block = 0, 0, None, None
+    if args.resume and not args.fresh:
+        start_step, seen_tok, last_wall, resumed_block = load_ckpt(pathlib.Path(args.resume), core, ar_h, sat_h, opt, scaler)
+        retie_weights(core, ar_h, tie_weights)
+        print(f"Resumed from step {start_step}" + (f", block_size={resumed_block}" if resumed_block else ""))
+
+    core, ar_h, sat_h = _maybe_compile_models(args, core, ar_h, sat_h, tie_weights)
+
+    step, seen_tok, last_wall = _train_phase(
+        args,
+        "pretrain",
+        core,
+        ar_h,
+        sat_h,
+        opt,
+        scaler,
+        start_step,
+        seen_tok,
+        last_wall,
+        cfg,
+        args.source,
+        args.steps,
+        (resumed_block if resumed_block and args.auto_grow else None) or args.block or DEFAULT_BLOCK,
+        args.batch_size or DEFAULT_BATCH,
+        chat_cfg={
+            "chat": args.chat,
+            "key": args.chat_messages_key,
+            "gen_prompt": args.sft_add_generation_prompt,
+            "text_field": args.dataset_field_text,
+        },
+        max_ckpts=args.max_ckpts,
+        target_tokens_override=args.target_tokens,
+        tie_weights=tie_weights,
+    )
+
+    if (not args.after_sft_source) and (args.after_sft_steps and args.after_sft_steps > 0):
+        args.after_sft_source = DEFAULT_AFTER_SFT_SOURCES
+        args.after_sft_chat = True
+        if args.after_sft_add_generation_prompt is None:
+            args.after_sft_add_generation_prompt = True
+        if not args.after_sft_block:
+            args.after_sft_block = DEFAULT_AFTER_SFT_BLOCK
+
+    if args.after_sft_source and args.after_sft_steps and args.after_sft_steps > 0:
+        print("\n[Orchestrator] Starting Post-Pretraining SFT Phase...")
+        _phase_freeze(
+            core,
+            freeze_core=args.after_sft_freeze_core,
+            unfreeze_ln=args.after_sft_unfreeze_ln,
+            train_emb=args.after_sft_train_emb,
+        )
+        opt = torch.optim.AdamW(
+            [
+                {"params": [p for p in core.parameters() if p.requires_grad], "lr": args.after_sft_lr_core or args.lr_core},
+                {"params": ar_h.parameters(), "lr": args.after_sft_lr_head or args.lr_head},
+                {"params": sat_h.parameters(), "lr": args.after_sft_lr_head or args.lr_head},
+            ]
+        )
+        step, seen_tok, last_wall = _train_phase(
+            args,
+            "sft",
+            core,
+            ar_h,
+            sat_h,
+            opt,
+            scaler,
+            step,
+            seen_tok,
+            last_wall,
+            cfg,
+            args.after_sft_source,
+            args.after_sft_steps,
+            args.after_sft_block or DEFAULT_AFTER_SFT_BLOCK,
+            args.batch_size or DEFAULT_BATCH,
+            chat_cfg={
+                "chat": args.after_sft_chat,
+                "key": args.after_sft_chat_messages_key,
+                "gen_prompt": args.after_sft_add_generation_prompt if args.after_sft_add_generation_prompt is not None else args.sft_add_generation_prompt,
+                "text_field": args.after_sft_dataset_field_text,
+            },
+            max_ckpts=args.max_ckpts,
+            target_tokens_override=None,
+            tie_weights=tie_weights,
+            streaming=False,
+        )
+
+    save_ckpt(
+        pathlib.Path(args.save_dir) / "final.pt",
+        core,
+        ar_h,
+        sat_h,
+        opt,
+        scaler,
+        meta={
+            "cfg": cfg,
+            "step": step,
+            "seen_tok": seen_tok,
+            "wall_time": time.time(),
+            "tie_weights": tie_weights,
+            "block_size": args.block or DEFAULT_BLOCK,
+            "batch_size": args.batch_size or DEFAULT_BATCH,
+        },
+    )
+    print("🎉 All Training Complete")
+
+
+# ───────────────────────── Sampling / inference ─────────────────────────
+def _apply_penalties(logits, ids, n, rep_p, pres_p, freq_p):
+    if ids.numel() == 0:
+        return logits
+    hist = ids[0, -n:].long() if n > 0 else ids[0].long()
+    uniq, counts = torch.unique(hist, return_counts=True)
+    if pres_p or freq_p:
+        logits[..., uniq] -= pres_p + freq_p * counts.to(logits.dtype)
+    if rep_p != 1.0:
+        sel = logits[..., uniq]
+        logits[..., uniq] = torch.where(sel > 0, sel / rep_p, sel * rep_p)
+    return logits
+
+
+
+def _sample(logits, T, top_k, top_p, min_p, greedy):
+    if greedy:
+        return logits.argmax(-1, keepdim=True)
+    probs = (logits / max(T, 1e-8)).softmax(-1)
+    if top_k:
+        v, i = torch.topk(probs, min(top_k, probs.size(-1)))
+        probs = torch.zeros_like(probs).scatter_(-1, i, v)
+    if top_p < 1.0:
+        s_probs, s_idx = torch.sort(probs, descending=True, dim=-1)
+        keep = (torch.cumsum(s_probs, -1) <= top_p).to(probs.dtype)
+        probs = torch.zeros_like(probs).scatter_(-1, s_idx, s_probs * keep)
+    if min_p > 0:
+        probs[probs < min_p] = 0
+    if probs.sum() == 0:
+        return logits.argmax(-1, keepdim=True)
+    return probs.div_(probs.sum()).multinomial(1)
+
+
+
+def _sample_on_cpu(logits_device, ids_device, args):
+    logits = logits_device.detach().float().cpu()
+    ids = ids_device.detach().cpu()
+    logits = _apply_penalties(
+        logits,
+        ids,
+        args.penalty_last_n,
+        args.repetition_penalty,
+        args.presence_penalty,
+        args.frequency_penalty,
+    )
+    nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
+    return nxt.to(DEV)
+
+
+@torch.no_grad()
+def _infer_tt_static(args, core, ar_h, sat_h, ids):
+    prompt_len = ids.size(1)
+    total_len = prompt_len + args.max_new
+    work = torch.full((1, total_len), PAD_ID, dtype=torch.long, device=DEV)
+    work[:, :prompt_len] = ids
+
+    if args.mode == "ar":
+        for step in range(args.max_new):
+            cur_len = prompt_len + step
+            h = core(work, causal_padded_mask(total_len, cur_len))
+            logits = ar_h(h)[:, cur_len - 1]
+            nxt = _sample_on_cpu(logits, work[:, :cur_len], args)
+            work[:, cur_len] = nxt.squeeze(-1)
+        return work
+
+    added = 0
+    while added < args.max_new:
+        cur_len = prompt_len + added
+        h = core(work, sat_padded_mask(total_len, cur_len))
+        start = max(0, cur_len - SAT_BLOCK)
+        h_last = h[:, start:cur_len]
+        if h_last.size(1) < SAT_BLOCK:
+            pad = torch.zeros(
+                h_last.size(0),
+                SAT_BLOCK - h_last.size(1),
+                h_last.size(2),
+                device=h_last.device,
+                dtype=h_last.dtype,
+            )
+            h_last = torch.cat([pad, h_last], dim=1)
+        logits_all, gate = sat_h(h_last)
+        stride = SAT_BLOCK if (not args.var or gate is None) else (gate.float().softmax(-1).cpu().multinomial(1).item() + 1)
+        for i in range(int(stride)):
+            if added >= args.max_new:
+                break
+            logits = logits_all[:, i]
+            nxt = _sample_on_cpu(logits, work[:, :cur_len], args)
+            work[:, cur_len] = nxt.squeeze(-1)
+            cur_len += 1
+            added += 1
+    return work
+
+
+@torch.no_grad()
+def infer(args):
+    setup_runtime(args)
+    if args.mode == "ar":
+        if args.temperature is None:
+            args.temperature = 0.7
+        if args.top_k is None:
+            args.top_k = 0
+        if args.repetition_penalty is None:
+            args.repetition_penalty = 1.3
+        if args.presence_penalty is None:
+            args.presence_penalty = 0.0
+        if args.frequency_penalty is None:
+            args.frequency_penalty = 0.3
+        if args.penalty_last_n is None:
+            args.penalty_last_n = 128
+        if args.var is None:
+            args.var = False
+    else:
+        if args.temperature is None:
+            args.temperature = 0.5
+        if args.top_k is None:
+            args.top_k = 30
+        if args.repetition_penalty is None:
+            args.repetition_penalty = 2.0
+        if args.presence_penalty is None:
+            args.presence_penalty = 0.6
+        if args.frequency_penalty is None:
+            args.frequency_penalty = 1.0
+        if args.penalty_last_n is None:
+            args.penalty_last_n = 200
+        if args.var is None:
+            args.var = True
+
+    path = _resolve_ckpt(pathlib.Path(args.ckpt)) or pathlib.Path(args.ckpt)
+    sd = torch.load(path, map_location="cpu")
+    cfg = sd["cfg"]
+    tie_weights = sd.get("tie_weights", False)
+    uk_time = get_uk_time()
+    ckpt_name = path.name
+
+    print("┌─────────────────────────────────────────────────┐")
+    print(f"│ INFERENCE @ {uk_time:<35s} │")
+    print("├─────────────────────────────────────────────────┤")
+    print(f"│ Checkpoint: {ckpt_name:<35s} │")
+    print("└─────────────────────────────────────────────────┘")
+    print_expansion_info(cfg, tie_weights)
+
+    core, ar_h, sat_h = _build_models(cfg, tie_weights=tie_weights)
+    core.load_state_dict(sd["core"])
+    ar_h.load_state_dict(sd["ar"])
+    sat_h.load_state_dict(sd["sat"])
+    retie_weights(core, ar_h, tie_weights)
+
+    if RUNTIME.is_tt and args.tt_dtype == "bf16":
+        core = core.to(dtype=torch.bfloat16)
+        ar_h = ar_h.to(dtype=torch.bfloat16)
+        sat_h = sat_h.to(dtype=torch.bfloat16)
+        retie_weights(core, ar_h, tie_weights)
+    elif getattr(args, "fp16", False):
+        core.half()
+        ar_h.half()
+        sat_h.half()
+        retie_weights(core, ar_h, tie_weights)
+        print(f"{Colors.INFO}Using fp16 inference{Colors.RESET}")
+
+    core, ar_h, sat_h = _move_models_to_device(core, ar_h, sat_h, tie_weights)
+    core.eval()
+    ar_h.eval()
+    sat_h.eval()
+
+    total_params = _count_enabled_params(core, ar_h, sat_h)
+    if total_params >= 1_000_000_000:
+        param_str = f"{total_params / 1_000_000_000:.2f}B"
+    elif total_params >= 1_000_000:
+        param_str = f"{total_params / 1_000_000:.2f}M"
+    elif total_params >= 1_000:
+        param_str = f"{total_params / 1_000:.2f}K"
+    else:
+        param_str = f"{total_params}"
+    print(f"Model size: {param_str} parameters ({total_params:,})")
+
+    prompt_tokens = tok.encode(args.prompt)
+    prompt_len = len(prompt_tokens)
+    ids = torch.tensor([prompt_tokens], device=DEV, dtype=torch.long)
+    if ids.size(1) == 0:
+        ids = torch.tensor([[EOS]], device=DEV, dtype=torch.long)
+        prompt_len = 1
+
+    mode_str = args.mode if args.mode == "ar" else f"sat-{'var' if args.var else 'fixed'}"
+    print(f"{Colors.INFO}Generating ({mode_str}) on backend={RUNTIME.backend}...{Colors.RESET}")
+
+    start = time.time()
+    if RUNTIME.is_tt:
+        ids = _infer_tt_static(args, core, ar_h, sat_h, ids)
+    elif args.mode == "ar":
+        h, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True, total_seq_len=ids.size(1))
+        for _ in range(args.max_new):
+            logits = ar_h(h)[:, -1]
+            logits = _apply_penalties(
+                logits,
+                ids,
+                args.penalty_last_n,
+                args.repetition_penalty,
+                args.presence_penalty,
+                args.frequency_penalty,
+            )
+            nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
+            ids = torch.cat([ids, nxt], 1)
+            h, kvs = core(ids[:, -1:], None, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1))
+    else:
+        cached_len = ids.size(1)
+        h, kvs = core(ids, sat_mask(ids.size(1)), use_cache=True, total_seq_len=cached_len)
+        added = 0
+        while added < args.max_new:
+            logits_all, gate = sat_h(h[:, -SAT_BLOCK:])
+            stride = SAT_BLOCK if (not args.var or gate is None) else (gate.softmax(-1).multinomial(1).item() + 1)
+            new_tokens = []
+            for i in range(int(stride)):
+                logits = logits_all[:, i]
+                logits = _apply_penalties(
+                    logits,
+                    ids,
+                    args.penalty_last_n,
+                    args.repetition_penalty,
+                    args.presence_penalty,
+                    args.frequency_penalty,
+                )
+                nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
+                new_tokens.append(nxt)
+                ids = torch.cat([ids, nxt], 1)
+                added += 1
+                if added >= args.max_new:
+                    break
+            if added >= args.max_new:
+                break
+            new_ids = torch.cat(new_tokens, dim=1)
+            mask = sat_mask_cached(new_ids.size(1), cached_len)
+            h, kvs = core(new_ids, mask, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1))
+            cached_len = ids.size(1)
+
+    if RUNTIME.is_tt:
+        RUNTIME.sync(wait=True)
+    elapsed = time.time() - start
+
+    all_tokens = ids[0].detach().cpu().tolist()
+    gen_tokens = len(all_tokens) - prompt_len
+    tok_per_sec = gen_tokens / elapsed if elapsed > 0 else 0.0
+    prompt_text = tok.decode(all_tokens[:prompt_len], skip_special_tokens=True)
+    gen_text = tok.decode(all_tokens[prompt_len:], skip_special_tokens=True)
+    print(f"{Colors.PROMPT}{prompt_text}{Colors.RESET}{gen_text}")
+    print(f"{Colors.INFO}[{elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s]{Colors.RESET}")
+
+
+# ───────────────────────── CLI ─────────────────────────
+def main():
+    ap = argparse.ArgumentParser(description="AGILLM Expansion Ratio Testing (CUDA / Tenstorrent / CPU)")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    tr = sub.add_parser("train")
+    tr.add_argument("--backend", choices=["auto", "cuda", "tt", "cpu"], default="auto")
+    tr.add_argument("--preset", choices=PRESETS.keys(), default="nano_3x")
+    tr.add_argument("--rank", type=int)
+    tr.add_argument("--block", type=int, default=DEFAULT_BLOCK)
+    tr.add_argument("--batch_size", type=int, default=DEFAULT_BATCH)
+    tr.add_argument("--source", default=DEFAULT_PRETRAIN_SOURCES)
+    tr.add_argument("--target_tokens", type=int)
+    tr.add_argument("--steps", type=int)
+    tr.add_argument("--amp", action="store_true")
+    tr.add_argument("--compile", action="store_true", help="Use torch.compile on CUDA. TT path skips this for stability.")
+    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")
+    tr.add_argument("--warmstart_from", type=str)
+    tr.add_argument("--fresh", action="store_true")
+    tr.add_argument("--max_ckpts", type=int, default=None)
+    tr.add_argument("--chilla_max_double", action="store_true")
+    tr.add_argument("--tie_weights", action="store_true")
+    tr.add_argument("--ar_only", action="store_true")
+    tr.add_argument("--freeze_core", action="store_true")
+    tr.add_argument("--unfreeze_ln", action="store_true")
+    tr.add_argument("--train_emb", action="store_true")
+    tr.add_argument("--lr_core", type=float, default=LR_CORE)
+    tr.add_argument("--lr_head", type=float, default=LR_HEAD)
+    tr.add_argument("--label_smoothing", type=float, default=0.1)
+    tr.add_argument("--chat", action="store_true")
+    tr.add_argument("--chat_messages_key", default="messages")
+    tr.add_argument("--dataset_field_text", default="text")
+    tr.add_argument("--sft_add_generation_prompt", action="store_true")
+    tr.add_argument("--auto_grow", action="store_true")
+    tr.add_argument("--grow_plan", default="576,640,768,896,1024,1122")
+    tr.add_argument("--grow_every_steps", type=int, default=50000)
+    tr.add_argument("--after_sft_source", default="")
+    tr.add_argument("--after_sft_steps", type=int, default=0)
+    tr.add_argument("--after_sft_chat", action="store_true")
+    tr.add_argument("--after_sft_chat_messages_key", default="messages")
+    tr.add_argument("--after_sft_dataset_field_text", default="text")
+    tr.add_argument("--after_sft_add_generation_prompt", type=bool, default=None)
+    tr.add_argument("--after_sft_block", type=int, default=0)
+    tr.add_argument("--after_sft_freeze_core", action="store_true")
+    tr.add_argument("--after_sft_unfreeze_ln", action="store_true")
+    tr.add_argument("--after_sft_train_emb", action="store_true")
+    tr.add_argument("--after_sft_lr_core", type=float, default=0.0)
+    tr.add_argument("--after_sft_lr_head", type=float, default=0.0)
+    tr.add_argument("--tt_dtype", choices=["fp32", "bf16"], default="bf16")
+    tr.add_argument("--tt_bfp8", action="store_true")
+    tr.add_argument("--tt_weight_bfp8", action="store_true")
+    tr.add_argument("--tt_optimization_level", type=int, default=1)
+    tr.add_argument("--tt_trace", action="store_true")
+    tr.add_argument("--tt_trace_region_size", type=int, default=10_000_000)
+    tr.add_argument("--tt_spmd", action="store_true", help="Experimental: shard batch across visible TT chips.")
+
+    inf = sub.add_parser("infer")
+    inf.add_argument("--backend", choices=["auto", "cuda", "tt", "cpu"], default="auto")
+    inf.add_argument("--mode", choices=["ar", "sat"], required=True)
+    inf.add_argument("--ckpt", required=True)
+    inf.add_argument("--prompt", required=True)
+    inf.add_argument("--max_new", type=int, default=120)
+    inf.add_argument("--temperature", type=float, default=None)
+    inf.add_argument("--greedy", action="store_true")
+    inf.add_argument("--top_k", type=int, default=None)
+    inf.add_argument("--top_p", type=float, default=0.9)
+    inf.add_argument("--min_p", type=float, default=0.0)
+    inf.add_argument("--repetition_penalty", type=float, default=None)
+    inf.add_argument("--presence_penalty", type=float, default=None)
+    inf.add_argument("--frequency_penalty", type=float, default=None)
+    inf.add_argument("--penalty_last_n", type=int, default=None)
+    inf.add_argument("--var", action="store_true", default=None)
+    inf.add_argument("--no-var", dest="var", action="store_false")
+    inf.add_argument("--fp16", action="store_true", help="Use fp16 inference on CUDA/CPU-like backends.")
+    inf.add_argument("--tt_dtype", choices=["fp32", "bf16"], default="bf16")
+    inf.add_argument("--tt_bfp8", action="store_true")
+    inf.add_argument("--tt_weight_bfp8", action="store_true")
+    inf.add_argument("--tt_optimization_level", type=int, default=1)
+    inf.add_argument("--tt_trace", action="store_true")
+    inf.add_argument("--tt_trace_region_size", type=int, default=10_000_000)
+    inf.add_argument("--tt_spmd", action="store_true")
+
+    sub.add_parser("status")
+
+    args = ap.parse_args()
+    if args.cmd == "train":
+        train(args)
+    elif args.cmd == "status":
+        show_status()
+    else:
+        infer(args)
+
+
+if __name__ == "__main__":
+    main()