Copy files from models/jayasriii/390M_5e4

train_390M_5e4.py

@@ -0,0 +1,1847 @@
+#!/usr/bin/env python3
+# =============================================================================
+# train_muon.py — the MUON arm of train.py. Same model, same data order, same
+# schedule shape, same seed: the ONLY thing that changes is how a gradient
+# becomes a weight update.
+#
+# Fork discipline — everything outside section 5 is byte-identical to train.py,
+# so `diff train.py train_muon.py` is exactly the experiment being run. If you
+# edit a kernel or the model in one file, port it to the other or the comparison
+# stops meaning anything.
+#
+# The optimizer itself is torch.optim.Muon — stock, not reimplemented. This file
+# supplies only the two things it does not do for you: the ROUTING (which
+# parameters it may have) and the fp32 MASTER WEIGHTS (it updates the master,
+# not the bf16 param). Requires a torch new enough to ship it; main() checks and
+# says so.
+#
+# What Muon does: AdamW rescales each weight ENTRY by its own gradient history.
+# Muon rescales each weight MATRIX by its singular values — it takes the
+# momentum buffer M and replaces it with the nearest semi-orthogonal matrix
+# (all singular values driven to ~1) via a 5-step Newton-Schulz iteration, then
+# steps along that. Directions the gradient has barely explored get the same
+# step size as the dominant ones, which is where the reported speedup comes
+# from. Only 2D hidden weights can be orthogonalized, so this is a HYBRID:
+#
+#     Muon   : every 2D Linear weight inside the blocks   (75.5M params, 66%)
+#     AdamW  : the token embedding / lm_head              (38.6M params)
+#              + RMSNorm weights + QK-norm gammas (1D)    (22.3k params)
+#
+# The embedding is deliberately NOT given to Muon: its rows are per-token, not a
+# linear map between two feature spaces, and orthogonalizing across the vocab
+# axis mixes unrelated tokens. That split is the standard Muon recipe, and it
+# keeps the AdamW arm's exact update on 34% of the parameters — so a difference
+# in the loss curve is attributable to the hidden matrices alone.
+#
+# Everything train_ddp.py / train_zero1.py import from blu_common / blu_data /
+# blu_kernels / blu_model is inlined here, in that order, so this file runs on
+# its own:
+#     section 1  fused QK-norm + RoPE (Triton)  <- blu_kernels.py
+#     section 2  config + WSD schedule          <- blu_common.py
+#     section 3  data loader                    <- blu_data.py
+#     section 4  model                          <- blu_model.py
+#     section 5  optimizer / clipping / checkpointing
+#     section 6  training
+#
+# Architecture:
+#   token embedding
+#   -> N x { RMSNorm -> [fused QK-norm + RoPE] -> causal GQA (flash) -> Wo -> +x
+#            RMSNorm -> SwiGLU FFN (gate_up -> silu(gate)*up -> down)  -> +x }
+#   -> final RMSNorm
+#   -> lm_head (weight-tied with the token embedding)
+#   -> cross-entropy
+#   -> backward -> AdamW (decoupled weight decay) on a WSD (warmup / stable /
+#      cosine-decay) schedule
+#
+# Kept from the AdamW arm, unchanged, because they are confounds if they move:
+#   * fp32 master weights for every non-fp32 param — Muon updates the MASTER,
+#     reading the same fp32 landing buffer AdamW reads, so both optimizers see
+#     bit-identical gradients;
+#   * decoupled weight decay 0.1 on dim>=2 params, applied to the master;
+#   * ONE global-norm gradient clip over every parameter, BEFORE the optimizer.
+#     (Muon's orthogonalization makes the update scale-invariant, so clipping is
+#     nearly a no-op for its group — but it still gates the AdamW group, and
+#     removing it would change two things at once.)
+#
+# Memory, at the 114M config this repo actually runs (printed at startup):
+#
+#     parameters (bf16 Linear + fp32 embed/norms)  0.28 GiB
+#     gradients (same dtypes as the params)        0.28 GiB
+#     fp32 master weights (the bf16 Linears)       0.28 GiB
+#     fp32 gradient landing buffers                0.28 GiB
+#     Muon momentum   (75.5M params, fp32 x1)      0.28 GiB
+#     AdamW m + v     (38.7M params, fp32 x2)      0.29 GiB
+#     -------------------------------------------------------
+#     static total                                 1.70 GiB   (AdamW arm: 1.98)
+#
+# Muon carries ONE state tensor per param where Adam carries two, so the saving
+# scales with how much of the model is hidden matrices: 0.28 GiB here (14% of
+# static), because this config spends a third of its parameters on a 50304-row
+# embedding that stays on AdamW. At the 1.73B config in the original header the
+# same split saves ~5.5 GiB.
+#
+# Activations dominate either way (63.96 GiB measured at B=16 x T=4096), so the
+# OOM escape hatches are unchanged: --B 8, --loss-chunks 4, --activation-checkpoint.
+#
+# Fidelity notes (why a run here is statistically but not bit-wise identical):
+#   * init draws come from torch's Philox, not the C++ single advancing
+#     std::mt19937, so the two runs start from different samples of the same
+#     distributions;
+#   * QK-norm reads the projection output in bf16 here (autocast) where the C++
+#     casts it up to fp32 first; the reduction and the rotation are fp32 in both.
+#
+# Determinism: two runs of the same launch line on the same GPU produce
+# bit-identical losses. Seeding is only half of that — the default CUDA kernels
+# for the embedding backward and the flash/mem-efficient attention backward
+# accumulate partial sums with atomicAdd, so the order they land in (and with it
+# the fp32 rounding) follows the scheduler, not the program. --determinism
+# strict (the default) forces the ordered variants of those kernels; it costs a
+# few percent of step time, which --determinism off buys back.
+#
+# Env overrides (the bluscript_single.cpp names, so an existing launch line keeps
+# working; the equivalent CLI flag wins if both are given):
+#   BLU_DEVICE=<idx>          CUDA device to train on            (default 0)
+#   BLU_MAX_STEPS=<n>         short benchmark run (rescales WSD) (default off)
+#   BLU_B=<n>                 micro-batch size                   (default 4)
+#   BLU_GLOBAL_BATCH=<tokens> tokens per optimizer step          (default 524288)
+#   BLU_DIAG=1                per-micro-batch loss + depth probes
+#
+# Multi-GPU: DDP, one process per GPU, via torchrun. RANK / LOCAL_RANK /
+# WORLD_SIZE (set by torchrun, not this file) switch the bootstrap in main();
+# their absence is the single-process path above, unchanged. --device /
+# BLU_DEVICE are ignored under torchrun — LOCAL_RANK picks each rank's GPU.
+# --global-batch is still the TOTAL tokens per optimizer step across every
+# rank (grad_accum = global_batch / (B*T*world_size)), so the same
+# --global-batch reproduces the same schedule at any world size. Data is
+# rank-sharded (DataLoaderLite reads a disjoint, rank-offset slice of the
+# shards); gradients and the reported train/val loss are averaged across
+# ranks; only rank 0 writes the CSV and checkpoints. Same launch line as
+# single-GPU otherwise:
+#   torchrun --standalone --nproc_per_node=8 train_muon.py
+#
+# Learning rate — the one knob that decides whether the comparison is fair.
+#
+# Muon's update is spectrally normalized, so its natural scale has nothing to do
+# with AdamW's. torch.optim.Muon ships both published conventions behind its
+# `adjust_lr_fn`; --muon-adjust-lr passes the string straight through:
+#
+#   match_rms_adamw      Moonshot's: gamma <- 0.2 * gamma * sqrt(max(A,B)). The
+#     (DEFAULT HERE)     Newton-Schulz output has singular values ~1, so its
+#                        entries have RMS 1/sqrt(max(A,B)); this lifts them to
+#                        RMS 0.2, which is what a converged AdamW update looks
+#                        like. The consequence: --muon-lr is on the SAME SCALE
+#                        as --max-lr, and the default (2e-4, equal to the AdamW
+#                        arm's max_lr) gives both arms the same step size on the
+#                        first run — only the update DIRECTION differs. That is
+#                        the controlled experiment; run it before tuning.
+#   original             Keller's modded-nanogpt scaling, sqrt(max(1, A/B)).
+#                        Entries stay at RMS 1/sqrt(max(A,B)), so an equivalent
+#                        LR is ~5.5x larger at d_model=768 (0.2*sqrt(768)).
+#                        Use it to reproduce published Muon numbers.
+#
+# NOTE: torch's own default for adjust_lr_fn is `original`. This file defaults
+# to `match_rms_adamw` instead, because an LR-matched first run is the point of
+# having two arms. If you cite a number from here, cite the flag with it.
+#
+# Muon usually tolerates 2-4x more LR than AdamW once it is working, so the
+# honest protocol is: baseline run, matched run, then a short --max-steps 200
+# sweep over --muon-lr before claiming a win.
+#
+# Launch (the two arms; separate CSV + checkpoint dirs so they never collide):
+#   python3 train.py                                  # AdamW baseline
+#   python3 train_muon.py                             # Muon, LR-matched
+#   python3 train_muon.py --muon-lr 6e-4              # after a sweep
+#   python3 train_muon.py --max-steps 200 --muon-lr 4e-4   # short sweep point
+# =============================================================================
+
+import argparse
+import glob
+import math
+import os
+import random
+import re
+import signal
+import time
+from contextlib import nullcontext
+from dataclasses import dataclass
+
+# cuBLAS reads CUBLAS_WORKSPACE_CONFIG once, when it creates its handle on the
+# first GEMM, and torch refuses to run fp32 cuBLAS GEMMs under
+# use_deterministic_algorithms() unless it is one of the two allowed values. It
+# has to be in the environment before torch touches CUDA, hence: before the
+# import, not inside main(). An explicit setting from the launch line wins.
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.checkpoint import checkpoint
+
+try:
+    import triton
+    import triton.language as tl
+    _HAS_TRITON = True
+except ImportError:                                   # eager fallback below
+    _HAS_TRITON = False
+
+
+# =============================================================================
+# 1. Fused QK-norm (RMSNorm over head_dim) + RoPE, in Triton.
+#
+# PyTorch counterpart of the C++ `autograd::fused_rms_rope`. One kernel does,
+# per (batch, head, position) row:
+#
+#     s = rsqrt(mean(x^2) + eps)          # RMS over head_dim
+#     y = x * s * gamma                   # learnable per-head_dim scale
+#     o = rope_neox(y, cos[t], sin[t])    # NeoX / Llama half-split rotation
+#
+# Attention is deliberately NOT fused in here — it stays a separate op
+# (F.scaled_dot_product_attention on the flash backend), mirroring the C++ split
+# between `fused_rms_rope` and `gqa_fused_flash_attn_forward/backward`.
+#
+# Why fuse: the eager formulation (`rms_rope_eager`) materialises an fp32 copy of
+# x, the squared tensor, the normalised tensor, the gamma-scaled tensor and a
+# `cat` of the two rotated halves — roughly 6 extra [B,H,T,D] buffers held for
+# backward. The fused version saves only x (which the projection's backward
+# already holds) plus one fp32 scalar per row.
+#
+# Layout notes
+#   - x is [B,H,T,D] and may be a non-contiguous transpose of a [B,T,H,D] view;
+#     arbitrary (b,h,t) strides are honoured, D must be stride-1.
+#   - D must be a power of two (the C++ kernel requires 64 or 128).
+#   - cos/sin are [T_ctx, D//2] fp32, as produced by `build_rope_cache`.
+#   - the row index is r = (b*H + h)*T + t, so the output is contiguous
+#     [B,H,T,D] in that same order.
+# =============================================================================
+if _HAS_TRITON:
+
+    @triton.jit
+    def _fused_rms_rope_fwd(
+        X, GAMMA, COS, SIN, OUT, RSTD,
+        sx_b, sx_h, sx_t,                   # input strides, in elements
+        T, H, N,                            # seq len, n_heads, total rows
+        eps,
+        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)          # [R]
+        mask_r = rows < N
+
+        # r = (b*H + h)*T + t  ->  recover t (for RoPE) and the (b, h) base.
+        t = rows % T
+        bh = rows // T
+        b = bh // H
+        h = bh % H
+        base = b * sx_b + h * sx_h + t * sx_t                  # [R]
+
+        dh = tl.arange(0, HALF)                               # [HALF]
+        p1 = X + base[:, None] + dh[None, :]
+        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        # RMS over the FULL head_dim (both halves), fp32 accumulation.
+        ss = tl.sum(x1 * x1, 1) + tl.sum(x2 * x2, 1)          # [R]
+        rstd = 1.0 / tl.sqrt(ss / D + eps)                    # [R]
+
+        g1 = tl.load(GAMMA + dh).to(tl.float32)
+        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
+        y1 = x1 * rstd[:, None] * g1[None, :]
+        y2 = x2 * rstd[:, None] * g2[None, :]
+
+        cp = t[:, None] * HALF + dh[None, :]
+        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
+        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)
+
+        # NeoX half-split rotation.
+        o1 = y1 * c - y2 * s
+        o2 = y2 * c + y1 * s
+
+        op = OUT + rows[:, None] * D + dh[None, :]
+        tl.store(op, o1.to(OUT.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(op + HALF, o2.to(OUT.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(RSTD + rows, rstd, mask=mask_r)
+
+    # -------------------------------------------------------------------------
+    # backward
+    #
+    #   RoPE is an orthogonal rotation, so its vjp is the inverse rotation:
+    #       dy1 = do1*cos + do2*sin ;  dy2 = do2*cos - do1*sin
+    #   then, with n = x * rstd (the normalised row),
+    #       dgamma += dy * n
+    #       dn      = dy * gamma
+    #       dx      = rstd * (dn - n * mean(dn * n))
+    # -------------------------------------------------------------------------
+    @triton.jit
+    def _fused_rms_rope_bwd(
+        X, GAMMA, COS, SIN, DOUT, RSTD, DX, DGP,
+        sx_b, sx_h, sx_t,                   # input strides
+        sd_b, sd_h, sd_t,                   # dx strides
+        T, H, N,
+        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
+        mask_r = rows < N
+
+        t = rows % T
+        bh = rows // T
+        b = bh // H
+        h = bh % H
+        base = b * sx_b + h * sx_h + t * sx_t
+
+        dh = tl.arange(0, HALF)
+        p1 = X + base[:, None] + dh[None, :]
+        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        rstd = tl.load(RSTD + rows, mask=mask_r, other=1.0)
+        n1 = x1 * rstd[:, None]
+        n2 = x2 * rstd[:, None]
+
+        dp = DOUT + rows[:, None] * D + dh[None, :]
+        do1 = tl.load(dp, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        do2 = tl.load(dp + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        cp = t[:, None] * HALF + dh[None, :]
+        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
+        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)
+
+        # inverse rotation
+        dy1 = do1 * c + do2 * s
+        dy2 = do2 * c - do1 * s
+
+        # per-program partial dgamma; summed by the caller (deterministic, no atomics)
+        tl.store(DGP + pid * D + dh, tl.sum(dy1 * n1, 0))
+        tl.store(DGP + pid * D + HALF + dh, tl.sum(dy2 * n2, 0))
+
+        g1 = tl.load(GAMMA + dh).to(tl.float32)
+        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
+        dn1 = dy1 * g1[None, :]
+        dn2 = dy2 * g2[None, :]
+
+        dot = (tl.sum(dn1 * n1, 1) + tl.sum(dn2 * n2, 1)) / D      # [R]
+        dx1 = rstd[:, None] * (dn1 - n1 * dot[:, None])
+        dx2 = rstd[:, None] * (dn2 - n2 * dot[:, None])
+
+        dbase = b * sd_b + h * sd_h + t * sd_t
+        q1 = DX + dbase[:, None] + dh[None, :]
+        tl.store(q1, dx1.to(DX.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(q1 + HALF, dx2.to(DX.dtype.element_ty), mask=mask_r[:, None])
+
+
+_BLOCK_R = 16
+_NUM_WARPS = 4
+
+
+def _launch_fwd(x, gamma, cos, sin, eps):
+    B, H, T, D = x.shape
+    N = B * H * T
+    out = torch.empty((B, H, T, D), dtype=x.dtype, device=x.device)
+    rstd = torch.empty(N, dtype=torch.float32, device=x.device)
+    grid = (triton.cdiv(N, _BLOCK_R),)
+    _fused_rms_rope_fwd[grid](
+        x, gamma, cos, sin, out, rstd,
+        x.stride(0), x.stride(1), x.stride(2),
+        T, H, N, eps,
+        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
+    )
+    return out, rstd
+
+
+def _launch_bwd(x, gamma, cos, sin, dout, rstd):
+    B, H, T, D = x.shape
+    N = B * H * T
+    dout = dout.contiguous()
+    dx = torch.empty_like(x)
+    nprog = triton.cdiv(N, _BLOCK_R)
+    dgp = torch.empty((nprog, D), dtype=torch.float32, device=x.device)
+    _fused_rms_rope_bwd[(nprog,)](
+        x, gamma, cos, sin, dout, rstd, dx, dgp,
+        x.stride(0), x.stride(1), x.stride(2),
+        dx.stride(0), dx.stride(1), dx.stride(2),
+        T, H, N,
+        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
+    )
+    return dx, dgp.sum(0).to(gamma.dtype)
+
+
+class _FusedRMSRoPE(torch.autograd.Function):
+    """QK-norm + RoPE for q and k in one autograd node (two kernel launches).
+    Mirrors the C++ `fused_rms_rope(q, k, q_gamma, k_gamma, cache)`: q/k keep
+    their incoming dtype (bf16 under autocast), the reduction and the rotation
+    run in fp32 internally, gamma and the cache stay fp32.
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, q_gamma, k_gamma, cos, sin, eps):
+        qo, q_rstd = _launch_fwd(q, q_gamma, cos, sin, eps)
+        ko, k_rstd = _launch_fwd(k, k_gamma, cos, sin, eps)
+        ctx.save_for_backward(q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd)
+        return qo, ko
+
+    @staticmethod
+    def backward(ctx, dqo, dko):
+        q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd = ctx.saved_tensors
+        dq, dqg = _launch_bwd(q, q_gamma, cos, sin, dqo, q_rstd)
+        dk, dkg = _launch_bwd(k, k_gamma, cos, sin, dko, k_rstd)
+        return dq, dk, dqg, dkg, None, None, None
+
+
+def fused_rms_rope(q, k, q_gamma, k_gamma, cos, sin, eps):
+    """q: [B,Hq,T,D], k: [B,Hkv,T,D] -> same shapes, contiguous.
+    `cos`/`sin` are [>=T, D//2] fp32; only the first T rows are read.
+    """
+    assert q.shape[-1] == k.shape[-1], "q and k must share head_dim"
+    D = q.shape[-1]
+    assert D % 2 == 0 and (D & (D - 1)) == 0, f"head_dim must be a power of two, got {D}"
+    assert q.stride(-1) == 1 and k.stride(-1) == 1, "head_dim must be stride-1"
+    T = q.shape[2]
+    return _FusedRMSRoPE.apply(q, k, q_gamma, k_gamma,
+                               cos[:T].contiguous(), sin[:T].contiguous(), eps)
+
+
+def rms_rope_eager(x, gamma, cos, sin, eps):
+    """Unfused QK-norm then NeoX RoPE, fp32 throughout — the reference the fused
+    kernel is measured against (`--rms-rope eager`), and the fallback when
+    Triton is unavailable."""
+    T, D = x.shape[2], x.shape[-1]
+    xf = x.float()
+    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * gamma
+    x1, x2 = xf[..., : D // 2], xf[..., D // 2:]
+    c = cos[:T][None, None, :, :]
+    s = sin[:T][None, None, :, :]
+    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)
+
+
+def build_rope_cache(seq_len, head_dim, theta, device):
+    """NeoX / Llama RoPE cache: cos, sin of shape [seq_len, head_dim//2], fp32."""
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
+                                             device=device) / head_dim))
+    pos = torch.arange(seq_len, dtype=torch.float32, device=device)
+    freqs = torch.outer(pos, inv_freq)                       # [T, D/2]
+    return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous()
+
+
+# =============================================================================
+# 2. Configuration — 1:1 with ModelConfig in bluscript_single.cpp — and the WSD
+#    learning-rate schedule.
+# =============================================================================
+@dataclass
+class ModelConfig:
+    # ---- model ----
+    d_model: int = 960
+    n_layers: int = 30
+    q_heads: int = 15
+    kv_heads: int = 5                     # grouped-query: G = q/kv = 4
+    head_dim: int = 64                   # the fused kernels need 64 or 128
+    ffn_hidden: int = 2560                # ~8/3 * d_model, rounded to /256
+    context_length: int = 4096            # RoPE cache length; must be >= T
+    rope_theta: float = 500000.0
+    vocab_size: int = 50304               # 50257 (tiktoken) padded for the GPU
+    weight_tying: bool = False
+
+    rms_eps: float = 1e-5                 # block RMSNorm epsilon
+    qk_eps: float = 1e-6                  # QK-norm (fused) epsilon
+
+    # ---- optimization / data ----
+    B: int = 4                            # micro-batch
+    T: int = 4096                         # sequence length (<= context_length)
+    global_batch: int = 524288            # tokens per optimizer step
+    grad_clip: float = 1.0
+    weight_decay: float = 0.1
+
+    # WSD: linear warmup -> constant at max_lr -> cosine decay to min_lr.
+    # Requires warmup_steps + stable_steps < max_steps - 1 (the remainder is
+    # the decay tail).
+    max_lr: float = 0.0005
+    min_lr: float = 0.00005                # 10% of max_lr
+    
+    muon_lr: float = 0.0005
+    muon_min_lr: float = 0.00005
+    
+    max_steps: int = 26703
+
+    warmup_steps: int = 500
+    stable_steps: int = int(max_steps * 0.9) - warmup_steps           # steps held at max_lr after warmup
+     
+
+    # ---- Muon (hidden 2D weights); the AdamW fields above still drive the
+    #      embedding / lm_head / 1D group, unchanged from the AdamW arm ----
+    # Default == max_lr: under adjust_lr_fn=match_rms_adamw the two are the same unit,
+    # so out of the box both arms take the same size step and differ only in
+    # direction. Same 10% floor, same WSD shape, so the schedules stay aligned.
+   
+    # The rest are torch.optim.Muon's own arguments, passed through unchanged.
+    # momentum/nesterov/ns_steps are torch's defaults; adjust_lr_fn is not (see
+    # the header — torch defaults to "original").
+    muon_momentum: float = 0.95           # beta on the momentum buffer
+    muon_ns_steps: int = 5                # Newton-Schulz iterations
+    muon_adjust_lr: str = "match_rms_adamw"   # "match_rms_adamw" | "original"
+
+    val_freq: int = 250
+    val_loss_steps: int = 20
+
+    # ---- checkpointing ----
+    checkpointing: bool = True
+    ckpt_freq: int = 1000                 # save every 1800 optimizer steps
+    ckpt_keep: int = 40                    # retain the most-recent N checkpoints
+    # Separate from the AdamW arm's dir: the two checkpoint formats differ (Muon
+    # momentum vs Adam m/v) and auto-resume picks the newest file it finds.
+    ckpt_dir: str = "ckpt_blu_muon"
+    ckpt_prefix: str = "blumodel_muon"
+
+    # ---- runtime knobs (not in the C++ struct) ----
+    fused_rms_rope: bool = True
+    activation_checkpoint: bool = False
+    loss_chunks: int = 1
+    # Depth probe: which layers --diag reports gradient / weight statistics for.
+    probe_layers: tuple = (0, 8, 17, 26, 35)
+
+
+def get_lr(step, cfg, max_lr=None, min_lr=None):
+    """WSD: linear Warmup -> Stable (constant max_lr) -> cosine Decay to min_lr.
+    Byte-for-byte the C++ get_lr().
+    `max_lr`/`min_lr` override the config pair so the Muon group can ride the
+    SAME schedule shape — same warmup length, same stable window, same cosine
+    tail — at its own amplitude. Called with no override it is identical to the
+    AdamW arm's get_lr(), which is what the embedding/norm group uses.
+    """
+    max_lr = cfg.max_lr if max_lr is None else max_lr
+    min_lr = cfg.min_lr if min_lr is None else min_lr
+    if step < cfg.warmup_steps:
+        return max_lr * (step + 1) / cfg.warmup_steps
+    if step < cfg.stable_steps + cfg.warmup_steps:
+        return max_lr
+    tail = cfg.max_steps - (cfg.stable_steps + cfg.warmup_steps) - 1
+    decay_ratio = (step - (cfg.stable_steps + cfg.warmup_steps)) / max(tail, 1)
+    decay_ratio = min(max(decay_ratio, 0.0), 1.0)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))       # 1 -> 0
+    return min_lr + coeff * (max_lr - min_lr)
+
+
+# =============================================================================
+# 3. Data — port of the C++ DataLoaderLite (Tensor-Implementations/include/
+#    DataLoader.h), including its pinned-host / side-stream prefetch ring.
+#
+# Shards are HEADERLESS flat uint16 token files. A shard is "val" if "val"
+# appears anywhere in its basename, "train" otherwise — this covers both an
+# explicit flux_train_*/flux_val_* naming and an unmarked-train naming like
+# flux_20b_shard_*.bin alongside a separately named *val* shard. For a batch
+# at position p:
+#     x = tokens[p : p+B*T]      y = tokens[p+1 : p+B*T+1]
+# p strides by B*T (world_size == 1 here). A shard tail that cannot fit a full
+# stride is dropped and the loader wraps to the next shard.
+#
+# The prefetch mirrors the C++ ring: NBUF pinned staging buffers and NBUF device
+# tensors, H2D issued on a side stream, CUDA events ordering producer against
+# consumer. Each slot also carries the (shard, token offset) it was filled from,
+# which is what --diag prints to name the bytes behind a loss spike.
+# =============================================================================
+class DataLoaderLite:
+    NBUF = 4
+
+    def __init__(self, B, T, split, data_root, device, master=True,
+                 max_tokens_per_shard=100_000_000, rank=0, world_size=1):
+        assert split in ("train", "val")
+        self.B, self.T = B, T
+        self.split = split
+        self.max_tokens = max_tokens_per_shard
+        self.device = device
+        # DDP: every rank reads the SAME shard list but a DISJOINT, interleaved
+        # slice of it (offset by rank, strided by world_size), so a training
+        # step's world_size*grad_accum micro-batches never repeat a window.
+        self.rank = rank
+        self.world_size = world_size
+
+        # Shard naming isn't consistent across datasets: some mark the split
+        # explicitly in both names (flux_train_*.bin / flux_val_*.bin), others
+        # only mark val and leave train shards unmarked (flux_20b_shard_*.bin
+        # + a separate *val* file). Treat "contains val" as the val split and
+        # everything else as train so both conventions resolve correctly.
+        self.shards = sorted(
+            p for p in glob.glob(os.path.join(data_root, "*.bin"))
+            if ("val" in os.path.basename(p).lower()) == (split == "val")
+        )
+        if not self.shards:
+            raise RuntimeError(f"no .bin shards found for split {split} in {data_root}")
+        if master:
+            print(f"found {len(self.shards)} shards for split {split}")
+
+        BT = B * T
+        self.pin_x = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
+        self.pin_y = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
+        self.dev_x = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
+        self.dev_y = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
+        self.copy_stream = torch.cuda.Stream(device=device)
+        self.h2d_done = [torch.cuda.Event() for _ in range(self.NBUF)]
+        self.h2d_recorded = [False] * self.NBUF
+        self.where = [(0, 0)] * self.NBUF        # (shard_index, token_offset)
+        self.reset()
+
+    # -- shard plumbing ------------------------------------------------------
+    def _load_shard(self, i):
+        self.cur = i
+        arr = np.memmap(self.shards[i], dtype=np.uint16, mode="r")
+        self.tok = arr
+        self.ntok = min(len(arr), self.max_tokens)
+
+    def _advance_shard(self):
+        self._load_shard((self.cur + 1) % len(self.shards))
+        self.pos = self.B * self.T * self.rank
+
+    def reset(self):
+        self.copy_stream.synchronize()
+        self._load_shard(0)
+        self.pos = self.B * self.T * self.rank
+        self.primed = False
+        self.buf = 0
+
+    # -- prefetch ------------------------------------------------------------
+    def _stage(self, idx):
+        BT = self.B * self.T
+        stride = BT * self.world_size
+        if self.ntok < BT + 1:
+            raise RuntimeError(f"shard too small for one batch: {self.shards[self.cur]}")
+        if self.pos + BT + 1 > self.ntok:
+            self._advance_shard()
+
+        # Host-side WAR guard: the pinned buffer may still be the source of the
+        # H2D issued the last time this slot was filled (that copy is gated on
+        # the device, not the host, so a host running ahead of the GPU would
+        # otherwise overwrite bytes still in flight and hand the model a batch
+        # spliced from two different windows).
+        if self.h2d_recorded[idx]:
+            self.h2d_done[idx].synchronize()
+
+        self.where[idx] = (self.cur, self.pos)
+        buf = np.asarray(self.tok[self.pos: self.pos + BT + 1]).astype(np.int64)
+        self.pin_x[idx].copy_(torch.from_numpy(buf[:-1]))
+        self.pin_y[idx].copy_(torch.from_numpy(buf[1:]))
+
+        with torch.cuda.stream(self.copy_stream):
+            self.dev_x[idx].view(-1).copy_(self.pin_x[idx], non_blocking=True)
+            self.dev_y[idx].view(-1).copy_(self.pin_y[idx], non_blocking=True)
+            self.h2d_done[idx].record(self.copy_stream)
+        self.h2d_recorded[idx] = True
+
+        self.pos += stride
+        if self.pos + BT + 1 > self.ntok:
+            self._advance_shard()
+
+    def next_batch(self):
+        """-> (x, y, shard_index, token_offset), x/y int64 [B,T] on device."""
+        if not self.primed:
+            for i in range(self.NBUF):
+                self._stage(i)
+            self.primed = True
+            self.buf = 0
+        else:
+            # Refill the slot returned LAST call. Its consumer's kernels are
+            # already queued on the compute stream, so the producer must wait
+            # for them before overwriting the device tensor.
+            self.copy_stream.wait_stream(torch.cuda.current_stream())
+            self._stage((self.buf - 1) % self.NBUF)
+
+        torch.cuda.current_stream().wait_event(self.h2d_done[self.buf])
+        i = self.buf
+        self.buf = (self.buf + 1) % self.NBUF
+        return self.dev_x[i], self.dev_y[i], self.where[i][0], self.where[i][1]
+
+    def skip_batches(self, n):
+        """Fast-forward `n` batches without touching the GPU (resume path)."""
+        self.copy_stream.synchronize()
+        self.primed = False
+        BT = self.B * self.T
+        stride = BT * self.world_size
+        for _ in range(n):
+            if self.pos + BT + 1 > self.ntok:
+                self._advance_shard()
+            self.pos += stride
+            if self.pos + BT + 1 > self.ntok:
+                self._advance_shard()
+
+
+# =============================================================================
+# 4. Model.
+#
+# Precision (bf16 mixed precision; both --precision modes run under autocast):
+#   bf16-params : Linear weights stored bf16, embedding / RMSNorm / QK-norm
+#                 scales stay fp32, fp32 master weights live in the optimizer
+#                 — the exact scheme bluscript_single.cpp uses (default).
+#   bf16        : all params fp32, autocast casts per-op (standard PyTorch AMP).
+# In both cases: matmuls + attention bf16, RMSNorm / QK-norm reductions, RoPE,
+# softmax and the loss in fp32. No GradScaler (bf16 has fp32's exponent range).
+# =============================================================================
+_FLASH = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm over d_model, fp32 weight. F.rms_norm is a fused ATen op and is
+    not autocast-eligible, so it runs in the input dtype — fp32 on the residual
+    stream, matching the C++ nn::RMSNorm."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.dim = (dim,)
+
+    def forward(self, x):
+        return F.rms_norm(x, self.dim, self.weight, self.eps)
+
+
+class CausalGQA(nn.Module):
+    """RMSNorm -> fused(QK-norm + RoPE) -> causal grouped-query attention -> Wo -> +x"""
+
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.cfg = cfg
+        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
+        self.wQ = nn.Linear(cfg.d_model, cfg.q_heads * cfg.head_dim, bias=False)
+        self.wK = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wV = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wO = nn.Linear(cfg.q_heads * cfg.head_dim, cfg.d_model, bias=False)
+        # QK-norm scales, initialised to 1.0 (identity RMS scale), fp32.
+        self.q_gamma = nn.Parameter(torch.ones(cfg.head_dim))
+        self.k_gamma = nn.Parameter(torch.ones(cfg.head_dim))
+        # GPT-2/Llama init: projections std=0.02, residual proj scaled by
+        # (2*n_layers)^-0.5 to keep residual-stream variance ~constant.
+        for lin in (self.wQ, self.wK, self.wV):
+            nn.init.normal_(lin.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.wO.weight, mean=0.0, std=0.02 * resid_scale)
+
+    def forward(self, x, cos, sin):
+        cfg = self.cfg
+        B, T, _ = x.shape
+        h = self.norm(x)
+
+        # [B,T,H*D] -> [B,H,T,D]. The transpose stays a view; the fused kernel
+        # honours arbitrary (b,h,t) strides, so no .contiguous() is paid here.
+        q = self.wQ(h).view(B, T, cfg.q_heads, cfg.head_dim).transpose(1, 2)
+        k = self.wK(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)
+        v = self.wV(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)
+
+        if cfg.fused_rms_rope:
+            q, k = fused_rms_rope(q, k, self.q_gamma, self.k_gamma, cos, sin, cfg.qk_eps)
+        else:
+            q = rms_rope_eager(q, self.q_gamma, cos, sin, cfg.qk_eps)
+            k = rms_rope_eager(k, self.k_gamma, cos, sin, cfg.qk_eps)
+
+        # Separate fused flash kernel. enable_gqa maps query head i to kv head
+        # i // (q_heads/kv_heads), the same grouping as the C++ kernel.
+        with sdpa_kernel(_FLASH):
+            o = F.scaled_dot_product_attention(
+                q.to(v.dtype), k.to(v.dtype), v, is_causal=True, enable_gqa=True)
+
+        o = o.transpose(1, 2).reshape(B, T, cfg.q_heads * cfg.head_dim)
+        return x + self.wO(o)
+
+
+class SwiGLUMLP(nn.Module):
+    """RMSNorm -> gate_up -> silu(gate)*up -> down -> +x   ([gate | up] layout)"""
+
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
+        self.gate_up = nn.Linear(cfg.d_model, 2 * cfg.ffn_hidden, bias=False)
+        self.down = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
+        nn.init.normal_(self.gate_up.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.down.weight, mean=0.0, std=0.02 * resid_scale)
+
+    def forward(self, x):
+        gate, up = self.gate_up(self.norm(x)).chunk(2, dim=-1)
+        return x + self.down(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.attn = CausalGQA(cfg, resid_scale)
+        self.mlp = SwiGLUMLP(cfg, resid_scale)
+
+    def forward(self, x, cos, sin):
+        return self.mlp(self.attn(x, cos, sin))
+
+
+class BLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        resid_scale = 1.0 / math.sqrt(2.0 * cfg.n_layers)
+
+        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        nn.init.normal_(self.wte.weight, mean=0.0, std=0.02)
+
+        self.blocks = nn.ModuleList([Block(cfg, resid_scale) for _ in range(cfg.n_layers)])
+        self.norm_f = RMSNorm(cfg.d_model, cfg.rms_eps)
+        if not cfg.weight_tying:
+            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+            nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)
+
+        cos, sin = build_rope_cache(cfg.context_length, cfg.head_dim, cfg.rope_theta, "cpu")
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def apply_param_dtype(self, mode):
+        """`bf16-params`: store Linear weights in bf16 (embedding, RMSNorm and
+        QK-norm scales stay fp32) — the two-dtype-group layout of the C++ model.
+        Under autocast this is transparent: the cast to bf16 becomes a no-op."""
+        if mode != "bf16-params":
+            return
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.weight.data = m.weight.data.to(torch.bfloat16)
+
+    def forward(self, idx, targets=None):
+        cfg = self.cfg
+        T = idx.shape[1]
+        cos, sin = self.rope_cos[:T], self.rope_sin[:T]
+
+        x = self.wte(idx)                                    # fp32 [B,T,d_model]
+        for blk in self.blocks:
+            if cfg.activation_checkpoint and self.training:
+                x = checkpoint(blk, x, cos, sin, use_reentrant=False)
+            else:
+                x = blk(x, cos, sin)
+        x = self.norm_f(x)
+
+        head_w = self.wte.weight if cfg.weight_tying else self.lm_head.weight
+        if targets is None:
+            return F.linear(x, head_w), None
+        if cfg.loss_chunks > 1:
+            return None, _chunked_ce(x, head_w, targets, cfg.loss_chunks)
+        # Matches the C++: bf16 matmul, logits promoted to fp32 for the loss
+        # (cross_entropy is an autocast-fp32 op, so the promotion is implicit).
+        logits = F.linear(x, head_w)
+        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1))
+        return logits, loss
+
+
+def _chunked_ce(x, weight, targets, chunks):
+    """lm_head + cross-entropy in row chunks, logits recomputed in backward.
+    Peak logit memory drops from [B*T, V] to [B*T/chunks, V] at the cost of one
+    extra lm_head matmul. Off by default (`--loss-chunks 1`) so the memory
+    profile matches the C++ script, which materialises the full fp32 logits.
+    """
+    xf = x.reshape(-1, x.shape[-1])
+    tf = targets.reshape(-1)
+    n = xf.shape[0]
+    size = (n + chunks - 1) // chunks
+    total = xf.new_zeros((), dtype=torch.float32)
+    for lo in range(0, n, size):
+        hi = min(lo + size, n)
+        part = checkpoint(_ce_chunk, xf[lo:hi], weight, tf[lo:hi], use_reentrant=False)
+        total = total + part * (hi - lo)
+    return total / n
+
+
+def _ce_chunk(x, weight, targets):
+    return F.cross_entropy(F.linear(x, weight).float(), targets)
+
+
+# =============================================================================
+# 5. Optimizer (Muon on the hidden matrices, AdamW on the rest), gradient
+#    clipping, checkpointing.
+# =============================================================================
+def _is_muon_param(name, p):
+    """Muon orthogonalizes a matrix, so it needs the parameter to BE a linear
+    map between two feature spaces:
+      * dim >= 2 — RMSNorm weights and QK-norm gammas are 1D and have no
+        singular values to equalise;
+      * not the embedding / lm_head — its rows index tokens, not features, so
+        the left singular vectors mix unrelated vocabulary entries. Weight
+        tying makes this one tensor serving both ends; either way it is AdamW's.
+    Everything else — wQ/wK/wV/wO, gate_up, down — is Muon's: 75.5M of the
+    114.2M parameters at the default config.
+    """
+    return p.dim() >= 2 and not name.startswith(("wte.", "lm_head."))
+
+
+def _state_bytes(opt):
+    """Bytes held in a torch optimizer's per-parameter state — Adam's m/v, or
+    Muon's momentum buffer. Reads 0 before the first step(): both allocate
+    lazily."""
+    return sum(v.numel() * v.element_size()
+               for st in opt.state.values()
+               for v in st.values() if torch.is_tensor(v))
+
+
+class MasterMuon:
+    """torch.optim.Muon on the hidden matrices + torch.optim.AdamW on the rest,
+    both driving fp32 master weights.
+    Neither optimizer is reimplemented here. What this class owns is the part
+    torch cannot know about:
+      * ROUTING — torch.optim.Muon accepts whatever you hand it and the docs are
+        explicit that only 2D hidden weights belong ("bias, and embedding,
+        should be optimized by a standard method such as AdamW"). The split is
+        `_is_muon_param`.
+      * MASTER WEIGHTS — the parameters are bf16 under --precision bf16-params,
+        and a bf16 tensor cannot absorb an lr*update of 1e-4 relative size. So
+        both optimizers are constructed over the fp32 MASTERS, fed from the fp32
+        landing buffers, and the masters are cast back over the params after the
+        step. Unchanged from the AdamW arm — that machinery is what makes the
+        two arms comparable, since both optimizers then read a bit-identical
+        fp32 gradient and write a bit-identical bf16 parameter.
+    Weight decay is decoupled in both (torch.optim.Muon applies it before the
+    update, as AdamW does) and is applied to the master for the same reason: a
+    decay written to the bf16 param is overwritten by the master write-back
+    microseconds later and silently no-ops.
+    State cost: Muon holds one momentum buffer per param, against Adam's two.
+    """
+
+    def __init__(self, named_params, cfg):
+        self.cfg = cfg
+        self.triples = []                 # (param, master, landing_grad), ALL params
+        self.muon_masters = []            # masters routed to Muon
+        decay, nodecay = [], []
+        self.n_muon = self.n_decay = self.n_nodecay = 0
+        for name, p in named_params:
+            if p.dtype == torch.float32:
+                m, g = p, None
+            else:
+                m = p.detach().float().clone()
+                g = torch.zeros_like(m)
+                m.grad = g
+            self.triples.append((p, m, g))
+            if _is_muon_param(name, p):
+                self.muon_masters.append(m)
+                self.n_muon += m.numel()
+            elif p.dim() >= 2:
+                decay.append(m)
+                self.n_decay += m.numel()
+            else:
+                nodecay.append(m)
+                self.n_nodecay += m.numel()
+        self.opt = torch.optim.AdamW(
+            [{"params": decay, "weight_decay": cfg.weight_decay},
+             {"params": nodecay, "weight_decay": 0.0}],
+            lr=cfg.max_lr, betas=(0.9, 0.95), eps=1e-8, fused=True)
+        # Everything below is a torch.optim.Muon default except adjust_lr_fn
+        # (torch: "original") and weight_decay, which is pinned to the AdamW
+        # arm's 0.1 so the decay is identical across both groups and both arms.
+        self.muon = torch.optim.Muon(
+            self.muon_masters,
+            lr=cfg.muon_lr,
+            weight_decay=cfg.weight_decay,
+            momentum=cfg.muon_momentum,
+            nesterov=True,
+            ns_steps=cfg.muon_ns_steps,
+            adjust_lr_fn=cfg.muon_adjust_lr)
+
+    @torch.no_grad()
+    def step(self, lr, muon_lr):
+        for p, m, g in self.triples:
+            if g is not None:
+                g.copy_(p.grad)               # bf16 grad -> fp32 landing buffer
+        # AdamW group (embedding / lm_head / norms / gammas) — identical to the
+        # baseline's update, on identical gradients.
+        for grp in self.opt.param_groups:
+            grp["lr"] = lr
+        self.opt.step()
+        # Muon group. Its nesterov path may consume the gradient in place; that
+        # is harmless here because grad_norm and the --diag probes are read
+        # before step() and the landing buffers are overwritten wholesale on the
+        # next one.
+        for grp in self.muon.param_groups:
+            grp["lr"] = muon_lr
+        self.muon.step()
+        for p, m, g in self.triples:
+            if g is not None:
+                p.data.copy_(m)               # fp32 master -> bf16 param
+
+    def state_breakdown(self):
+        """-> (master_bytes, landing_grad_bytes, adam_mv_bytes, muon_mom_bytes).
+        The first two are zero under --precision bf16, where the master IS the
+        parameter and the gradient is already fp32. Both optimizers allocate
+        their state lazily on the first step(), so the last two read 0 until
+        then — the startup report uses the expected_* predictors instead.
+        """
+        master = landing = 0
+        for p, m, g in self.triples:
+            if m is not p:
+                master += m.numel() * m.element_size()
+                landing += g.numel() * g.element_size()
+        mv = _state_bytes(self.opt)
+        mom = _state_bytes(self.muon)
+        return master, landing, mv, mom
+
+    def state_bytes(self):
+        return sum(self.state_breakdown())
+
+    def expected_mv_bytes(self):
+        """Adam m/v are lazily allocated, so the startup report has to predict
+        them rather than measure them. Muon's params are excluded — they carry a
+        momentum buffer instead, reported separately."""
+        muon_ids = {id(m) for m in self.muon_masters}
+        return sum(m.numel() * m.element_size() * 2
+                   for _, m, _ in self.triples if id(m) not in muon_ids)
+
+    def expected_muon_mom_bytes(self):
+        """One fp32 momentum buffer per Muon parameter — the prediction the
+        startup report prints before step 0 has allocated it."""
+        return sum(m.numel() * m.element_size() for m in self.muon_masters)
+
+    # -- checkpointing -------------------------------------------------------
+    def state_dict(self):
+        # `optimizer` tags the arm: resuming a Muon run from an AdamW checkpoint
+        # would restore the weights and then silently drop the momentum, so
+        # CheckpointManager refuses on a mismatch rather than producing a run
+        # that looks fine and is not comparable to either arm.
+        return {"optimizer": "muon",
+                "adam": self.opt.state_dict(),
+                "muon": self.muon.state_dict(),
+                "masters": [None if m is p else m for p, m, g in self.triples]}
+
+    @torch.no_grad()
+    def load_state_dict(self, blob):
+        """-> True if the fp32 masters were restored from the checkpoint."""
+        if blob.get("optimizer") != "muon":
+            raise SystemExit(
+                "ERROR: this checkpoint was written by the AdamW arm (train.py). "
+                "Point --ckpt-dir somewhere else, or resume it with train.py.")
+        self.opt.load_state_dict(blob["adam"])
+        if blob.get("muon") is not None:
+            self.muon.load_state_dict(blob["muon"])
+        masters = blob.get("masters")
+        if masters is None or len(masters) != len(self.triples):
+            return False
+        for (p, m, g), saved in zip(self.triples, masters):
+            if saved is not None and m is not p:
+                m.copy_(saved.to(m.device))
+        return True
+
+    @torch.no_grad()
+    def resync_masters(self):
+        """Re-seed every master from its parameter.
+        Needed after loading a checkpoint that carries no masters: the masters
+        were built in __init__ from the RANDOM INIT, so without this the first
+        step() would write them straight back over the restored weights and
+        silently throw the run away.
+        """
+        for p, m, g in self.triples:
+            if m is not p:
+                m.copy_(p.detach().float())
+
+
+def clip_grads_(params, max_norm):
+    """ONE global L2 norm over every parameter, then one rescale — the textbook
+    definition. (It differs from the ZeRO script, which clips each dtype group
+    against its own norm and can therefore let the true global norm reach
+    sqrt(n_groups) * grad_clip.) The per-tensor norms are promoted to fp32
+    before they are combined, or a reduction over ~250 bf16 norms would lose
+    ~3 digits."""
+    grads = [p.grad for p in params if p.grad is not None]
+    norms = torch._foreach_norm(grads)
+    total = torch.linalg.vector_norm(torch.stack([n.float() for n in norms]))
+    if max_norm and max_norm > 0:
+        clip = (max_norm / (total + 1e-6)).clamp(max=1.0)
+        torch._foreach_mul_(grads, clip)
+    return total
+
+
+class CheckpointManager:
+    """Synchronous save / keep-N / auto-resume, mirroring the C++
+    CheckpointManager. A checkpoint holds the model parameters, Adam's m/v, the
+    fp32 masters, and the step it should resume at."""
+
+    def __init__(self, ckpt_dir, prefix, keep):
+        self.dir, self.prefix, self.keep = ckpt_dir, prefix, keep
+        os.makedirs(ckpt_dir, exist_ok=True)
+
+    def _path(self, step):
+        return os.path.join(self.dir, f"{self.prefix}_step{step:08d}.pt")
+
+    def _existing(self):
+        out = []
+        for p in glob.glob(os.path.join(self.dir, f"{self.prefix}_step*.pt")):
+            m = re.search(r"_step(\d+)\.pt$", os.path.basename(p))
+            if m:
+                out.append((int(m.group(1)), p))
+        return sorted(out)
+
+    def save(self, step, model, opt, loss, last_val, cfg):
+        tmp = self._path(step) + ".tmp"
+        torch.save({"step": step, "loss": loss, "last_val": last_val,
+                    # B/T/global_batch pin the data stride this checkpoint's
+                    # step count assumes. world_size is deliberately NOT
+                    # recorded: grad_accum = global_batch // (B*T*world_size)
+                    # already makes tokens-per-step (and skip_batches'
+                    # target offset) independent of how many ranks are
+                    # running, so a GPU-count change alone is not a mismatch.
+                    "run_cfg": {"B": cfg.B, "T": cfg.T,
+                                "global_batch": cfg.global_batch},
+                    "model": model.state_dict(),
+                    "opt": opt.state_dict()}, tmp)
+        # Rename last: a crash mid-write leaves a .tmp, never a truncated
+        # checkpoint that the next load_latest would pick as the newest.
+        os.replace(tmp, self._path(step))
+        if self.keep > 0:
+            for _, stale in self._existing()[:-self.keep]:
+                os.remove(stale)
+
+    def load_latest(self, model, opt, cfg, master=True):
+        """-> (start_step, loss, last_val). (0, 0.0, nan) when there is
+        nothing to resume."""
+        existing = self._existing()
+        if not existing:
+            return 0, 0.0, float("nan")
+        _, path = existing[-1]
+        return self.load(path, model, opt, cfg, master)
+
+    def load(self, path, model, opt, cfg, master=True):
+        """Load a specific checkpoint file (e.g. from --resume-from), rather
+        than auto-discovering the latest one in self.dir.
+        -> (start_step, loss, last_val)."""
+        blob = torch.load(path, map_location="cpu", weights_only=True)
+        model.load_state_dict(blob["model"])
+        if not opt.load_state_dict(blob["opt"]):
+            opt.resync_masters()
+        run_cfg = blob.get("run_cfg")
+        if master and run_cfg is not None:
+            mismatches = [f"{k} checkpoint={v!r} current={getattr(cfg, k)!r}"
+                         for k, v in run_cfg.items() if getattr(cfg, k) != v]
+            if mismatches:
+                print("[WARN] resume config differs from the checkpoint's — "
+                      "the data offset this step count assumes will be wrong:")
+                for m in mismatches:
+                    print(f"         {m}")
+        print(f"[resume] loaded {path}")
+        return (int(blob["step"]), float(blob["loss"]),
+                float(blob.get("last_val", float("nan"))))
+
+
+# =============================================================================
+# 6. Training.
+# =============================================================================
+def _env_int(name, default=0):
+    v = os.environ.get(name)
+    if v in (None, ""):
+        return default
+    try:
+        return int(v)
+    except ValueError:
+        print(f"[WARN] ignoring non-integer {name}={v!r}")
+        return default
+
+
+def seed_everything(seed):
+    """Every RNG the run can draw from, off one seed.
+    torch's CPU generator is the one that matters here (the model is built on
+    CPU and moved), but `random` / numpy are seeded too so that anything added
+    later — a shuffled shard order, a sampled probe — is reproducible by
+    default rather than by accident.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def apply_determinism(mode):
+    """Pin the run-to-run drift that seeding does not cover.
+    Three sources, all on the backward pass:
+      * nn.Embedding's dense backward scatters one atomicAdd per token into
+        [vocab, d_model]; with weight tying that is also the lm_head gradient,
+        so it is the largest reduction in the step;
+      * the flash / mem-efficient attention backward splits dq over the key
+        blocks and atomicAdds them together;
+      * cuDNN picks kernels by wall-clock timing when `benchmark` is on, so the
+        chosen algorithm can differ between runs of the same shape.
+    Each is a float reduction in an order the scheduler chooses, so two runs
+    diverge in the last mantissa bits of the first step and, through the
+    optimizer, visibly by a few hundred steps.
+    `strict` raises if some op has no deterministic kernel — better than
+    reporting a reproducible run that is not one. `warn` downgrades that to a
+    warning and keeps training; `off` restores the default fast kernels.
+    """
+    if mode == "off":
+        return
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    # Routes embedding-backward and SDPA-backward to their ordered variants, and
+    # enforces the cuBLAS workspace setting made at the top of this file.
+    torch.use_deterministic_algorithms(True, warn_only=(mode == "warn"))
+
+
+def build_argparser():
+    ap = argparse.ArgumentParser(
+        description="BLU single-GPU training — PyTorch port of bluscript_single.cpp")
+    ap.add_argument("--data-root",
+                    default=("/root/.cache/huggingface/hub/datasets--jayasriii--flux_20b/"
+                             "/snapshots/ddc6cc7e43365a6b9096c30053df0fb001c94f4c"))
+    ap.add_argument("--device", type=int, default=None,
+                    help="CUDA device index (env: BLU_DEVICE)")
+    ap.add_argument("--precision", choices=["bf16-params", "bf16"], default="bf16-params",
+                    help="bf16-params (default): Linear weights stored bf16 with "
+                         "fp32 master weights in the optimizer — the C++ scheme. "
+                         "bf16: fp32 params + autocast (standard PyTorch AMP).")
+    ap.add_argument("--rms-rope", choices=["fused", "eager"], default="fused",
+                    help="fused = Triton QK-norm+RoPE kernel (default), "
+                         "eager = unfused reference, for A/B on the kernel")
+    ap.add_argument("--compile", action="store_true", help="torch.compile each block")
+    ap.add_argument("--activation-checkpoint", action="store_true",
+                    help="recompute each block in backward (trades ~30%% time for memory)")
+    ap.add_argument("--loss-chunks", type=int, default=1,
+                    help="chunk lm_head+CE into N pieces with recompute; 1 (default) "
+                         "materialises the full logits, as the C++ does")
+    # bluscript_single.cpp defaults unless overridden.
+    ap.add_argument("--B", type=int, default=None, help="micro-batch (env: BLU_B)")
+    ap.add_argument("--T", type=int, default=None)
+    ap.add_argument("--d-model", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--n-layers", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--q-heads", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--kv-heads", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--ffn-hidden", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--vocab-size", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--global-batch", type=int, default=None,
+                    help="tokens per optimizer step (env: BLU_GLOBAL_BATCH)")
+    ap.add_argument("--max-steps", type=int, default=None,
+                    help="caps the run for benchmarking; warmup/stable shrink "
+                         "proportionally so the WSD decay tail stays non-empty "
+                         "(env: BLU_MAX_STEPS)")
+    # ---- Muon ----
+    ap.add_argument("--muon-lr", type=float, default=None,
+                    help="peak LR for the hidden 2D weights. Under the default "
+                         "--muon-adjust-lr match_rms_adamw this is the same unit as "
+                         "--max-lr, and defaults to the same value (2e-4) so "
+                         "both arms take a matched first run; sweep it upward "
+                         "from there (Muon usually takes 2-4x AdamW's LR).")
+    ap.add_argument("--muon-min-lr", type=float, default=None,
+                    help="floor of the Muon cosine tail (default 10%% of --muon-lr)")
+    ap.add_argument("--muon-momentum", type=float, default=None,
+                    help="torch.optim.Muon momentum (default 0.95, torch's own)")
+    ap.add_argument("--muon-ns-steps", type=int, default=None,
+                    help="torch.optim.Muon ns_steps, Newton-Schulz iterations "
+                         "(default 5, torch's own); fewer = looser "
+                         "orthogonalization, cheaper")
+    ap.add_argument("--muon-adjust-lr", choices=["match_rms_adamw", "original"],
+                    default=None,
+                    help="torch.optim.Muon adjust_lr_fn. match_rms_adamw "
+                         "(default HERE): 0.2*sqrt(max(A,B)), putting --muon-lr "
+                         "on the same scale as --max-lr. original (default in "
+                         "torch): sqrt(max(1,A/B)), needs a ~5.5x larger "
+                         "--muon-lr at d_model=768.")
+    ap.add_argument("--val-freq", type=int, default=None)
+    ap.add_argument("--val-loss-steps", type=int, default=None)
+    ap.add_argument("--diag", action="store_true",
+                    help="per-micro-batch loss + depth probes (env: BLU_DIAG=1)")
+    ap.add_argument("--warmup-measure", type=int, default=3,
+                    help="steps excluded from the steady-state throughput summary")
+    ap.add_argument("--csv", default="blu_muon_log.csv",
+                    help="CSV log path; appended (not truncated) on resume. The "
+                         "first 8 columns match the AdamW arm's CSV exactly, so "
+                         "the two files plot against each other directly; "
+                         "muon_lr is appended as a 9th.")
+    ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--determinism", choices=["strict", "warn", "off"], default="strict",
+                    help="strict (default): deterministic CUDA kernels, so a rerun "
+                         "of the same launch line reproduces the loss curve bit for "
+                         "bit; raises if an op has no deterministic implementation. "
+                         "warn: same, but non-deterministic ops warn instead of "
+                         "raising. off: fastest kernels, results drift between runs.")
+    # checkpointing
+    ap.add_argument("--no-checkpointing", action="store_true")
+    ap.add_argument("--ckpt-dir", default=None)
+    ap.add_argument("--ckpt-prefix", default=None)
+    ap.add_argument("--ckpt-freq", type=int, default=None)
+    ap.add_argument("--ckpt-keep", type=int, default=None)
+    ap.add_argument("--resume-from", default=None,
+                    help="path to a specific *_step*.pt checkpoint to resume "
+                         "from, overriding auto-discovery of the latest "
+                         "checkpoint in --ckpt-dir")
+    return ap
+
+
+def config_from_args(args):
+    cfg = ModelConfig()
+    for f in ("B", "T", "d_model", "n_layers", "q_heads", "kv_heads", "ffn_hidden",
+              "vocab_size", "global_batch", "val_freq", "val_loss_steps",
+              "ckpt_dir", "ckpt_prefix", "ckpt_freq", "ckpt_keep",
+              "muon_lr", "muon_min_lr", "muon_momentum", "muon_ns_steps",
+              "muon_adjust_lr"):
+        v = getattr(args, f, None)
+        if v is not None:
+            setattr(cfg, f, v)
+    # A --muon-lr given without a floor keeps the 10% ratio the WSD tail assumes,
+    # instead of silently decaying to the default 2e-5 from wherever it starts.
+    if args.muon_lr is not None and args.muon_min_lr is None:
+        cfg.muon_min_lr = 0.1 * cfg.muon_lr
+    cfg.fused_rms_rope = (args.rms_rope == "fused")
+    cfg.activation_checkpoint = args.activation_checkpoint
+    cfg.loss_chunks = args.loss_chunks
+    cfg.checkpointing = not args.no_checkpointing
+    cfg.context_length = max(cfg.context_length, cfg.T)
+
+    # Mirrors the BLU_MAX_STEPS branch in bluscript_single.cpp: shrink warmup and
+    # stable with the run, or a short benchmark leaves no decay tail (the C++
+    # asserts on it; get_lr here would clamp, which is worse — it would silently
+    # report a schedule the config does not describe).
+    if args.max_steps is not None:
+        ms = args.max_steps
+        cfg.max_steps = ms
+        cfg.warmup_steps = max(1, ms // 10)
+        cfg.stable_steps = max(1, ms // 2)
+        if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
+            cfg.stable_steps = max(1, cfg.max_steps - cfg.warmup_steps - 2)
+        print(f"[bench] max_steps={ms} (warmup={cfg.warmup_steps} "
+              f"stable={cfg.stable_steps})")
+    return cfg
+
+
+def validate_config(cfg, world_size=1, master=True):
+    """Cheap checks; each one is a silent-corruption or out-of-bounds failure
+    mode if it does not hold."""
+    bad = []
+    if cfg.T > cfg.context_length:
+        bad.append(f"T ({cfg.T}) > context_length ({cfg.context_length}): "
+                   "the RoPE cache is too short")
+    if cfg.q_heads % cfg.kv_heads != 0:
+        bad.append(f"q_heads ({cfg.q_heads}) not divisible by kv_heads ({cfg.kv_heads})")
+    if cfg.head_dim not in (64, 128):
+        bad.append(f"the fused kernels support head_dim 64 or 128, got {cfg.head_dim}")
+    # Under DDP each rank runs B*T tokens per micro-step, so the world's tokens
+    # per micro-step are B*T*world_size; global_batch must land on a whole
+    # number of those or grad_accum below truncates and quietly under-shoots
+    # the configured optimizer-step token count.
+    if cfg.global_batch % (cfg.B * cfg.T * world_size) != 0:
+        bad.append(f"global_batch ({cfg.global_batch}) not divisible by "
+                   f"B*T*world_size ({cfg.B * cfg.T * world_size})")
+    if cfg.warmup_steps < 1:
+        bad.append("warmup_steps must be >= 1 (get_lr divides by it)")
+    if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
+        bad.append(f"warmup_steps + stable_steps "
+                   f"({cfg.warmup_steps + cfg.stable_steps}) leaves no decay tail "
+                   f"before max_steps ({cfg.max_steps})")
+    if not hasattr(torch.optim, "Muon"):
+        bad.append(f"torch {torch.__version__} has no torch.optim.Muon — it first "
+                   "appears in the 2.9 docs. Upgrade torch, or run the AdamW "
+                   "arm (train.py).")
+    if cfg.muon_ns_steps < 1:
+        bad.append("muon_ns_steps must be >= 1 (0 leaves the update unorthogonalized, "
+                   "which is not Muon)")
+    if cfg.muon_momentum <= 0.0 or cfg.muon_momentum >= 1.0:
+        bad.append(f"muon_momentum must be in (0,1), got {cfg.muon_momentum}")
+    if cfg.q_heads * cfg.head_dim != cfg.d_model and master:
+        print(f"WARNING: q_heads*head_dim ({cfg.q_heads * cfg.head_dim}) "
+              f"!= d_model ({cfg.d_model})")
+    if bad:
+        raise SystemExit("ERROR: " + "\nERROR: ".join(bad))
+
+
+def main():
+    args = build_argparser().parse_args()
+
+    # Env overrides, for parity with the C++ launch line. An explicit CLI flag
+    # wins over the environment.
+    if args.B is None:
+        args.B = _env_int("BLU_B", 0) or None
+    if args.global_batch is None:
+        args.global_batch = _env_int("BLU_GLOBAL_BATCH", 0) or None
+    if args.max_steps is None:
+        args.max_steps = _env_int("BLU_MAX_STEPS", 0) or None
+    diag = args.diag or _env_int("BLU_DIAG", 0) != 0
+
+    if not torch.cuda.is_available():
+        raise SystemExit("ERROR: no CUDA device available")
+
+    # ---- DDP bootstrap ----
+    # torchrun sets RANK/LOCAL_RANK/WORLD_SIZE; their absence means a plain
+    # single-process launch, where --device / BLU_DEVICE pick the GPU instead.
+    # torch.cuda.set_device MUST happen before init_process_group: NCCL creates
+    # its context on whatever device is current at that call, so doing it after
+    # would leave every rank on this node contending for device 0.
+    ddp = int(os.environ.get("RANK", -1)) != -1
+    if ddp:
+        ddp_rank = int(os.environ["RANK"])
+        ddp_local_rank = int(os.environ["LOCAL_RANK"])
+        ddp_world_size = int(os.environ["WORLD_SIZE"])
+        master_process = ddp_rank == 0
+        if ddp_local_rank >= torch.cuda.device_count():
+            raise SystemExit(f"ERROR: LOCAL_RANK={ddp_local_rank} but only "
+                             f"{torch.cuda.device_count()} CUDA device(s) visible")
+        torch.cuda.set_device(ddp_local_rank)
+        device = torch.device("cuda", ddp_local_rank)
+        dist.init_process_group(backend="nccl")
+        if master_process and (args.device is not None
+                               or os.environ.get("BLU_DEVICE") not in (None, "")):
+            print("[WARN] --device / BLU_DEVICE is ignored under torchrun — "
+                  "LOCAL_RANK selects each rank's GPU.")
+    else:
+        ddp_rank, ddp_local_rank, ddp_world_size = 0, 0, 1
+        master_process = True
+        if args.device is None:
+            args.device = _env_int("BLU_DEVICE", 0)
+        if args.device >= torch.cuda.device_count():
+            raise SystemExit(f"ERROR: BLU_DEVICE={args.device} but only "
+                             f"{torch.cuda.device_count()} CUDA device(s) visible")
+        torch.cuda.set_device(args.device)
+        device = torch.device("cuda", args.device)
+
+    seed_everything(args.seed)
+    apply_determinism(args.determinism)
+    # TF32 loses mantissa bits but is deterministic, so it is orthogonal to the
+    # setting above and stays on.
+    torch.set_float32_matmul_precision("high")
+
+    cfg = config_from_args(args)
+    if cfg.fused_rms_rope and not _HAS_TRITON:
+        if master_process:
+            print("[WARN] Triton not available — falling back to the eager QK-norm+RoPE")
+        cfg.fused_rms_rope = False
+    validate_config(cfg, ddp_world_size, master_process)
+
+    # Tokens consumed per optimizer step = B * T * grad_accum * world_size.
+    grad_accum = cfg.global_batch // (cfg.B * cfg.T * ddp_world_size)
+    tokens_per_step = cfg.B * cfg.T * grad_accum * ddp_world_size
+
+    if master_process:
+        print("================ Configuration ================")
+        mode_str = (f"DDP x{ddp_world_size} (torchrun, NCCL)" if ddp
+                   else "SINGLE GPU (no torchrun / DDP / ZeRO)")
+        print(f"  mode            : {mode_str}")
+        print(f"  cuda device     : {ddp_local_rank if ddp else args.device} "
+              f"({torch.cuda.get_device_name(device)})")
+        print(f"  d_model         : {cfg.d_model}")
+        print(f"  n_layers        : {cfg.n_layers}")
+        print(f"  q_heads/kv_heads: {cfg.q_heads} / {cfg.kv_heads} "
+              f"(G={cfg.q_heads // cfg.kv_heads})")
+        print(f"  head_dim        : {cfg.head_dim}")
+        print(f"  ffn_hidden      : {cfg.ffn_hidden}  (SwiGLU)")
+        print(f"  context_length  : {cfg.context_length}")
+        print(f"  rope_theta      : {cfg.rope_theta:g}  (NeoX half-split)")
+        print(f"  vocab_size      : {cfg.vocab_size}")
+        print(f"  weight_tying    : {'ON' if cfg.weight_tying else 'OFF'}")
+        print(f"  B x T           : {cfg.B} x {cfg.T}")
+        print(f"  global_batch    : {cfg.global_batch}  (grad_accum={grad_accum}, "
+              f"world_size={ddp_world_size})")
+        print("  optimizer       : MUON (hidden 2D) + AdamW (embed/head/1D)")
+        print(f"  max_lr / min_lr : {cfg.max_lr:g} / {cfg.min_lr:g}   (AdamW group)")
+        print(f"  muon lr / min   : {cfg.muon_lr:g} / {cfg.muon_min_lr:g}   "
+              f"(adjust_lr_fn={cfg.muon_adjust_lr})")
+        print(f"  muon momentum   : {cfg.muon_momentum:g}  | NS steps {cfg.muon_ns_steps}")
+        print(f"  warmup / steps  : {cfg.warmup_steps} / {cfg.max_steps}")
+        print(f"  WSD stable_steps: {cfg.stable_steps}  "
+              f"(decay tail={cfg.max_steps - cfg.warmup_steps - cfg.stable_steps})")
+        print(f"  precision       : {args.precision}")
+        print(f"  seed / determ.  : {args.seed} / {args.determinism}")
+        print(f"  rms+rope        : {'fused' if cfg.fused_rms_rope else 'eager'}"
+              f" | compile {args.compile} | act-ckpt {cfg.activation_checkpoint}"
+              f" | loss_chunks {cfg.loss_chunks}")
+        print("===============================================", flush=True)
+
+    # ---- model ----
+    # Built and seeded identically on every rank (same seed, deterministic CPU
+    # init); DDP's constructor broadcasts rank 0's parameters/buffers anyway,
+    # so this is belt-and-suspenders against any future source of divergence.
+    if master_process:
+        print(f"Initializing model on CUDA device {device.index}...", flush=True)
+    raw_model = BLU(cfg).to(device)
+    raw_model.apply_param_dtype(args.precision)
+    params = list(raw_model.parameters())
+    n_params = sum(p.numel() for p in params)
+    if master_process:
+        print(f"Number of parameters: {n_params}")
+    if args.compile:
+        for blk in raw_model.blocks:
+            blk.compile()
+
+    # DDP wraps the (already dtype-cast, already compiled) raw module. Buffer
+    # broadcasting is off: rope_cos/rope_sin are deterministic functions of
+    # cfg, identical on every rank already, so syncing them every forward
+    # would just be wasted communication.
+    # if ddp:
+    #     model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
+    # else:
+    #     model = raw_model
+
+    if ddp:
+        model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
+        torch.cuda.synchronize()   # let DDP's constructor-time NCCL broadcast fully land
+                                # before DataLoaderLite's independent copy_stream starts
+                                # issuing its own async H2D work — see [crash investigation]
+    else:
+        model = raw_model
+
+
+    # ---- optimizer ----
+    # named_parameters(), not parameters(): the routing is by name (the
+    # embedding and lm_head stay on AdamW). Duplicates are dropped by default,
+    # so a weight-tied head is visited once, as `wte.weight`. Built from
+    # raw_model so names never carry DDP's "module." prefix.
+    opt = MasterMuon(raw_model.named_parameters(), cfg)
+    if master_process:
+        print(f"muon params: {opt.n_muon} | adamw decayed: {opt.n_decay}, "
+              f"non-decayed: {opt.n_nodecay}")
+    # Static footprint. Nothing here is sharded — plain DDP replicates the
+    # full optimizer state on every rank (the whole cost of not running
+    # ZeRO-1), so this is the per-GPU charge before a single activation is
+    # allocated.
+    gb = 2 ** 30
+    param_bytes = sum(p.numel() * p.element_size() for p in params)
+    grad_bytes = param_bytes                       # grads share the param dtype
+    master_bytes, landing_bytes, _, _ = opt.state_breakdown()
+    # Both optimizers allocate their state on the first step(), so at this point
+    # the measured values are 0 and the report has to predict them.
+    mv_bytes = opt.expected_mv_bytes()
+    mom_bytes = opt.expected_muon_mom_bytes()
+    static_bytes = (param_bytes + grad_bytes + master_bytes + landing_bytes
+                    + mv_bytes + mom_bytes)
+    if master_process:
+        print(f"[{'ddp' if ddp else 'single'}] static footprint  (per GPU: "
+              f"params + grads + optimizer state)")
+        print(f"    parameters              : {param_bytes / gb:7.2f} GiB")
+        print(f"    gradients               : {grad_bytes / gb:7.2f} GiB")
+        print(f"    fp32 master weights     : {master_bytes / gb:7.2f} GiB")
+        print(f"    fp32 grad landing bufs  : {landing_bytes / gb:7.2f} GiB")
+        print(f"    Muon momentum           : {mom_bytes / gb:7.2f} GiB")
+        print(f"    Adam exp_avg + exp_avg_sq: {mv_bytes / gb:6.2f} GiB")
+        print(f"    static total            : {static_bytes / gb:7.2f} GiB", flush=True)
+
+    # ---- data ----
+    # Sharded by rank: each rank reads a disjoint, interleaved slice of the
+    # same shard list (see DataLoaderLite), so DDP's cross-rank gradient
+    # average is over world_size DIFFERENT windows, not world_size copies of
+    # the same one.
+    train_loader = DataLoaderLite(cfg.B, cfg.T, "train", args.data_root, device,
+                                  master=master_process, rank=ddp_rank,
+                                  world_size=ddp_world_size)
+    # rank=0, world_size=1 UNCONDITIONALLY, on every rank: val_loss_steps
+    # windows at fixed token offsets 0, BT, 2BT, ... Sharding this by
+    # ddp_rank/ddp_world_size like the train loader would make the reported
+    # val_loss a function of how many GPUs happen to be running (world_size=2
+    # averages 2x as many, differently-offset windows as world_size=1) —
+    # the two numbers would not be comparable across a GPU-count change, which
+    # is what produced the apparent val_loss jump on the 2-GPU -> 1-GPU
+    # resume. Every rank now redundantly evaluates the identical fixed
+    # windows, so the metric means the same thing regardless of world_size.
+    val_loader = DataLoaderLite(cfg.B, cfg.T, "val", args.data_root, device,
+                                master=master_process, rank=0, world_size=1)
+
+    # ---- checkpointing: manager + auto-resume ----
+    # Every rank independently loads the same checkpoint file from the same
+    # (assumed shared) filesystem — a read, so no race — which keeps model,
+    # optimizer state, and start_step identical across ranks without an
+    # explicit broadcast. Only rank 0 ever WRITES one (below).
+    start_step, ckpt = 0, None
+    last_val = float("nan")
+    if cfg.checkpointing or args.resume_from:
+        ckpt = CheckpointManager(cfg.ckpt_dir, cfg.ckpt_prefix, cfg.ckpt_keep)
+        if args.resume_from:
+            if master_process:
+                print(f"[resume] loading explicit checkpoint: {args.resume_from}")
+            start_step, resume_loss, last_val = ckpt.load(
+                args.resume_from, raw_model, opt, cfg, master_process)
+        else:
+            start_step, resume_loss, last_val = ckpt.load_latest(
+                raw_model, opt, cfg, master_process)
+        if start_step > 0:
+            if master_process:
+                print(f"[resume] at step {start_step} (loss {resume_loss:.6f}, "
+                      f"last val {last_val:.6f})")
+            # Skip the batches already consumed by steps [0, start_step).
+            # global_batch // (B*T*world_size) makes tokens-per-step (and so
+            # this target offset) independent of world_size, so a GPU-count
+            # change alone lands on the correct byte offset as long as B, T,
+            # and global_batch are unchanged (checked above).
+            train_loader.skip_batches(start_step * grad_accum)
+
+    # ---- CSV log ----
+    # Appended (not truncated) when resuming so a restart does not erase the
+    # history of the run it is continuing. Rank 0 only — every rank would
+    # otherwise open and write the same path.
+    log_file = None
+    if args.csv and master_process:
+        resuming = start_step > 0 and os.path.exists(args.csv)
+        log_file = open(args.csv, "a" if resuming else "w")
+        if not resuming:
+            log_file.write("step,train_loss,val_loss,tok_per_sec,grad_norm,dt_ms,lr,"
+                           "elapsed_min,muon_lr\n")
+        log_file.flush()
+
+    amp = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+    probes = [L for L in cfg.probe_layers if L < cfg.n_layers]
+    steady_dts = []
+    t_train_start = time.perf_counter()
+
+    # ---- Ctrl-C: finish the in-flight step, save a checkpoint, exit clean ----
+    # Registered here (not earlier) so a Ctrl-C during setup/data-loading still
+    # gets Python's default immediate-abort behaviour — only the training loop
+    # below gets the graceful path. A second Ctrl-C re-arms the default handler
+    # and re-raises, for when the checkpoint save itself needs interrupting.
+    interrupted = False
+
+    def _handle_sigint(signum, frame):
+        nonlocal interrupted
+        if interrupted:
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            raise KeyboardInterrupt
+        interrupted = True
+        if master_process:
+            print("\n[ctrl-c] finishing this step, then checkpointing and exiting "
+                  "(press again to abort immediately)...", flush=True)
+
+    # A terminal Ctrl-C is delivered to every process in the foreground
+    # process group, so under torchrun every rank sets `interrupted` at
+    # roughly the same time and the ranks stay in lockstep for the collective
+    # ops (DDP's grad all-reduce, the checkpoint barrier below) still ahead.
+    signal.signal(signal.SIGINT, _handle_sigint)
+
+    if master_process:
+        print("\nStarting training" +
+              (f" (resumed at step {start_step})" if start_step else "") + "...\n", flush=True)
+
+    for step in range(start_step, cfg.max_steps):
+        # ---------- periodic validation ----------
+        if cfg.val_loss_steps > 0 and (step % cfg.val_freq == 0 or step == cfg.max_steps - 1):
+            model.eval()
+            val_loader.reset()
+            acc = 0.0
+            with torch.no_grad(), amp:
+                for _ in range(cfg.val_loss_steps):
+                    x, y, _, _ = val_loader.next_batch()
+                    acc += model(x, y)[1].item() / cfg.val_loss_steps
+            model.train()
+            # val_loader is rank=0/world_size=1 on every rank (see its
+            # construction above), so every rank just computed the identical
+            # fixed windows redundantly — no cross-rank reduction needed, and
+            # val_loss no longer depends on how many ranks are running.
+            last_val = acc
+            if master_process:
+                print(f"  [val] step {step} | val_loss {last_val:.6f}", flush=True)
+
+        # ---------- training step ----------
+        t0 = time.perf_counter()
+        model.zero_grad(set_to_none=True)
+        # Fresh zeros rather than an in-place scale of last step's accumulator:
+        # multiplying by zero does NOT clear a NaN, so a single non-finite
+        # micro-batch would poison the accumulator for the rest of the run.
+        loss_accum = torch.zeros((), device=device)
+        micro_loss, micro_where = [], []
+
+        for micro in range(grad_accum):
+            x, y, shard, offset = train_loader.next_batch()
+            # DDP allreduces gradients on every .backward(); for grad_accum > 1
+            # only the LAST micro-batch of the step should pay for that — the
+            # others accumulate purely locally under no_sync().
+            sync_ctx = (nullcontext() if (not ddp or micro == grad_accum - 1)
+                       else model.no_sync())
+            with sync_ctx:
+                with amp:
+                    _, loss = model(x, y)
+                loss_accum += loss.detach()
+                # Per-micro-batch loss (--diag). The reported train loss is the mean
+                # over grad_accum micro-steps, so a single pathological micro-batch
+                # is diluted while its gradient dominates the accumulated norm
+                # almost undiluted — which is how a modest loss bump can sit next to
+                # a large norm spike. `micro_where` names the shard and token offset
+                # so the responsible bytes can be decoded offline.
+                #
+                # Costs a device sync per micro-step, so it is off by default.
+                if diag:
+                    micro_loss.append(loss.item())
+                    micro_where.append((shard, offset))
+                (loss / grad_accum).backward()
+
+        if ddp:
+            # Average each rank's accumulated (summed) loss BEFORE dividing by
+            # grad_accum, so loss_val is the mean over world_size*grad_accum
+            # micro-batches — the true loss behind this optimizer step, not
+            # just rank 0's slice. Also doubles as a sync point: a NaN on any
+            # rank poisons the average for every rank, so the check below
+            # raises everywhere instead of only where it originated (which
+            # would otherwise hang the survivors on the next collective).
+            dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
+        loss_val = loss_accum.item() / grad_accum
+        if math.isnan(loss_val) or math.isinf(loss_val):
+            if log_file:
+                log_file.close()
+            # Every rank reaches this identically (the all_reduce above is a
+            # sync point and NaN propagates through AVG), so every rank exits
+            # here together rather than the survivors hanging on the next
+            # collective.
+            if ddp:
+                dist.destroy_process_group()
+            raise SystemExit(f"ERROR: NaN/Inf loss at step {step}")
+
+        # Depth probe (--diag), read BEFORE the update, so these are the
+        # pre-update gradients. Compare the profile on a spiking step against a
+        # quiet one: a blow-up that begins mid-stack and grows toward the loss
+        # points at the forward, while a flat multiplicative lift across all
+        # depths points at the optimizer/update.
+        depth = []
+        if diag:
+            for L in probes:
+                a = raw_model.blocks[L].attn
+                gn = a.wQ.weight.grad
+                depth.append((L,
+                              gn.float().norm().item() if gn is not None else -1.0,
+                              a.q_gamma.abs().max().item(),
+                              a.wQ.weight.abs().max().float().item()))
+
+        grad_norm = clip_grads_(params, cfg.grad_clip).item()
+        lr = get_lr(step, cfg)                                   # AdamW group
+        muon_lr = get_lr(step, cfg, cfg.muon_lr, cfg.muon_min_lr)
+        opt.step(lr, muon_lr)
+
+        torch.cuda.synchronize()   # make the per-step timing reflect completed GPU work
+        dt = time.perf_counter() - t0
+        tok_per_sec = tokens_per_step / dt
+        elapsed_min = (time.perf_counter() - t_train_start) / 60.0
+        if step - start_step >= args.warmup_measure:
+            steady_dts.append(dt)
+
+        # ---------- per-step report (rank 0 only — every rank's numbers agree
+        # on loss/grad_norm post-allreduce, so printing from every rank would
+        # only duplicate lines, not add information) ----------
+        if master_process:
+            vs = "     -  " if last_val != last_val else f"{last_val:.6f}"
+            print(f"step {step:6d} | train {loss_val:.6f} | val {vs} "
+                  f"| tok/s {tok_per_sec:10.2f} | norm {grad_norm:.4f} "
+                  f"| dt {dt * 1e3:8.2f} ms | lr {lr:.4e} | mlr {muon_lr:.4e} "
+                  f"| mem {torch.cuda.max_memory_allocated() / gb:5.2f}/"
+                  f"{torch.cuda.max_memory_reserved() / gb:5.2f} GiB", flush=True)
+
+            if diag and micro_loss:
+                # Flag micro-batches that stand out against this step's own median.
+                med = sorted(micro_loss)[len(micro_loss) // 2]
+                marks = " ".join(("*" if m > med + 0.5 else "") + f"{m:.3f}"
+                                 for m in micro_loss)
+                print(f"  [micro] med {med:.4f} | {marks}")
+                for m, (ml, (shard, offset)) in enumerate(zip(micro_loss, micro_where)):
+                    if ml > med + 0.5:
+                        print(f"  [micro] OUTLIER m={m} loss {ml:.4f} (median {med:.4f}) "
+                              f"from shard {shard} token offset {offset}")
+            if diag and depth:
+                print("  [depth] ||g(wQ)|| by layer:"
+                      + "".join(f" L{L} {g:.3e}" for L, g, _, _ in depth)
+                      + " | max|q_gamma|:"
+                      + "".join(f" L{L} {gm:.4f}" for L, _, gm, _ in depth)
+                      + " | max|wQ|:"
+                      + "".join(f" L{L} {w:.4f}" for L, _, _, w in depth), flush=True)
+
+            if log_file:
+                log_file.write(f"{step},{loss_val:.6f},{last_val:.6f},{tok_per_sec:.6f},"
+                               f"{grad_norm:.6f},{dt * 1e3:.6f},{lr:.6f},{elapsed_min:.6f},"
+                               f"{muon_lr:.6f}\n")
+                log_file.flush()
+
+        if ddp:
+            # `interrupted` is set by this process's own SIGINT handler, so in
+            # principle one rank could see it before another (a signal isn't
+            # guaranteed simultaneous delivery, even though in practice a
+            # terminal Ctrl-C hits every rank in the same process group at
+            # once). Left asymmetric, it would desync the checkpoint barrier
+            # and the `break` below against ranks still issuing DDP's regular
+            # per-step collectives — a hang, not just a wrong answer. This
+            # runs unconditionally every step (never gated on a per-rank
+            # value), so it can never itself mismatch across ranks.
+            flag = torch.tensor(1.0 if interrupted else 0.0, device=device)
+            dist.all_reduce(flag, op=dist.ReduceOp.MAX)
+            interrupted = flag.item() > 0
+
+        # ---------- periodic checkpoint (also fires once on Ctrl-C) ----------
+        # Labeled step+1 (the next iteration to run) so load_latest resumes
+        # exactly where this save left off — no duplicated or skipped iteration.
+        # Captures weights + Adam m/v AFTER this step's update. Rank 0 only —
+        # every rank's raw_model/opt state is identical post-allreduce, so a
+        # second writer would only race the first on the same path. The
+        # barrier holds every rank at this point until the save (or its
+        # failure) is done, so a slow write can't let other ranks race ahead
+        # into the next step's collectives before the checkpoint is durable.
+        if cfg.checkpointing and ((step + 1) % cfg.ckpt_freq == 0
+                                  or step == cfg.max_steps - 1
+                                  or interrupted):
+            if master_process:
+                try:
+                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, cfg)
+                except Exception as e:                   # noqa: BLE001 — advisory
+                    print(f"[WARN] checkpoint save at step {step + 1} failed: {e}")
+            if ddp:
+                dist.barrier()
+
+        if interrupted:
+            if not cfg.checkpointing and master_process:
+                print(f"[ctrl-c] checkpointing is disabled (--no-checkpointing); "
+                      f"stopping at step {step} without saving.", flush=True)
+            break
+
+        # Reclaim the extra allocator high-water mark from the first couple of
+        # steps' one-time warmup allocations (autograd graph caches, etc.).
+        if step in (start_step, start_step + 1):
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+    torch.cuda.synchronize()
+    if log_file:
+        log_file.close()
+
+    # ---- steady-state summary ----
+    # dt/throughput are this rank's own timing; under DDP all ranks run in
+    # lockstep (grad sync + the checkpoint barrier), so rank 0's numbers are
+    # representative and `tokens_per_step` already counts every rank's tokens.
+    if steady_dts and master_process:
+        xs = sorted(steady_dts)
+        med = xs[len(xs) // 2]
+        mean = sum(xs) / len(xs)
+        alloc = torch.cuda.max_memory_allocated() / gb
+        reserved = torch.cuda.max_memory_reserved() / gb
+        # Measured now that Adam's m/v exist, rather than the startup estimate.
+        static = (param_bytes + grad_bytes + opt.state_bytes()) / gb
+        mode_str = f"DDP x{ddp_world_size}" if ddp else "single GPU"
+        print("\n" + "=" * 78)
+        print(f"  STEADY-STATE SUMMARY — {mode_str}  optimizer=muon+adamw  "
+              f"precision={args.precision}  "
+              f"rms+rope={'fused' if cfg.fused_rms_rope else 'eager'}")
+        print(f"  torch.optim.Muon  lr={cfg.muon_lr:g} "
+              f"adjust_lr_fn={cfg.muon_adjust_lr} "
+              f"momentum={cfg.muon_momentum:g} ns_steps={cfg.muon_ns_steps}")
+        print(f"  ({len(xs)} steps, first {args.warmup_measure} excluded)")
+        print("=" * 78)
+        print(f"  step time          : {med * 1e3:10.2f} ms  (mean {mean * 1e3:.2f})")
+        print(f"  throughput         : {tokens_per_step / med:10.1f} tok/s")
+        print(f"  tokens / step      : {tokens_per_step:10d}")
+        print(f"  static (p+g+opt)   : {static:10.2f} GiB  (per GPU)")
+        print(f"  activations (peak) : {max(alloc - static, 0.0):10.2f} GiB")
+        print(f"  peak mem allocated : {alloc:10.2f} GiB")
+        print(f"  peak mem reserved  : {reserved:10.2f} GiB")
+        print("=" * 78)
+
+    if ddp:
+        dist.destroy_process_group()
+    if master_process:
+        print("\n=== Training complete ===")
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+