training.py

# --- nemo_automodel/recipes/llm/prism_entry.py (trainer) ---
"""PRISM AutoModel 2.0 entry: self-contained build_model + train.

Modern decoder LM (RMSNorm, RoPE, SwiGLU, GQA, QK-norm, value-residual, U-net
skips, tied-free head, logit soft-cap) trained from random init with Muon +
AdamW (per-group LRs), a WSD schedule, and FlexAttention document masking.
The harness resolves this module and calls build_model(ctx) then train(model, ctx).
"""

import math
import os
import sys
import time

os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")

import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2TokenizerFast

try:
    from torch.nn.attention.flex_attention import flex_attention as _flex
    _flex = torch.compile(_flex)
except Exception:
    _flex = None

try:
    from torch.nn.attention.flex_attention import create_block_mask as _create_block_mask
except Exception:
    _create_block_mask = None

try:
    import prism_telemetry
except ImportError:
    class _TelemetryFallback:
        @staticmethod
        def report(**_kwargs):
            return None

        @staticmethod
        def finish_evaluation():
            return None

    prism_telemetry = _TelemetryFallback()


VOCAB = 8192  # default; real value comes from the shipped tokenizer via ctx["vocab_size"]
BLOCK = 1024
D_MODEL = 768
N_LAYER = 15
N_HEAD = 12
N_KV_HEAD = 4
FFN_MULT = 4
DROPOUT = 0.0
LOGIT_SOFTCAP = 15.0


def _ctx_get(ctx, key, default):
    if hasattr(ctx, "get"):
        try:
            v = ctx.get(key, default)
            return default if v is None else v
        except Exception:
            return default
    return getattr(ctx, key, default)


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

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


def _rope(head_dim: int, seq: int, base: float = 10000.0):
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    t = torch.arange(seq, dtype=torch.float32)
    f = torch.outer(t, inv)
    return torch.cos(f), torch.sin(f)


def _apply_rope(x, cos, sin):
    # x: (b, h, t, hd). Interleaved even/odd rotation.
    xe, xo = x[..., 0::2], x[..., 1::2]
    cos = cos[None, None, :, :].to(x.dtype)
    sin = sin[None, None, :, :].to(x.dtype)
    re = xe * cos - xo * sin
    ro = xe * sin + xo * cos
    return torch.stack([re, ro], dim=-1).flatten(-2)


class Attention(nn.Module):
    def __init__(self, d: int, n_head: int, n_kv: int):
        super().__init__()
        self.n_head, self.n_kv = n_head, n_kv
        self.hd = d // n_head
        self.groups = n_head // n_kv
        self.wq = nn.Linear(d, n_head * self.hd, bias=False)
        self.wk = nn.Linear(d, n_kv * self.hd, bias=False)
        self.wv = nn.Linear(d, n_kv * self.hd, bias=False)
        self.wo = nn.Linear(n_head * self.hd, d, bias=False)
        # QK-norm: normalize queries/keys per head for stable attention logits.
        self.q_norm = RMSNorm(self.hd)
        self.k_norm = RMSNorm(self.hd)
        # Value-residual: mix layer-0 values into this layer (sigmoid(lam) in (0,1),
        # init 0.5). Speeds convergence by giving every block direct access to
        # first-layer value features (ResFormer / value-residual learning).
        self.lam = nn.Parameter(torch.zeros(1))

    def forward(self, x, cos, sin, v0=None, block_mask=None):
        b, t, _ = x.shape
        q = self.q_norm(self.wq(x).view(b, t, self.n_head, self.hd)).transpose(1, 2)
        k = self.k_norm(self.wk(x).view(b, t, self.n_kv, self.hd)).transpose(1, 2)
        v = self.wv(x).view(b, t, self.n_kv, self.hd)
        if v0 is None:
            v0_out = v
        else:
            lam = torch.sigmoid(self.lam)
            v = lam * v + (1.0 - lam) * v0
            v0_out = v0
        v = v.transpose(1, 2)
        q = _apply_rope(q, cos, sin)
        k = _apply_rope(k, cos, sin)
        k = k.to(q.dtype)
        v = v.to(q.dtype)
        if block_mask is not None and _flex is not None:
            # Fast document-masked attention (GQA handled natively).
            o = _flex(q, k, v, block_mask=block_mask, enable_gqa=True)
        else:
            if self.groups > 1:
                k = k.repeat_interleave(self.groups, dim=1)
                v = v.repeat_interleave(self.groups, dim=1)
            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        o = o.transpose(1, 2).contiguous().view(b, t, self.n_head * self.hd)
        return self.wo(o), v0_out


class SwiGLU(nn.Module):
    def __init__(self, d: int, hidden: int):
        super().__init__()
        self.w_gate = nn.Linear(d, hidden, bias=False)
        self.w_up = nn.Linear(d, hidden, bias=False)
        self.w_down = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, d, n_head, n_kv, ffn, dropout=0.0):
        super().__init__()
        self.n1 = RMSNorm(d)
        self.attn = Attention(d, n_head, n_kv)
        self.n2 = RMSNorm(d)
        self.mlp = SwiGLU(d, ffn)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, cos, sin, v0=None, block_mask=None):
        a, v0_out = self.attn(self.n1(x), cos, sin, v0, block_mask)
        x = x + self.drop(a)
        return x + self.drop(self.mlp(self.n2(x))), v0_out


class DecoderLM(nn.Module):
    def __init__(self, vocab, d, n_layer, n_head, n_kv, block, ffn_mult):
        super().__init__()
        self.block = block
        self.tok_emb = nn.Embedding(vocab, d)
        ffn = ((int(d * ffn_mult * 2 / 3) + 63) // 64) * 64
        self.blocks = nn.ModuleList(Block(d, n_head, n_kv, ffn, DROPOUT) for _ in range(n_layer))
        self.emb_drop = nn.Dropout(DROPOUT)
        self.skip_weights = nn.Parameter(torch.zeros(n_layer // 2))
        self.ln_f = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        # untied head (zero-init), Adam high-LR

        cos, sin = _rope(d // n_head, block)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(self._init)
        # Scaled residual init for depth stability.
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w_down.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layer))
        nn.init.zeros_(self.head.weight)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, ids, block_mask=None):
        t = min(ids.shape[1], self.block)
        ids = ids[:, -t:]
        if self.rope_cos.device != ids.device:
            self.rope_cos = self.rope_cos.to(ids.device)
            self.rope_sin = self.rope_sin.to(ids.device)
        cos, sin = self.rope_cos[:t], self.rope_sin[:t]
        x = self.emb_drop(self.tok_emb(ids))
        v0 = None
        n = len(self.blocks)
        n_enc = n // 2
        skips = []
        for i, blk in enumerate(self.blocks):
            if i >= n - n_enc:
                x = x + self.skip_weights[i - (n - n_enc)] * skips.pop()
            x, v0 = blk(x, cos, sin, v0, block_mask)
            if i < n_enc:
                skips.append(x)
        x = self.ln_f(x)
        logits = self.head(x)
        return LOGIT_SOFTCAP * torch.tanh(logits / LOGIT_SOFTCAP)


def _resolve_tokenizer(ctx):
    """Prefer the harness-resolved+validated tokenizer (ctx['tokenizer']); else
    load the shipped in-domain tokenizer from the local source tree. Never hits
    the network (the miner child runs under unshare --net)."""
    tok = _ctx_get(ctx, "tokenizer", None)
    if tok is not None:
        return tok
    import os as _os
    here = _os.path.dirname(_os.path.abspath(__file__))
    cands = [
        _ctx_get(ctx, "tokenizer_dir", None),
        _os.path.join(here, "tokenizer"),
        _os.path.join(here, "tok8k"),
        _os.path.join(_os.getcwd(), "tokenizer"),
        _os.path.join(_os.getcwd(), "submission", "tokenizer"),
    ]
    for d in cands:
        if d and _os.path.isfile(_os.path.join(d, "tokenizer.json")):
            try:
                from transformers import AutoTokenizer
                return AutoTokenizer.from_pretrained(d, local_files_only=True)
            except Exception:
                from transformers import PreTrainedTokenizerFast
                return PreTrainedTokenizerFast(
                    tokenizer_file=_os.path.join(d, "tokenizer.json"),
                    eos_token="<|eos|>", pad_token="<|eos|>")
    # Last-resort fallback: the pinned GPT-2 default (keeps local smoke tests
    # runnable even without shipped files; never reached in-pod).
    from transformers import GPT2TokenizerFast
    return GPT2TokenizerFast.from_pretrained("gpt2")


def build_model(ctx):
    """Recipe contract entrypoint: return an nn.Module given the context."""
    torch.manual_seed(int(_ctx_get(ctx, "seed", 0)))
    d = int(_ctx_get(ctx, "d_model", D_MODEL))
    n_layer = int(_ctx_get(ctx, "n_layer", N_LAYER))
    n_head = int(_ctx_get(ctx, "n_head", N_HEAD))
    n_kv = int(_ctx_get(ctx, "n_kv_head", N_KV_HEAD))
    block = int(_ctx_get(ctx, "block", BLOCK))
    # Embedding/head are sized from the resolved (shipped) tokenizer, not a
    # hardcoded constant — the harness passes vocab_size for the tokenizer it
    # resolved from submission/tokenizer/. Fail loud on mismatch would be ideal,
    # but the harness guarantees this equals len(shipped tokenizer).
    vocab = int(_ctx_get(ctx, "vocab_size", VOCAB))
    return DecoderLM(vocab, d, n_layer, n_head, n_kv, block, FFN_MULT)


HARD_MAX_STEPS = 20000
HARD_MAX_SECONDS = 21600
MUON_LR = 0.036       # peak LR for the transformer weight matrices (Muon)
# Embedding LR scales inversely with per-token gradient density. At 50k vocab
# (sparse: ~4 hits/token/step) 0.3 was optimal; at 4096 vocab embeddings update
# ~12x denser, so hold the effective per-token update constant -> ~0.03.
EMB_LR = 0.05         # input embedding LR (tuned for 8k in-domain vocab)
HEAD_LR = 3.6e-3      # peak LR for the untied output head (AdamW ~ 0.1/sqrt(d))
SCALAR_LR = 0.015     # peak LR for 1-D params: norms, skip/value-residual scalars
MIN_LR_RATIO = 0.01
WARMUP_STEPS = 300
WEIGHT_DECAY = 0.1
GRAD_CLIP = 1.0
REPORT_EVERY = 20
EOT_ID = 50256  # gpt2 <|endoftext|>
# Rows of unique data to read (memory-safe as uint16); ~420M tokens, reused a few epochs.
MAX_DATA_ROWS = 700000
# Large effective batch via gradient accumulation (the run is step-capped, so
# tokens/step is the lever); target a few epochs over the shard, not more.
EFFECTIVE_TOKENS_PER_STEP = 262144
MICRO_TOKENS = 16384
TARGET_EPOCHS = 2.0


def _ctx_get(ctx, key, default):
    if hasattr(ctx, "get"):
        try:
            v = ctx.get(key, default)
            return default if v is None else v
        except Exception:
            return default
    return getattr(ctx, key, default)


def _budget(ctx):
    max_steps, max_seconds = HARD_MAX_STEPS, HARD_MAX_SECONDS
    b = getattr(ctx, "budget", None)
    if callable(b):
        try:
            bb = b()
            max_steps = int(getattr(bb, "max_steps", max_steps))
            max_seconds = float(getattr(bb, "max_seconds", max_seconds))
        except Exception:
            pass
    max_steps = int(_ctx_get(ctx, "max_train_steps", max_steps))
    max_seconds = float(_ctx_get(ctx, "max_seconds", max_seconds))
    return min(max_steps, HARD_MAX_STEPS), min(max_seconds, HARD_MAX_SECONDS)


def _packed_tokens(ctx, tok, block, want_tokens):
    # Concatenate tokenized docs (EOT-separated) into one uint16 stream, then reshape
    # into (n, block) windows. Tokens are kept as numpy uint16 (gpt2 vocab < 65536) so
    # ~hundreds of millions of tokens fit in memory instead of a huge Python int list.
    import numpy as np

    path = _ctx_get(ctx, "dataset_path", None)
    tr = int(_ctx_get(ctx, "train_rows", 2048))
    vr = int(_ctx_get(ctx, "val_rows", 256))
    # Cache the packed token stream so a re-run in the same pod skips re-tokenizing.
    cache = f"/tmp/prism_pack_{MAX_DATA_ROWS}_{tr}_{vr}.npy"
    if os.path.exists(cache):
        stream = np.load(cache)
    else:
        # Only materialize the rows we use. The harness scores texts[tr:tr+vr]; the
        # pinned shard has no <100-char rows, so raw row order matches its filtered
        # order and slicing raw keeps train disjoint from the scored val cut.
        col = pq.read_table(path, columns=["text"]).column("text")
        head = col.slice(0, tr).to_pylist()
        body = col.slice(tr + vr, MAX_DATA_ROWS).to_pylist()
        texts = [t for t in head + body if isinstance(t, str) and len(t) >= 100]
        eot = np.array([EOT_ID], dtype=np.uint16)
        parts = []
        total = 0
        for start in range(0, len(texts), 1024):
            for enc in tok(texts[start : start + 1024], add_special_tokens=False)["input_ids"]:
                parts.append(np.asarray(enc, dtype=np.uint16))
                parts.append(eot)
                total += len(enc) + 1
            if total >= want_tokens:
                break
        stream = np.concatenate(parts) if parts else np.zeros(block + 1, dtype=np.uint16)
        try:
            np.save(cache, stream)
        except OSError:
            pass
    n = max(1, (len(stream) - 1) // block)
    stream = stream[: n * block + 1]
    x = torch.from_numpy(stream[:-1].astype(np.int32)).view(n, block)
    y = torch.from_numpy(stream[1:].astype(np.int32)).view(n, block)
    return x, y


def _batch_size(device, block):
    if not str(device).startswith("cuda"):
        return 4
    # ~32k tokens/step; back off if VRAM is tight (handled by the OOM guard).
    return max(1, min(48, 32768 // max(block, 1)))


def _doc_block_mask(xb):
    # Block-sparse mask: attend only within the same document (EOT-split) + causal.
    if _create_block_mask is None:
        return None
    eot = (xb == EOT_ID)
    doc_id = torch.cumsum(eot.int(), dim=1) - eot.int()
    def mm(b, h, q_idx, kv_idx):
        return (q_idx >= kv_idx) & (doc_id[b, q_idx] == doc_id[b, kv_idx])
    return _create_block_mask(mm, xb.shape[0], None, xb.shape[1], xb.shape[1], device=xb.device)


def _sched(elapsed, budget, warmup_frac=0.03, decay_frac=0.20):
    # Warmup -> stable (hold peak LR) -> cosine decay over the last decay_frac.
    # WSD keeps LR high through most of the run (faster convergence when val is
    # still descending at the cap) and only anneals near the end.
    warm = warmup_frac * budget
    if elapsed < warm:
        return elapsed / max(warm, 1e-9)
    decay_start = (1.0 - decay_frac) * budget
    if elapsed < decay_start:
        return 1.0
    progress = min(1.0, (elapsed - decay_start) / max(1e-9, budget - decay_start))
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return MIN_LR_RATIO + (1.0 - MIN_LR_RATIO) * cosine


def _is_no_decay(name, p):
    return p.dim() < 2 or "norm" in name.lower() or "emb" in name.lower()


def _newtonschulz5(G, steps=5):
    # Orthogonalize the gradient matrix via a quintic Newton-Schulz iteration
    # (Muon; Bernstein/Jordan). Runs in bf16; coefficients tuned for LM training.
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    transposed = X.shape[0] > X.shape[1]
    if transposed:
        X = X.T
    X = X / (X.norm() + 1e-7)
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    return X.T if transposed else X


class Muon(torch.optim.Optimizer):
    """Muon: momentum + Newton-Schulz orthogonalization for 2D weight matrices.
    Trains transformer matrices markedly faster than AdamW at the same compute."""

    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
        super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps))

    @torch.no_grad()
    def step(self):
        for grp in self.param_groups:
            lr, mom, nest, ns = grp["lr"], grp["momentum"], grp["nesterov"], grp["ns_steps"]
            for p in grp["params"]:
                if p.grad is None:
                    continue
                st = self.state[p]
                if "buf" not in st:
                    st["buf"] = torch.zeros_like(p.grad)
                buf = st["buf"]
                buf.mul_(mom).add_(p.grad)
                g = p.grad.add(buf, alpha=mom) if nest else buf
                g = _newtonschulz5(g, steps=ns)
                # Scale so the RMS update size is ~lr regardless of matrix shape.
                scale = (max(1.0, p.size(0) / p.size(1)) ** 0.5)
                p.add_(g.to(p.dtype), alpha=-lr * scale)


def train(model, ctx):
    """Recipe contract entrypoint: trains in place, returns a metrics dict."""
    device = _ctx_get(ctx, "device", "cpu")
    torch.manual_seed(int(_ctx_get(ctx, "seed", 0)))
    guard = _ctx_get(ctx, "guard", lambda: None)
    if str(device).startswith("cuda"):
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    max_steps, max_seconds = _budget(ctx)
    warmup = min(WARMUP_STEPS, max(1, max_steps // 20))

    tok = _resolve_tokenizer(ctx)
    # The frozen-val bpb is scored with THIS tokenizer; the doc separator in the
    # packed training stream + the doc-mask must use its real eos id.
    global EOT_ID
    _eos = getattr(tok, "eos_token_id", None)
    EOT_ID = int(_eos) if _eos is not None else (len(tok) - 1)
    block = getattr(model, "block", 512)
    micro = max(1, int(_ctx_get(ctx, "batch_size", MICRO_TOKENS // max(block, 1))))
    accum = max(1, EFFECTIVE_TOKENS_PER_STEP // (micro * block))

    x, y = _packed_tokens(ctx, tok, block, want_tokens=1_200_000_000)
    # Exact replica of the harness val cut (texts[train_rows:train_rows+val_rows])
    # for a live bpb estimate; these rows are excluded from training data.
    val_ids = []
    try:
        _vp = _ctx_get(ctx, "dataset_path", None)
        _tr = int(_ctx_get(ctx, "train_rows", 2048)); _vr = int(_ctx_get(ctx, "val_rows", 256))
        _vcol = pq.read_table(_vp, columns=["text"]).column("text").slice(_tr, _vr).to_pylist()
        _vtexts = [t for t in _vcol if isinstance(t, str) and len(t) >= 100]
        for _t in _vtexts:
            _ids = tok(_t, add_special_tokens=False)["input_ids"]
            if len(_ids) >= 2:
                val_ids.append(torch.tensor(_ids, dtype=torch.long))
    except Exception:
        val_ids = []
    n_windows = x.shape[0]
    # Train for the whole compute budget (time governs the schedule); reshuffled
    # windows give ~as many epochs as fit, and the LR decays to its floor by the cap.
    time_budget = max_seconds * 0.94

    # Muon for the 2D transformer weight matrices; AdamW for embeddings + norms.
    muon_params, adam_emb, adam_head, adam_scalar = [], [], [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if "tok_emb" in n:
            adam_emb.append(p)
        elif "head" in n:
            adam_head.append(p)
        elif p.ndim >= 2:
            muon_params.append(p)
        else:
            adam_scalar.append(p)
    opt_muon = Muon(muon_params, lr=MUON_LR, momentum=0.95)
    opt_adam = torch.optim.AdamW(
        [{"params": adam_emb, "lr": EMB_LR},
         {"params": adam_head, "lr": HEAD_LR},
         {"params": adam_scalar, "lr": SCALAR_LR}],
        betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0,
        fused=str(device).startswith("cuda"),
    )
    optimizers = [opt_muon, opt_adam]
    use_bf16 = str(device).startswith("cuda") and torch.cuda.is_bf16_supported()
    model.to(device)
    model.train()
    EMA_DECAY = 0.999
    ema_params = {n: p.detach().clone().float() for n, p in model.named_parameters() if p.requires_grad}
    # Compile the TRAINING forward for throughput (fixed shape). The harness scores
    # val on the original `model` (variable lengths), so it never recompiles.
    fwd = model
    if str(device).startswith("cuda"):
        try:
            fwd = torch.compile(model)
        except Exception:
            fwd = model

    g = torch.Generator().manual_seed(int(_ctx_get(ctx, "seed", 0)))
    order = torch.randperm(n_windows, generator=g)
    t0 = time.time()
    last = 0.0
    ema = 0.0
    cursor = 0
    step = 0

    def _next_micro(n):
        nonlocal cursor, order
        if cursor + n > n_windows:
            order = torch.randperm(n_windows, generator=g)
            cursor = 0
        idx = order[cursor : cursor + n]
        cursor += n
        return (x[idx].to(device, non_blocking=True).long(),
                y[idx].to(device, non_blocking=True).long())

    while step < max_steps:
        guard()
        elapsed = time.time() - t0
        if elapsed >= time_budget:
            break
        mult = _sched(elapsed, time_budget)
        lr = MUON_LR * mult
        opt_muon.param_groups[0]["lr"] = MUON_LR * mult
        opt_adam.param_groups[0]["lr"] = EMB_LR * mult
        opt_adam.param_groups[1]["lr"] = HEAD_LR * mult
        opt_adam.param_groups[2]["lr"] = SCALAR_LR * mult
        for o in optimizers:
            o.zero_grad(set_to_none=True)
        try:
            for a in range(accum):
                xb, yb = _next_micro(micro)
                bm = _doc_block_mask(xb)
                if use_bf16:
                    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                        logits = fwd(xb, block_mask=bm)
                        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), yb.reshape(-1))
                else:
                    logits = fwd(xb, block_mask=bm)
                    loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), yb.reshape(-1))
                (loss / accum).backward()
            grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP))
            for o in optimizers:
                o.step()
            with torch.no_grad():
                for _n, _p in model.named_parameters():
                    if _n in ema_params:
                        ema_params[_n].mul_(EMA_DECAY).add_(_p.detach().float(), alpha=1.0 - EMA_DECAY)
        except RuntimeError as exc:
            if "out of memory" in str(exc).lower() and micro > 1:
                micro = max(1, micro // 2)
                accum = max(1, EFFECTIVE_TOKENS_PER_STEP // (micro * block))
                if str(device).startswith("cuda"):
                    torch.cuda.empty_cache()
                continue
            raise

        last = float(loss.item())
        ema = last if step == 0 else (0.98 * ema + 0.02 * last)
        step += 1
        if step == 1 or step % REPORT_EVERY == 0:
            prism_telemetry.report(
                loss=last, step=step, grad_norm=grad_norm,
                layer_stats={"model": {"grad_norm": grad_norm}},
            )
        if step == 1 or step % 100 == 0:
            print(f"[train] step {step}/{max_steps} loss {last:.3f} ema {ema:.3f} lr {lr:.2e} "
                  f"gnorm {grad_norm:.2f} t {time.time()-t0:.0f}s", file=sys.stderr, flush=True)
        if val_ids and (step == 1 or step % 1000 == 0):
            try:
                model.eval()
                vloss = []
                with torch.no_grad():
                    for _ids in val_ids:
                        _i = _ids.to(device).unsqueeze(0)
                        if use_bf16:
                            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                                _lg = model(_i[:, :-1])
                                _tg = _i[:, 1:][:, -_lg.shape[1]:]
                                _l = F.cross_entropy(_lg.reshape(-1, _lg.shape[-1]).float(), _tg.reshape(-1))
                        else:
                            _lg = model(_i[:, :-1])
                            _tg = _i[:, 1:][:, -_lg.shape[1]:]
                            _l = F.cross_entropy(_lg.reshape(-1, _lg.shape[-1]), _tg.reshape(-1))
                        vloss.append(float(_l))
                model.train()
                _ce = sum(vloss) / max(1, len(vloss))
                print(f"[val] step {step} val_ce {_ce:.4f} bpb {_ce/0.6931471805599453:.4f} "
                      f"(top 4.2881) t {time.time()-t0:.0f}s", file=sys.stderr, flush=True)
            except Exception:
                model.train()

    with torch.no_grad():
        for _n, _p in model.named_parameters():
            if _n in ema_params:
                _p.copy_(ema_params[_n].to(_p.dtype))
    prism_telemetry.report(loss=last, step=step, grad_norm=0.0)
    prism_telemetry.finish_evaluation()
    return {"train_loss": last, "train_steps": step, "train_seconds": time.time() - t0}