cortex/train.py

"""Training loop: data-parallel, bf16-ready, Orbax checkpoint/resume.

Device-count agnostic: builds a 1D mesh over all visible devices and shards the
batch along it (params/opt-state replicated). Runs on 1 CPU here, 8 chips on a
v5e-8 with no code change. Grad accumulation is a ``lax.scan`` over micro-batches
so only one gradient accumulator lives in memory.
"""

from __future__ import annotations

import os
import time
from concurrent.futures import ThreadPoolExecutor
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import nnx
from jax.sharding import NamedSharding, PartitionSpec as P
import orbax.checkpoint as ocp

from .config import ModelConfig, TrainConfig, RunConfig
from .model import CortexLM
from .latent_diffusion import CortexLatentDiffusion
from .losses import chunked_lm_loss


def build_model(mcfg: ModelConfig, rngs: nnx.Rngs):
    """Pick the architecture family by mcfg.arch."""
    if mcfg.arch == "latent_diffusion":
        return CortexLatentDiffusion(mcfg, rngs=rngs)
    return CortexLM(mcfg, rngs=rngs)
from .neuron_probe import dormancy, recycle
from .optim import build_optimizer, wsd_schedule


def make_mesh():
    return jax.make_mesh((len(jax.devices()),), ("data",))


def count_params(params) -> int:
    return sum(int(x.size) for x in jax.tree.leaves(params))


def _merge_safetensors_by_path(params, sft_path, repl):
    """Load a bf16 safetensors (keys '<idx>|<keystr>') BY PATH into `params`; leaves whose path is
    absent (brand-new params, e.g. learnable sinks) keep their current fresh init. Replicated out."""
    from safetensors.flax import load_file
    by_path = {k.split("|", 1)[1]: v for k, v in load_file(str(sft_path)).items()}
    paths_leaves, treedef = jax.tree_util.tree_flatten_with_path(params)
    out, n_load, n_fresh = [], 0, 0
    for path, leaf in paths_leaves:
        ps = jax.tree_util.keystr(path)
        if ps in by_path:
            out.append(jnp.asarray(by_path[ps]).astype(leaf.dtype)); n_load += 1
        else:
            out.append(leaf); n_fresh += 1
    print(f"[resume] safetensors merge: {n_load} tensors loaded, {n_fresh} kept fresh (new params)")
    return jax.device_put(jax.tree_util.tree_unflatten(treedef, out), repl)


def make_step(graphdef, update_fn, mcfg: ModelConfig, tcfg: TrainConfig, mesh):
    """update_fn(grads, opt_state, params) -> (new_params, new_opt_state, grad_norm).
    Standard (optax) and ZeRO-1 sharded optimizers both satisfy this signature."""
    latent = mcfg.arch == "latent_diffusion"

    def loss_fn(params, batch, key):
        model = nnx.merge(graphdef, params)
        if latent:                                               # Cortex-A 0.6: cosine + denoise + CE
            return model.compute_loss(batch, key, ce_chunk=tcfg.ce_chunk,
                                      z_loss_coef=tcfg.z_loss_coef)
        main_h, mtp_h, gate_aux = model.forward_train_hidden(batch)   # pre-head hidden + MoD aux
        total, metrics = chunked_lm_loss(model._head, main_h, mtp_h, batch, chunk=tcfg.ce_chunk,
                                         z_loss_coef=tcfg.z_loss_coef, mtp_weight=mcfg.mtp_loss_weight,
                                         mtp_decay=mcfg.mtp_decay)
        total = total + mcfg.ffn_gate_aux_weight * gate_aux       # MoD causal-predictor BCE (0 if gate off)
        return total, {**metrics, "gate_aux": gate_aux}

    def local_grads(params, micro_batches, key):          # micro_batches: [G, b_local, T]
        def accum(carry, xb):
            b, idx = xb
            # per-microbatch + per-shard noise key (diffusion); the AR path ignores it.
            k_i = jax.random.fold_in(jax.random.fold_in(key, jax.lax.axis_index("data")), idx)
            (loss, metrics), g = jax.value_and_grad(loss_fn, has_aux=True)(params, b, k_i)
            return jax.tree.map(lambda c, gg: c + gg.astype(c.dtype), carry, g), (loss, metrics)

        # accumulation dtype is explicit (cotangents follow the param dtype, which is bf16
        # under host-master): fp32 by default; tcfg.accum_dtype="bfloat16" frees the other
        # 1.17G of the accumulator at ~sqrt(G)*2^-9 relative grad noise (keystone-gated).
        adt = jnp.dtype(tcfg.accum_dtype)
        init = jax.tree.map(lambda x: jnp.zeros(x.shape, adt), params)
        G = micro_batches.shape[0]
        grads, (losses, metrics) = jax.lax.scan(accum, init, (micro_batches, jnp.arange(G)))
        grads = jax.tree.map(lambda x: x / G, grads)
        # average grads/metrics across the data-parallel shards
        grads = jax.lax.pmean(grads, "data")
        losses = jax.lax.pmean(losses, "data")
        metrics = jax.tree.map(lambda x: jax.lax.pmean(x, "data"), metrics)
        return grads, losses, metrics

    # shard_map: each device runs the forward/backward on its local batch shard
    # (params replicated) -> no ambiguous-sharding gather on the embedding lookup.
    sharded = jax.shard_map(
        local_grads, mesh=mesh,
        in_specs=(P(), P(None, "data", None), P()), out_specs=(P(), P(), P()),
        check_vma=False)

    def _sr_bf16(x32, key):
        """Stochastically round fp32 -> bf16: add 16 random low bits, truncate. Unbiased,
        so per-step deltas SMALLER than a bf16 ulp (ours: ~1e-4 vs ulp ~4e-3) still move
        the params in expectation instead of being swallowed by round-to-nearest."""
        bits = jax.lax.bitcast_convert_type(x32, jnp.uint32)
        rnd = jax.random.bits(key, x32.shape, jnp.uint32) & jnp.uint32(0xFFFF)
        return jax.lax.bitcast_convert_type(
            (bits + rnd) & jnp.uint32(0xFFFF0000), jnp.float32).astype(jnp.bfloat16)

    # Donate params + opt_state: they are both inputs and outputs, so without
    # donation XLA double-buffers them (old + new live at once during the update),
    # which is the param-scale peak that OOMs a 584M Muon model on a 16G v5e core.
    # Donation lets the update run in place, freeing ~(params fp32 + opt_state) ~= 3.8G.
    # Host-master mode: update_fn returns the DELTAS; the device applies them to its own
    # bf16 params in-jit via STOCHASTIC ROUNDING (so the param chain never waits on the
    # host) and returns both -- the host folds the same deltas into fp32 masters
    # asynchronously, and the periodic resync corrects any accumulated rounding.
    @partial(jax.jit, donate_argnums=(0, 1))
    def step(params, opt_state, micro_batches, key):
        grads, losses, metrics = sharded(params, micro_batches, key)
        out_tree, opt_state, gnorm = update_fn(grads, opt_state, params)  # new params OR deltas
        out = {k: jnp.mean(v) for k, v in metrics.items()}
        out["loss"] = jnp.mean(losses)
        out["grad_norm"] = gnorm
        if tcfg.host_master:
            pl, tdef = jax.tree_util.tree_flatten(params)
            ul = jax.tree_util.tree_leaves(out_tree)
            if tcfg.host_param_dtype == "bfloat16":
                nl = [_sr_bf16(p.astype(jnp.float32) + u.astype(jnp.float32),
                               jax.random.fold_in(key, i)) for i, (p, u) in enumerate(zip(pl, ul))]
            else:                                        # fp32 device params (keystone): exact add
                nl = [(p + u).astype(p.dtype) for p, u in zip(pl, ul)]
            return (jax.tree_util.tree_unflatten(tdef, nl), out_tree), opt_state, out
        return out_tree, opt_state, out

    return step


def run_training(mcfg: ModelConfig, tcfg: TrainConfig, run: RunConfig, loader, val_sets=None, tok=None):
    mesh = make_mesh()
    repl = NamedSharding(mesh, P())
    dshard = NamedSharding(mesh, P(None, "data", None))   # shard batch axis of [G, B, T]
    eshard = NamedSharding(mesh, P("data", None))         # shard batch axis of a [B, T] val batch

    latent = mcfg.arch == "latent_diffusion"
    model = build_model(mcfg, rngs=nnx.Rngs(run.seed))
    graphdef, params = nnx.split(model)
    params = jax.device_put(params, repl)

    if tcfg.zero_optimizer:                      # ZeRO-1: distributed Muon (whole matrices/device)
        from .optim_sharded import make_sharded_optimizer
        init_fn, update_fn = make_sharded_optimizer(tcfg, params, mesh,
                                                    emit_updates=tcfg.host_master)
        opt_state = init_fn(params)              # mu sharded across the data mesh
    else:
        tx = build_optimizer(tcfg)
        opt_state = jax.device_put(tx.init(params), repl)
        def update_fn(grads, opt_state, params):
            gnorm = optax.global_norm(grads)
            updates, opt_state = tx.update(grads, opt_state, params)
            if tcfg.host_master:                 # host-master: deltas out, host applies to fp32
                hdt = jnp.dtype(tcfg.host_param_dtype)
                return jax.tree.map(lambda u: u.astype(hdt), updates), opt_state, gnorm
            return optax.apply_updates(params, updates), opt_state, gnorm

    step_jit = make_step(graphdef, update_fn, mcfg, tcfg, mesh)
    lr_sched = wsd_schedule(tcfg.muon_lr, tcfg.warmup_steps, tcfg.stable_steps,
                            tcfg.decay_steps, tcfg.end_lr_frac, tcfg.decay_type)

    @jax.jit
    def probe_jit(params, batch):                        # per-neuron mean|act|: [n_layers, d_ff]
        return nnx.merge(graphdef, params).collect_mlp_acts(batch)

    def _eval_local(params, batch):                     # batch: local [b, T] shard
        m = nnx.merge(graphdef, params)
        if latent:                                       # dual CE: WITH the planner plan vs WITHOUT it
            ce_w = m.eval_ce(batch, ce_chunk=tcfg.ce_chunk)                      # backbone on
            ce_o = m.eval_ce(batch, ce_chunk=tcfg.ce_chunk, ablate_plan=True)    # backbone ablated
            return jax.lax.pmean(jnp.stack([ce_w, ce_o]), "data")               # [2]
        hidden = m._trunk(batch, route="prefill")        # inference routing; no MTP / no gate BCE
        _, met = chunked_lm_loss(m._head, hidden, [], batch, chunk=tcfg.ce_chunk,
                                 z_loss_coef=0.0, mtp_weight=0.0)
        return jax.lax.pmean(met["ce"], "data")
    eval_ce = jax.jit(jax.shard_map(                     # shard_map: same reason as the train step
        _eval_local, mesh=mesh, in_specs=(P(), P("data", None)), out_specs=P(), check_vma=False))

    def evaluate(params):                                # CE per held-out set (val_sets: name -> [nB,B,T])
        res = {}
        for name, arr in (val_sets or {}).items():
            vs = [np.asarray(eval_ce(params, jax.device_put(np.asarray(arr[i]), eshard)))
                  for i in range(arr.shape[0])]
            mean = np.mean(vs, axis=0)                    # scalar (cortex) or [2] (latent: with/without plan)
            res[name] = (float(mean[0]), float(mean[1])) if latent else float(mean)
        return res

    # ---- val text generation (latent arch): run the REAL inference pipeline -- cached
    # latent-AR planner + AR writer greedy decode -- on a few fixed prompts each eval,
    # so the log shows actual model output as it trains. ----
    sample_text = None
    if latent and tok is not None:
        from .model import compute_rope as _crope
        from .latent_diffusion import backbone_cfg as _bbcfg, writer_cfg as _wcfg
        from .config import BOS_ID, EOT_ID
        Pp, dM = mcfg.patch_size, mcfg.d_model
        GEN_NEW, GEN_MAXT = 48, 512                       # generate up to 48 tokens
        GEN_MAXC = GEN_MAXT // Pp
        GEN_TEMP, GEN_TOPP, GEN_REPPEN = 0.8, 0.95, 1.3   # val sampling: temperature / nucleus / rep-penalty
        _grng = np.random.default_rng(run.seed)
        _R = lambda x: jax.device_put(x, repl)            # fully replicate (same as the training params)
        _cc, _sc = _crope(GEN_MAXC, mcfg.head_dim, mcfg.rope_base, _bbcfg(mcfg), jnp.float32)
        _ct, _st = _crope(GEN_MAXT, mcfg.head_dim, mcfg.rope_base, _wcfg(mcfg), jnp.float32)
        PROMPTS = ["The capital of France is", "Once upon a time, there was a",
                   "Water is made of"]

        # B=1 generation interleaves jitted calls with host-side control; run each op as
        # shard_map with ALL specs replicated so there's no mesh sharding to reconcile.
        def _sm(fn, n_in, n_out=1):
            ins = tuple(P() for _ in range(n_in))
            outs = P() if n_out == 1 else tuple(P() for _ in range(n_out))
            return jax.jit(jax.shard_map(fn, mesh=mesh, in_specs=ins, out_specs=outs, check_vma=False))

        _g_embed = _sm(lambda p, ids: nnx.merge(graphdef, p).embed(ids), 2)
        _g_pred = _sm(lambda p, cm, pos, c: nnx.merge(graphdef, p).predict_chunk(cm, pos, c, _cc, _sc), 4, 2)
        _g_write = _sm(lambda p, tk, cv, pos, c: nnx.merge(graphdef, p).writer_step(tk, cv, pos, c, _ct, _st), 5, 2)
        _zc = _R(jnp.zeros((1, 1, dM), jnp.float32))      # null hint for chunk-0 tokens

        def _pick(lg, prev):                              # host-side top-p + repetition-penalty sample
            lg = np.asarray(lg, np.float32).copy()
            if prev and GEN_REPPEN != 1.0:                # penalize already-generated tokens (breaks greedy loops)
                idx = np.fromiter(set(prev), np.int64)
                lg[idx] = np.where(lg[idx] > 0, lg[idx] / GEN_REPPEN, lg[idx] * GEN_REPPEN)
            lg = lg / max(GEN_TEMP, 1e-3)
            p = np.exp(lg - lg.max()); p /= p.sum()
            order = np.argsort(-p)                         # nucleus: smallest set with cumprob >= top_p
            keep = order[:max(int(np.searchsorted(np.cumsum(p[order]), GEN_TOPP)) + 1, 1)]
            return int(_grng.choice(keep, p=p[keep] / p[keep].sum()))

        def sample_text(params, prompt, ablate_plan=False):
            ids = [BOS_ID] + tok.encode(prompt, add_eot=False)
            if len(ids) % Pp:                             # left-pad to a chunk boundary
                ids = [BOS_ID] * (Pp - len(ids) % Pp) + ids
            L = len(ids)
            nchunk = L // Pp
            emb_p = _g_embed(params, _R(jnp.asarray([ids], jnp.int32)))     # [1,L,d]
            pcache = _R(nnx.merge(graphdef, params).init_chunk_cache(1, GEN_MAXC))
            wcache = _R(nnx.merge(graphdef, params).init_writer_cache(1, GEN_MAXT))
            chats = []                                    # chats[c] = planner hint for chunk c+1
            for c in range(nchunk):
                cm = emb_p[:, c * Pp:(c + 1) * Pp].mean(1, keepdims=True)
                pc, pcache = _g_pred(params, cm, _R(jnp.int32(c)), pcache)
                chats.append(pc)

            def hint_for(t):                              # cond[t] = chats[(t+1)//Pp - 1] (causal)
                if ablate_plan:                           # no-backbone generation: zero the planner hint
                    return _zc
                si = (t + 1) // Pp - 1
                return chats[si] if 0 <= si < len(chats) else _zc

            logits = None
            for t in range(L):                            # prime the writer over the prompt
                logits, wcache = _g_write(params, _R(jnp.asarray([[ids[t]]], jnp.int32)),
                                          hint_for(t), _R(jnp.int32(t)), wcache)
            out, buf, pos, cur_c = [], [], L, nchunk
            for _ in range(GEN_NEW):
                if pos >= GEN_MAXT:
                    break
                nxt = _pick(logits[0], out)               # top-p + rep-penalty (was greedy argmax)
                if nxt == EOT_ID:
                    break
                out.append(nxt); buf.append(nxt)
                if len(buf) == Pp:                        # chunk complete -> planner step
                    cm = _g_embed(params, _R(jnp.asarray([buf], jnp.int32))).mean(1, keepdims=True)
                    pc, pcache = _g_pred(params, cm, _R(jnp.int32(cur_c)), pcache)
                    chats.append(pc); cur_c += 1; buf = []
                logits, wcache = _g_write(params, _R(jnp.asarray([[nxt]], jnp.int32)),
                                          hint_for(pos), _R(jnp.int32(pos)), wcache)
                pos += 1
            return tok.decode(out)

    out_dir = os.path.abspath(run.out_dir)
    os.makedirs(out_dir, exist_ok=True)
    mgr = ocp.CheckpointManager(
        out_dir, options=ocp.CheckpointManagerOptions(
            save_interval_steps=run.ckpt_every, max_to_keep=run.max_to_keep))

    start = 0
    if run.reinit_from_weights:
        # structure-changing resume (e.g. enabling learnable sinks): path-merge the master weights
        # from the bf16 snapshot, keep brand-new params at init, fresh optimizer, continue the schedule.
        sft = os.path.join(out_dir, "model.safetensors")
        assert os.path.exists(sft), f"--reinit-from-weights set but {sft} is missing"
        latest = mgr.latest_step() or 0
        params = _merge_safetensors_by_path(params, sft, repl)
        opt_state = init_fn(params) if tcfg.zero_optimizer else jax.device_put(tx.init(params), repl)
        if tcfg.zero_optimizer:
            opt_state = opt_state._replace(count=jax.device_put(jnp.asarray(latest + 1, jnp.int32), repl))
        start = latest + 1
        print(f"[resume] reinit-from-weights: master path-merged at step {latest}; new params fresh; "
              f"optimizer fresh (momentum re-accumulates); continuing at step {start}")
    elif mgr.latest_step() is not None:
        latest = mgr.latest_step()
        try:
            restored = mgr.restore(latest,
                                   args=ocp.args.StandardRestore({"params": params, "opt_state": opt_state}))
            params, opt_state = restored["params"], restored["opt_state"]
            print(f"[resume] restored step {latest}, continuing at {latest + 1}")
        except Exception:
            # opt-state structure changed (e.g. continuing a standard run under --zero-optimizer):
            # load the master weights, then re-init this optimizer fresh (momentum re-accumulates).
            assert tcfg.zero_optimizer, "checkpoint opt_state structure mismatch"
            std_tmpl = jax.device_put(build_optimizer(tcfg).init(params), repl)
            restored = mgr.restore(latest,
                                   args=ocp.args.StandardRestore({"params": params, "opt_state": std_tmpl}))
            params = restored["params"]
            # free the standard opt-state buffers (unused under ZeRO) NOW -- otherwise the
            # leftover ~3G squats in HBM and the jit_step program can't be allocated at runtime.
            for buf in jax.tree.leaves((std_tmpl, restored["opt_state"])):
                buf.delete()
            del std_tmpl, restored
            opt_state = init_fn(params)
            # seed the schedule counter to the resumed step so the WSD LR continues (no re-warmup)
            opt_state = opt_state._replace(
                count=jax.device_put(jnp.asarray(latest + 1, jnp.int32), repl))
            print(f"[resume] loaded params from a standard checkpoint at step {latest}; "
                  f"re-initialized the ZeRO optimizer (momentum fresh, schedule continued)")
        start = latest + 1

    # ---- resume defrag: checkpoint restore allocates the restored tree while the init
    # template is still live, leaving the device heap fragmented; the jit_step program
    # then needs CONTIGUOUS arenas and can fail allocation with gigabytes nominally
    # free (v24: the rung that ran v22 fresh OOM'd at compile after a resume).  A host
    # round-trip (pull -> free everything -> re-upload) compacts the heap once. ----
    if start > 0:
        t_df = time.time()
        shardings = jax.tree.map(lambda x: x.sharding, (params, opt_state))  # keep ZeRO shards
        host = jax.tree.map(lambda x: np.array(jax.device_get(x)), (params, opt_state))
        for buf in jax.tree.leaves((params, opt_state)):
            buf.delete()
        params, opt_state = jax.device_put(host, shardings)
        del host
        print(f"[resume] device heap defrag (host round-trip) in {time.time()-t_df:.1f}s")

    # ---- host-master offload: fp32 masters -> host numpy; device keeps a bf16 working copy.
    # MUST run after the resume/init block: optimizer init + checkpoint restore see fp32
    # params (Adam's v inits as zeros_like(param) -- a bf16 v would gut the accumulator).
    masters, _hdt = None, None
    if tcfg.host_master:
        import ml_dtypes
        _hdt = np.dtype(ml_dtypes.bfloat16) if tcfg.host_param_dtype == "bfloat16" else np.float32
        masters = jax.tree.map(                          # np.array: owned + writable (device_get
            lambda x: np.array(jax.device_get(x), np.float32), params)   # views are read-only)
        for buf in jax.tree.leaves(params):
            buf.delete()                                 # free the fp32 masters from HBM NOW
        params = jax.device_put(jax.tree.map(lambda m: m.astype(_hdt), masters), repl)
        sz = sum(m.nbytes for m in jax.tree.leaves(masters)) / 2**30
        print(f"[host-master] fp32 masters ({sz:.2f}G) -> host numpy; device params "
              f"{tcfg.host_param_dtype}; deltas applied on host each step")

    N = count_params(params)
    # Effective params/TOKEN for the MFU estimate: the latent backbone processes only
    # T/patch_size positions, so counting all params at full token rate would overstate
    # MFU ~(backbone share)x.  N_eff = backbone/P + everything-else (decoder, embed, head).
    N_eff = N
    if latent:
        bb_names = ("'subs'", "'block_kv'", "'aggregators'", "'latent_norm'", "'latent_head'")
        flat, _ = jax.tree_util.tree_flatten_with_path(params)
        bb = sum(int(x.size) for p, x in flat
                 if any(n in jax.tree_util.keystr(p) for n in bb_names))
        N_eff = bb // mcfg.patch_size + (N - bb)
        print(f"[train] latent arch: backbone {bb/1e6:.1f}M @ 1/{mcfg.patch_size} token rate, "
              f"rest {(N-bb)/1e6:.1f}M @ full -> N_eff {N_eff/1e6:.1f}M/token for MFU")
    G, B, T = max(1, tcfg.grad_accum), tcfg.micro_batch, tcfg.seq_len
    print(f"[train] params={N/1e6:.1f}M  devices={len(jax.devices())}  "
          f"micro_batch={B} seq={T} grad_accum={G}  tokens/step={G*B*T:,}")

    it = iter(loader)

    def next_macro():
        # The loader may hand back a CPU-prefetched [G, B, T] macro-batch directly
        # (overlaps gather with the device step); else stack G [B, T] micro-batches.
        x = next(it)
        if getattr(x, "ndim", 0) == 3:
            return x
        return np.stack([x] + [next(it) for _ in range(G - 1)], axis=0)

    history = []
    t0, toks = time.time(), 0
    session_start = time.time()
    final_step = max(start - 1, 0)
    nan_abort = False                                    # catastrophic-NaN early stop (save compute)
    ce_best = float("inf")                                # divergence early-stop: best CE seen this run
    def _host_apply(deltas):
        upd_np = jax.device_get(deltas)                  # one replicated copy of the deltas (D2H)
        def _ap(m, u):
            m += np.asarray(u, np.float32)               # fp32 master += delta, in place
            return m
        jax.tree.map(_ap, masters, upd_np)

    K = max(1, tcfg.master_sync_every)
    pool = ThreadPoolExecutor(max_workers=1) if (tcfg.host_master and K > 1) else None
    pending = []                                         # in-flight async master applies

    def _drain():
        for f in pending:
            f.result()
        pending.clear()

    for step in range(start, run.max_steps):
        host_mb = next_macro()
        mb = jax.device_put(host_mb, dshard)
        if tcfg.host_master:
            (params, deltas), opt_state, metrics = step_jit(
                params, opt_state, mb, jax.random.fold_in(jax.random.PRNGKey(run.seed), step))
            if pool is None:                             # K=1: synchronous (v19 behavior)
                _host_apply(deltas)
                params = jax.device_put(jax.tree.map(lambda m: m.astype(_hdt), masters), repl)
            else:                                        # lazy sync: device params stay hot;
                pending.append(pool.submit(_host_apply, deltas))   # masters track asynchronously
                if (step - start) % K == K - 1:          # resync wipes accumulated bf16 rounding
                    _drain()
                    params = jax.device_put(jax.tree.map(lambda m: m.astype(_hdt), masters), repl)
        else:
            params, opt_state, metrics = step_jit(
                params, opt_state, mb, jax.random.fold_in(jax.random.PRNGKey(run.seed), step))
        toks += G * B * T

        # dormant-neuron recycling: probe MLP utilization, reclaim dead units (>trigger)
        if tcfg.recycle_enabled and step > 0 and step % tcfg.recycle_every == 0:
            pb = jax.device_put(np.asarray(host_mb[0]), repl)           # one micro-batch [B, T]
            frac, masks = dormancy(np.asarray(probe_jit(params, pb)), tau=tcfg.dormant_tau)
            params, opt_state, info = recycle(params, opt_state, masks,
                                              key=jax.random.key(run.seed + step),
                                              trigger=tcfg.recycle_trigger)
            if info:
                params = jax.device_put(params, repl)
                opt_state = jax.device_put(opt_state, repl)
            print(f"[recycle] step {step:6d} | dormant max={frac.max()*100:4.0f}% "
                  f"mean={frac.mean()*100:4.0f}% (tau={tcfg.dormant_tau}) | "
                  f"recycled {sum(info.values())} units in {len(info)} layers "
                  + (str({L: info[L] for L in sorted(info)}) if info else "(none > trigger)"))

        if step % run.log_every == 0 or step == run.max_steps - 1:
            metrics = jax.block_until_ready(metrics)
            dt = max(time.time() - t0, 1e-6)
            tps = toks / dt
            mfu = 6 * N_eff * tps / (run.peak_tflops * 1e12)
            m = {k: float(v) for k, v in metrics.items()}
            # Catastrophic-NaN stop: if the LOSS itself is non-finite the model has truly broken
            # (the per-leaf grad guard handles localized grad NaNs and keeps the loss finite), so
            # halt NOW rather than burn the rest of the session -- and do NOT upload, so the last
            # good checkpoint on the Hub is preserved.
            if not np.isfinite(m["loss"]):
                print(f"[abort] step {step}: loss is non-finite ({m['loss']}) -> stopping immediately "
                      f"to save compute (NOT uploading; last good checkpoint preserved)")
                nan_abort = True
                break
            # Divergence early-stop: a finite-but-exploding CE (e.g. a numerically wrong attention
            # grad walking the writer to uniform output, ce -> ln(vocab)) is just as fatal as a NaN
            # and the per-leaf grad guard does NOT catch it. If CE blows past 2x its best AND clears
            # an absolute margin (so tiny-loss noise can't trip it), halt NOW and preserve the Hub
            # checkpoint -- this would have caught v18 at step ~720 instead of burning 600 steps.
            ce_cur = m.get("ce", m["loss"])
            ce_best = min(ce_best, ce_cur)
            if step >= start + 40 and ce_cur > 2.0 * ce_best and ce_cur > ce_best + 1.0:
                print(f"[abort] step {step}: CE diverged ({ce_cur:.4f} vs best {ce_best:.4f}) -> stopping "
                      f"to save compute (NOT uploading; last good checkpoint preserved)")
                nan_abort = True
                break
            aux = (f"cos {m['cos']:.4f} pw {m.get('pw_ce', 0.0):.4f}"   # latent: planner cos + MTP draft CE
                   + (f" mtp {float(m['mtp_ce']):.4f}" if float(m.get('mtp_ce', 0.0)) > 0 else "")
                   if "cos" in m else f"mtp {m['mtp_ce']:.4f}")          # AR transformer MTP metric
            print(f"step {step:6d} | loss {m['loss']:.4f} ce {m['ce']:.4f} {aux} "
                  f"| gnorm {m['grad_norm']:.3f} | lr {float(lr_sched(step)):.2e} "
                  f"| {tps:,.0f} tok/s | MFU {mfu*100:.1f}%")
            history.append((step, m))
            t0, toks = time.time(), 0

        if run.eval_every and val_sets and (step % run.eval_every == 0 or step == run.max_steps - 1):
            vres = evaluate(params)
            if latent:                                   # dual: CE with the planner vs with it ablated
                line = " | ".join(f"{k} CE {cw:.4f} (no-plan {co:.4f} Δ{co-cw:+.3f})"
                                  for k, (cw, co) in vres.items())
            else:
                line = " | ".join(f"{k} CE {v:.4f}" for k, v in vres.items())
            print("[val] step %6d | " % step + line)
            if sample_text is not None:                  # two generations: with vs without the backbone
                for pr in PROMPTS:
                    try:
                        cw = sample_text(params, pr).replace("\n", " ")
                        co = sample_text(params, pr, ablate_plan=True).replace("\n", " ")
                        print(f"[gen]        step {step:6d} | {pr!r} -> {cw[:180]!r}")
                        print(f"[gen-noplan] step {step:6d} | {pr!r} -> {co[:180]!r}")
                    except Exception as e:
                        print(f"[gen] step {step:6d} | sample failed: {type(e).__name__}: {str(e)[:120]}")
            if hasattr(loader, "position"):              # data-stream health: should GROW each eval
                print(f"[data] step {step:6d} | stream pos {loader.position()} | mix {loader.ratios()}")
            ms = jax.local_devices()[0].memory_stats() or {}
            if ms.get("peak_bytes_in_use"):              # HBM headroom: answers fit questions from the log
                print(f"[mem]  step {step:6d} | peak {ms['peak_bytes_in_use']/2**30:.2f}G "
                      f"/ limit {ms.get('bytes_limit', 0)/2**30:.2f}G")
            hist_v = {}
            for k, v in vres.items():
                if latent:
                    hist_v[f"val/{k}"], hist_v[f"val/{k}_noplan"] = v
                else:
                    hist_v[f"val/{k}"] = v
            history.append((step, hist_v))

        def _ckpt_tree():
            if not tcfg.host_master:
                return params
            _drain()                                     # masters current + quiescent; copy so the
            return jax.tree.map(np.copy, masters)        # async orbax write can't race the worker
        if step % run.ckpt_every == 0:                   # mirror the manager's own save interval
            mgr.save(step, args=ocp.args.StandardSave({"params": _ckpt_tree(),
                                                       "opt_state": opt_state}))
        final_step = step
        if run.max_train_seconds and (time.time() - session_start) >= run.max_train_seconds:
            print(f"[budget] reached {run.max_train_seconds/3600:.2f}h this session at step {step} "
                  f"-> stopping to checkpoint + upload")
            break

    if tcfg.host_master:
        _drain()
    ck = masters if tcfg.host_master else params
    if nan_abort:                                        # broken weights: keep the Hub checkpoint intact
        print("[abort] skipping checkpoint save + upload (preserving the last good checkpoint)")
        return ck, opt_state, start, history
    if mgr.latest_step() != final_step:                  # avoid re-saving an already-checkpointed step
        mgr.save(final_step, args=ocp.args.StandardSave({"params": ck, "opt_state": opt_state}), force=True)
    mgr.wait_until_finished()
    if run.hf_repo and run.upload_at_end:
        _finalize_and_upload(run, mcfg, ck, loader, final_step, time.time() - session_start)
    return ck, opt_state, start, history


def _finalize_and_upload(run: RunConfig, mcfg: ModelConfig, params, loader, step: int, elapsed_s: float):
    """Write bf16 inference weights + resume.json into out_dir, then push the whole
    checkpoint (Orbax master+optimizer, bf16, resume state, val cache) to the HF repo."""
    from . import hfsync
    out_dir = os.path.abspath(run.out_dir)
    try:
        hfsync.save_bf16_safetensors(params, os.path.join(out_dir, "model.safetensors"))
        hfsync.write_resume(out_dir, {
            "step": int(step),
            "arch": run.arch,                            # architecture tag (resume-compatibility guard)
            "writer": _writer_tag(mcfg),                 # STRUCTURE tag: the launcher path-merges (reinit)
            #            when the repo's tag != the code's structure, then resumes CLEAN once the new tag
            #            is written.  v30="ar"; v31 (full-cond fusion + parallel MTP head)="ar+fc+mtp".
            "loader": loader.state() if hasattr(loader, "state") else {},
            "ratios": loader.ratios() if hasattr(loader, "ratios") else {},
            "elapsed_session_s": float(elapsed_s),
            "seed": run.seed,
        })
        hfsync.upload(run.hf_repo, out_dir, msg=f"phase1 step {step}")
        hfsync.purge_old_steps(run.hf_repo, keep=2)   # keep latest 2 step folders -> bounded repo + small resume download
    except Exception as e:
        print(f"[hfsync] finalize/upload FAILED ({type(e).__name__}: {str(e)[:160]}) -- "
              f"local checkpoint in {out_dir} is intact")


def _writer_tag(mcfg: ModelConfig):
    """resume.json structure tag for the latent arch.  The launcher reinits (path-merges the
    bf16 snapshot, fresh optimizer) whenever the repo's tag differs from the running code's
    structure, then resumes clean.  Encodes the param-tree-changing flags so a structure switch
    (e.g. adding the full-cond fusion or the parallel MTP head) is detected even though the base
    is still the AR writer."""
    if mcfg.arch != "latent_diffusion":
        return None
    tag = "ar"
    if mcfg.use_full_cond:
        tag += "+fc"
    if mcfg.use_parallel_writer:
        tag += "+mtp"
    return tag