train_390M.py

#!/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>                 largest micro-batch                (default 6)
#   BLU_GLOBAL_BATCH=<tokens> constant tokens per optimizer step; replaces
#                             the batch schedule                 (default off)
#   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.
# Every batch in the batch schedule (or --global-batch) is the TOTAL tokens
# per optimizer step across every rank (grad_accum = batch / (micro_B*T*
# world_size)), so a batch that splits evenly runs the same at any world size;
# resolve_batch_schedule() rounds one that doesn't. Data is rank-sharded
# (DataLoaderLite gives each rank a disjoint slice of every block 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 = 6                          # largest micro-batch; each batch stage picks its own <= this
    T: int = 4096                         # sequence length (<= context_length)
    global_batch: int =786432     # tokens per optimizer step, only when batch_schedule is empty
    grad_clip: float = 1.0
    weight_decay: float = 0.1

    # Run length in TOKENS. The batch changes during the run, so the step count
    # (max_steps, and with it the WSD phases) is derived from this in main().
    total_tokens: int = 14_000_000_000    # what the fixed-batch arms ran: 17,802 x 786,432
    # The train shards are cut into windows of data_grain * T tokens, the way
    # the fixed-batch arms' loader cut them with micro-batch B = data_grain
    # (6 in the 786K arm), so this run trains on the same tokens in the same
    # order whatever its batch, micro-batch or GPU count (see DataLoaderLite).
    data_grain: int = 6

    # Batch-size schedule by tokens consumed: ((start_tokens, tokens/step), ...).
    # From tripartite_fit.py (BATCH_SIZE_STUDY.md §4.3). For each target loss
    # L_k the fit gives B_opt(L_k), the batch that reaches L_k in the fewest
    # tokens, and E_min(L_k), how many tokens that takes. B_opt grows as the
    # loss falls, so each stage runs the B_opt of the target it is heading for,
    # starting where the run should have reached the one before it:
    #     tokens 0      -> 1.967B   B_opt(2.6906) = 578K
    #            1.967B -> 4.014B   B_opt(2.5540) = 741K
    #            4.014B -> 5.489B   B_opt(2.5054) = 782K
    #            5.489B -> 6.554B   B_opt(2.4807) = 778K
    #            6.554B -> 8.324B   B_opt(2.4501) = 800K
    #            8.324B -> end      B_opt(2.4376) = 807K  (the last target fitted)
    # Zhou et al. (arXiv:2601.05034, §4.3.2) also switch on tokens consumed and
    # keep the LR fixed (LR ~ sqrt(B) gave them nothing); train_390M_dynamicb.py
    # is that arm. THIS arm tests the alternative: the LR follows sqrt(batch)
    # at every switch (see lr_ref_batch below). The fitted values are
    # not whole numbers of sequences; resolve_batch_schedule() rounds them for
    # the GPU count and main() prints what runs. Empty = global_batch throughout.
    batch_schedule: tuple = ((0,             578_286),
                             (1_966_588_302, 741_051),
                             (4_014_246_572, 782_304),
                             (5_488_766_465, 777_500),
                             (6_554_208_302, 799_579),
                             (8_324_144_799, 807_193))

    # LR ~ sqrt(batch). Every step's LR — both groups, the whole WSD shape, the
    # min_lr floor included — is multiplied by sqrt(stage batch / lr_ref_batch),
    # so it steps up with the batch at each switch (807K vs 578K -> x1.18).
    # 0 = the first stage's batch: the run starts at max_lr / muon_lr exactly
    # like the fixed-LR dynamic arm, and the two arms differ only from the
    # first switch on. The batches are the rounded ones that run, not targets.
    lr_ref_batch: int = 0

    # 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.001
    min_lr: float = 0.0001                # 10% of max_lr
    
    muon_lr: float = 0.001
    muon_min_lr: float = 0.0001
    
    max_steps: int = 0                    # derived in main(): steps to consume total_tokens

    warmup_steps: int = 500
    stable_frac: float = 0.99             # stable phase ends at this fraction of max_steps
    stable_steps: int = 0                 # derived: int(max_steps * stable_frac) - warmup_steps


    # ---- 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                 # also save every N steps, on top of the
                                           # target-loss saves below; 0 = milestones only
    ckpt_keep: int = 40                    # retain the most-recent N interval checkpoints;
                                           # target-loss checkpoints are never pruned
    # 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_dynamic"
    ckpt_prefix: str = "blumodel_muon_dynamic"

    # ---- runtime knobs (not in the C++ struct) ----
    fused_rms_rope: bool = True
    activation_checkpoint: bool = False
    loss_chunks: int = 1
    # Milestones to record (step + elapsed hours) AND checkpoint on, as the
    # EMA-smoothed train loss passes each one — the same shared target-loss
    # set used across this project's batch-size ablation arms. Raw per-step
    # loss is noisy, so the comparison runs against an EMA, not loss_val
    # directly — same convention as batchsize_scaling.py's EMA_BETA=0.99.
    # A checkpoint is saved the step each one is first crossed, in addition
    # to the every-ckpt_freq-steps saves: this ablation's whole point is
    # comparing batch sizes at matched loss levels, so these are the
    # checkpoints that matter — keep-N pruning never deletes them.
    target_losses: tuple = (2.690595, 2.554042, 2.505404,
                            2.480678, 2.450056, 2.437589)
    target_loss_ema_beta: float = 0.99
    # 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)


@dataclass(frozen=True)
class BatchStage:
    start_tokens: int      # used from the first step that starts with this many tokens consumed
    target: int            # tokens/step the schedule asks for
    batch: int             # tokens/step actually run
    micro_B: int           # sequences per rank per micro-step
    grad_accum: int        # micro-steps per optimizer step


def resolve_batch_schedule(cfg, world_size):
    """cfg.batch_schedule -> [BatchStage], split over `world_size` GPUs.
    A step is micro_B * grad_accum sequences on every rank, so a batch must be
    a whole number of sequences per rank (a multiple of T * world_size) that
    factors as micro_B * grad_accum with micro_B <= cfg.B. A target that does
    is run exactly. Any other — the fitted B_opt values never are — is rounded
    to the nearest batch that splits with micro_B >= cfg.B / 2, because an odd
    per-rank count would otherwise force micro_B = 1 and run that stage at a
    fraction of the throughput.
    """
    if cfg.B < 1:
        raise SystemExit(f"ERROR: B must be >= 1, got {cfg.B}")
    schedule = cfg.batch_schedule or ((0, cfg.global_batch),)
    starts = [int(s) for s, _ in schedule]
    if starts[0] != 0 or any(a >= b for a, b in zip(starts, starts[1:])):
        raise SystemExit(f"ERROR: batch_schedule must start at 0 tokens and its starts "
                         f"must increase, got {list(schedule)}")
    seq = cfg.T * world_size                     # one sequence on every rank
    stages = []
    for start, target in schedule:
        target = int(target)
        if target <= 0:
            raise SystemExit(f"ERROR: batch_schedule batch must be > 0, got {target}")
        if target % seq == 0:
            n = target // seq
        else:
            floor = (cfg.B + 1) // 2
            # The range always holds a multiple of cfg.B (micro == cfg.B), so it
            # never comes up empty. Ties go to the larger batch.
            n = min((k for k in range(1, target // seq + cfg.B + 1)
                     if _micro_batch(k, cfg.B) >= floor),
                    key=lambda k: (abs(k * seq - target), -k))
        b = _micro_batch(n, cfg.B)
        stages.append(BatchStage(int(start), target, n * seq, b, n // b))
    return stages


def _micro_batch(n, B_max):
    """Largest micro-batch <= B_max that divides n sequences per rank."""
    return max(d for d in range(1, min(n, B_max) + 1) if n % d == 0)


def keep_batch_schedule(saved, stages, cfg, world_size):
    """On resume: the batches a checkpoint's run was running, split over this
    many GPUs, or None if the schedule config changed since (its starts or
    targets differ from `stages`) or a batch no longer splits evenly.
    resolve_batch_schedule() rounds for the GPU count, so a run restarted on a
    different count would otherwise switch to that count's rounding mid-run.
    Going down by a factor (8 -> 4, 2, 1) always keeps the batches.
    """
    if not saved or [[s.start_tokens, s.target] for s in stages] != [e[:2] for e in saved]:
        return None
    seq = cfg.T * world_size
    if any(b % seq for _, _, b in saved):
        return None
    kept = []
    for start, target, batch in saved:
        b = _micro_batch(batch // seq, cfg.B)
        kept.append(BatchStage(start, target, batch, b, batch // seq // b))
    return kept


def stage_at(stages, tokens):
    """The stage a step runs in, given the tokens consumed before it."""
    k = 0
    while k + 1 < len(stages) and tokens >= stages[k + 1].start_tokens:
        k += 1
    return stages[k]


def walk_schedule(stages, total_tokens, max_steps=None, start_tokens=0):
    """Step through the schedule by the training loop's own rule (stage_at,
    then consume that stage's batch) from start_tokens until total_tokens are
    consumed or max_steps are taken.
    -> ([[stage, n_steps], ...] in order, tokens consumed).
    """
    runs, tokens, steps = [], start_tokens, 0
    while tokens < total_tokens and (max_steps is None or steps < max_steps):
        s = stage_at(stages, tokens)
        if runs and runs[-1][0] is s:
            runs[-1][1] += 1
        else:
            runs.append([s, 1])
        tokens += s.batch
        steps += 1
    return runs, tokens


def lr_ref_batch(cfg, stages):
    """The batch max_lr / muon_lr are set for (ModelConfig.lr_ref_batch)."""
    return cfg.lr_ref_batch or stages[0].batch


def lr_batch_scale(batch, ref_batch):
    """sqrt(B) LR scaling: the factor a stage's batch puts on the WSD LR."""
    return math.sqrt(batch / ref_batch)


def print_batch_schedule(stages, runs, steps_label, ref_batch):
    """One row per stage: where it starts, the batch asked for and run, how it
    splits, the LR factor it runs at, and its steps in `runs` (from walk_schedule)."""
    steps_in = {id(s): n for s, n in runs}
    print(f"    {'from tokens':>15} {'target':>9} {'batch':>9} {'':>6} "
          f"{'micro x accum':>13} {'lr x':>7} {steps_label:>10}")
    for s in stages:
        off = f"{100.0 * (s.batch - s.target) / s.target:+.1f}%" if s.batch != s.target else ""
        print(f"    {s.start_tokens:>15,} {s.target:>9,} {s.batch:>9,} {off:>6} "
              f"{f'{s.micro_B} x {s.grad_accum}':>13} "
              f"{lr_batch_scale(s.batch, ref_batch):>7.4f} {steps_in.get(id(s), 0):>10}")


# =============================================================================
# 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]
#
# The token stream. Every shard is cut into windows of grain*T tokens from its
# start; each window that fits (+1 token for the shifted targets) is used, the
# shard's tail after the last one is not. That is exactly the set of tokens the
# fixed-batch arms' loader (train_390M_bs.py: rank r reads windows r, r+W, ...
# of B*T tokens) read with micro-batch B = grain. The stream is those windows'
# sequences, shard after shard, and each micro-step takes the next
# world_size*B sequences of it, rank r the r-th run of B, crossing into the
# next shard where one runs out. So which tokens a step reads depends only on
# how many tokens came before it — not on B, grad_accum or world_size: the
# batch schedule, and a resume on another GPU count, change how steps are cut
# from the stream, never the stream itself. After E tokens this run is E
# tokens into the stream, which is where the fixed-batch arms' ranks are
# spread around (theirs drift apart by a few M tokens across shards).
#
# The position (shard, offset) is the same on every rank, and a checkpoint
# stores it to resume from.
#
# 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, grain=None):
        assert split in ("train", "val")
        # B is also the largest micro-batch set_micro_batch() can switch to:
        # the buffers below are sized for it.
        self.B = self.B_max = B
        self.T = T
        self.split = split
        self.max_tokens = max_tokens_per_shard
        self.device = device
        # DDP: every rank reads the SAME stream but a DISJOINT run of each
        # micro-step's sequences (offset by rank), 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}")

        # Tokens of each shard the stream uses: its whole windows (see above).
        G = (grain or B) * T
        self.usable = [max(0, (min(os.path.getsize(p) // 2, max_tokens_per_shard) - 1) // G) * G
                       for p in self.shards]
        if not any(self.usable):
            raise RuntimeError(f"no {split} shard holds one window of {G} tokens")
        self.tok, self.tok_shard = None, None

        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.block = [(0, 0)] * self.NBUF        # stream position each slot's micro-step started at
        self.reset()

    # -- the stream ----------------------------------------------------------
    def _tokens(self, s):
        if s != self.tok_shard:
            self.tok, self.tok_shard = np.memmap(self.shards[s], dtype=np.uint16, mode="r"), s
        return self.tok

    def _advance(self, n):
        """Move n sequences along the stream. -> [(shard, offset, n_seq), ...],
        the pieces of shards they lie in (two when they cross a shard end)."""
        pieces = []
        while n > 0:
            if self.off >= self.usable[self.cur]:
                self.cur, self.off = (self.cur + 1) % len(self.shards), 0
                continue
            k = min(n, (self.usable[self.cur] - self.off) // self.T)
            pieces.append((self.cur, self.off, k))
            self.off += k * self.T
            n -= k
        return pieces

    def stream_offset(self, shard, off, tokens):
        """How far (shard, off) is past `tokens` tokens into the stream (signed,
        modulo one pass over the shards)."""
        d = (sum(self.usable[:shard]) + min(off, self.usable[shard]) - tokens) % sum(self.usable)
        return d - sum(self.usable) if d > sum(self.usable) // 2 else d

    def reset(self):
        self.seek(0, 0)

    # -- position (the same on every rank) -----------------------------------
    def state(self):
        """-> (shard_index, offset) of the next micro-step next_batch() will
        return. Prefetched ones it has not returned yet do not count as consumed."""
        return self.block[self.buf] if self.primed else (self.cur, self.off)

    def seek(self, shard, off):
        """Continue from a state(). Drops anything already prefetched."""
        self.copy_stream.synchronize()
        self.cur, self.off = shard, -(-off // self.T) * self.T     # a sequence boundary
        self.primed = False
        self.buf = 0

    def set_micro_batch(self, B):
        """Serve B sequences per rank per micro-step from the next one on."""
        if not 1 <= B <= self.B_max:
            raise ValueError(f"micro-batch {B} outside 1..{self.B_max}")
        if B != self.B:
            self.seek(*self.state())
            self.B = B

    def skip_tokens(self, n):
        """Fast-forward n tokens of the stream without touching the GPU."""
        self.seek(*self.state())
        self._advance(n // self.T)

    # -- prefetch ------------------------------------------------------------
    def _stage(self, idx):
        B = self.B
        BT = B * self.T

        # 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.block[idx] = (self.cur, self.off)
        self._advance(self.rank * B)                               # lower ranks' runs
        pieces = self._advance(B)                                  # this rank's
        self._advance((self.world_size - 1 - self.rank) * B)       # higher ranks'
        self.where[idx] = pieces[0][:2]
        at = 0
        for s, off, k in pieces:
            buf = np.asarray(self._tokens(s)[off: off + k * self.T + 1]).astype(np.int64)
            self.pin_x[idx][at: at + k * self.T].copy_(torch.from_numpy(buf[:-1]))
            self.pin_y[idx][at: at + k * self.T].copy_(torch.from_numpy(buf[1:]))
            at += k * self.T

        with torch.cuda.stream(self.copy_stream):
            self.dev_x[idx][:B].view(-1).copy_(self.pin_x[idx][:BT], non_blocking=True)
            self.dev_y[idx][:B].view(-1).copy_(self.pin_y[idx][:BT], non_blocking=True)
            self.h2d_done[idx].record(self.copy_stream)
        self.h2d_recorded[idx] = True

    def next_batch(self):
        """-> (x, y, shard_index, token_offset), x/y int64 [B,T] on device."""
        if not self.primed:
            # After a seek() the device buffers may still be read by kernels
            # queued before it; the refill waits for them.
            self.copy_stream.wait_stream(torch.cuda.current_stream())
            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.B], self.dev_y[i][:self.B],
                self.where[i][0], self.where[i][1])


# =============================================================================
# 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, run_cfg, milestones=None,
             progress=None):
        tmp = self._path(step) + ".tmp"
        torch.save({"step": step, "loss": loss, "last_val": last_val,
                    # T, total_tokens and the batch schedule as resolved for
                    # this GPU count: what the step count and the LR schedule
                    # were planned from. load() warns if they differ.
                    "run_cfg": run_cfg,
                    # tokens_seen and data_pos (the loader's state()): where
                    # the run is, independent of B and world_size, so a resume
                    # continues the data and the schedule exactly.
                    **(progress or {}),
                    # Target losses already crossed (loss -> (step, hours, tokens)).
                    # Without this, a resume re-initializes the EMA from
                    # scratch at whatever loss the model is already at, which
                    # falsely reports every target below that value as
                    # "reached" at the resume step, all at once.
                    "milestones": milestones or {},
                    "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:
            # Keep-N rotates only the interval checkpoints. A target-loss
            # checkpoint is labeled (crossing step + 1), and with interval
            # saves running alongside, the early ones would otherwise be the
            # first to age out.
            pinned = {m[0] + 1 for m in (milestones or {}).values()}
            rotating = [(s, p) for s, p in self._existing() if s not in pinned]
            for _, stale in rotating[:-self.keep]:
                os.remove(stale)

    def load_latest(self, model, opt, run_cfg, master=True):
        """-> (start_step, loss, last_val, milestones, progress).
        (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, run_cfg, master)

    def load(self, path, model, opt, run_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, milestones, progress), progress being
        {"tokens_seen", "data_pos", "batch_schedule"}, or {} for a checkpoint
        that predates them. The batch schedule is main()'s to compare: it
        keeps the checkpoint's batches when it can (keep_batch_schedule)."""
        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()
        saved = blob.get("run_cfg") or {}
        if master:
            mismatches = [f"{k} checkpoint={v!r} current={run_cfg[k]!r}"
                         for k, v in saved.items()
                         if k != "batch_schedule" and k in run_cfg and run_cfg[k] != v]
            if mismatches:
                print("[WARN] resume config differs from the checkpoint's:")
                for m in mismatches:
                    print(f"         {m}")
        print(f"[resume] loaded {path}")
        progress = ({"tokens_seen": int(blob["tokens_seen"]),
                     "data_pos": tuple(blob["data_pos"]),
                     "batch_schedule": saved.get("batch_schedule")}
                    if "data_pos" in blob else {})
        return (int(blob["step"]), float(blob["loss"]),
                float(blob.get("last_val", float("nan"))),
                blob.get("milestones", {}), progress)


# =============================================================================
# 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="largest micro-batch; each batch stage picks its own "
                         "<= this (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="a constant tokens per optimizer step for the whole run, "
                         "in place of the batch schedule (env: BLU_GLOBAL_BATCH)")
    ap.add_argument("--batch-schedule", nargs="+", default=None, metavar="TOKENS:BATCH",
                    help="batch stages as start-tokens:tokens-per-step pairs, the "
                         "first starting at 0, e.g. '0:589824 2e9:786432 8.3e9:819200'. "
                         "Default: the B_opt schedule in ModelConfig.batch_schedule.")
    ap.add_argument("--lr-ref-batch", type=int, default=None,
                    help="the batch (tokens/step) --max-lr / --muon-lr are set for; "
                         "every step's LR is multiplied by sqrt(batch / this). "
                         "Default 0 = the first stage's batch, so the LR is "
                         "max_lr / muon_lr until the first batch switch")
    ap.add_argument("--total-tokens", type=float, default=None,
                    help="run length in tokens (default 14e9); the step count is "
                         "derived from it and the batch schedule")
    ap.add_argument("--data-grain", type=int, default=None,
                    help="cut the train shards into windows of this many sequences, "
                         "as the fixed-batch arms' loader did with that micro-batch "
                         "(default 6, the 786K arm's)")
    ap.add_argument("--realign-data", action="store_true",
                    help="on resume from a checkpoint an earlier loader wrote, move "
                         "the data position back to tokens-seen into the stream, "
                         "where the fixed-batch arms are; the tokens in between are "
                         "seen twice. Default: continue from the stored position.")
    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)")
    ap.add_argument("--target-losses", type=float, nargs="+", default=None,
                    help="loss milestones to record (step + elapsed hours) and "
                         "checkpoint on, as the EMA-smoothed train loss passes "
                         "each one; default is the shared target-loss set used "
                         "across this project's batch-size ablation arms.")
    ap.add_argument("--target-loss-ema-beta", type=float, default=None,
                    help="EMA smoothing beta applied before comparing against "
                         "--target-losses (default 0.99)")
    # ---- 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, then tokens (consumed "
                         "before the step, = step x batch for a fixed batch) and "
                         "global_batch.")
    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,
                    help="also checkpoint every N steps, alongside the target-loss "
                         "checkpoints (default 1000; 0 = target losses only)")
    ap.add_argument("--ckpt-keep", type=int, default=None,
                    help="keep the newest N interval checkpoints (default 40); "
                         "target-loss checkpoints are never pruned")
    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", "data_grain", "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", "target_losses", "target_loss_ema_beta",
              "lr_ref_batch"):
        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)

    if args.total_tokens is not None:
        cfg.total_tokens = int(args.total_tokens)
    if args.global_batch is not None and args.batch_schedule is not None:
        raise SystemExit("ERROR: pass --global-batch (constant) or --batch-schedule, not both")
    if args.global_batch is not None:
        cfg.batch_schedule = ()
    if args.batch_schedule is not None:
        try:
            cfg.batch_schedule = tuple((int(float(t)), int(float(b)))
                                       for t, b in (s.split(":") for s in args.batch_schedule))
        except ValueError:
            raise SystemExit(f"ERROR: --batch-schedule takes TOKENS:BATCH pairs, "
                             f"got {' '.join(args.batch_schedule)}")

    # 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, 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}")
    # Batch divisibility over B*T*world_size is resolve_batch_schedule()'s job:
    # every stage it returns splits exactly.
    if cfg.total_tokens < 1:
        bad.append(f"total_tokens must be >= 1, got {cfg.total_tokens}")
    if cfg.lr_ref_batch < 0:
        bad.append(f"lr_ref_batch must be >= 0 (0 = first stage's batch), got {cfg.lr_ref_batch}")
    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

    # ---- batch schedule -> stages for this GPU count -> step count ----
    # Each step's tokens = its stage's batch = micro_B * T * grad_accum *
    # world_size. The step count follows from total_tokens by walking the
    # schedule with the training loop's own rule, so the WSD phases (in steps)
    # end where the tokens run out. --max-steps (a benchmark cap) keeps its own.
    stages = resolve_batch_schedule(cfg, ddp_world_size)
    stage_runs, planned_tokens = walk_schedule(stages, cfg.total_tokens, args.max_steps)
    if args.max_steps is None:
        cfg.max_steps = sum(n for _, n in stage_runs)
        cfg.stable_steps = int(cfg.max_steps * cfg.stable_frac) - cfg.warmup_steps
    validate_config(cfg, master_process)
    # What a checkpoint records about the plan. On resume, T / total_tokens are
    # checked, and the batches are kept if they still split (keep_batch_schedule).
    run_cfg = {"T": cfg.T, "total_tokens": cfg.total_tokens,
               "batch_schedule": [[s.start_tokens, s.target, s.batch] for s in stages]}

    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}  (micro-batch set per stage)")
        print(f"  total tokens    : {cfg.total_tokens:,}  ({planned_tokens:,} in "
              f"{cfg.max_steps} steps)")
        print(f"  batch schedule  : by tokens consumed, world_size={ddp_world_size}; "
              f"LR x sqrt(batch / {lr_ref_batch(cfg, stages):,})")
        print_batch_schedule(stages, stage_runs, "steps", lr_ref_batch(cfg, stages))
        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})")
        tl_str = ", ".join(f"{t:g}" for t in cfg.target_losses)
        print(f"  target losses   : [{tl_str}]  (ema beta {cfg.target_loss_ema_beta:g}, "
              f"checkpointed on crossing, never pruned)")
        ck_str = f"every {cfg.ckpt_freq} steps" if cfg.ckpt_freq > 0 else "off"
        print(f"  interval ckpts  : {ck_str}  (keep newest {cfg.ckpt_keep})")
        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, grain=cfg.data_grain)
    # 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.
    #
    # Batch size is hardcoded to 4 here (NOT cfg.B) for the same reason: the
    # batch-size ablation varies cfg.B/global_batch across runs, and if val
    # windows scaled with it (B*T tokens per val_loss_steps iteration) the
    # val_loss numbers would stop being comparable across ablation arms. A
    # fixed val batch size means every run evaluates the same fixed set of
    # tokens (4*T*val_loss_steps), regardless of what B that run trains with.
    val_loader = DataLoaderLite(4, 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")
    resumed_milestones = {}
    tokens_seen = 0          # consumed by steps [0, step); picks each step's stage
    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, resumed_milestones, progress = ckpt.load(
                args.resume_from, raw_model, opt, run_cfg, master_process)
        else:
            start_step, resume_loss, last_val, resumed_milestones, progress = ckpt.load_latest(
                raw_model, opt, run_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})")
            if progress:
                # The checkpoint says where the data and the schedule are; both
                # are independent of B and world_size, so this is exact.
                tokens_seen = progress["tokens_seen"]
                train_loader.seek(*progress["data_pos"])
                # 0 when this loader wrote the checkpoint: E tokens in, the run
                # is E tokens into the stream. An earlier loader dropped more
                # of each shard's tail, which left the run ahead of that.
                ahead = train_loader.stream_offset(*progress["data_pos"], tokens_seen)
                if args.realign_data and ahead:
                    train_loader.reset()
                    train_loader.skip_tokens(tokens_seen)
                    if master_process:
                        print(f"[resume] --realign-data: moved the data position back "
                              f"{ahead:,} tokens to {tokens_seen:,} tokens into the "
                              f"stream; those tokens are seen a second time")
                elif master_process and ahead:
                    print(f"[resume] the data position is {ahead:+,} tokens from "
                          f"{tokens_seen:,} tokens into the stream (left by an "
                          f"earlier loader); it stays at that offset from here")
                # Keep running the batches the run started with when this GPU
                # count can (always true going 8 -> 4, 2, 1), rather than this
                # count's own rounding: changing GPUs then changes nothing but
                # how each step is split.
                kept = keep_batch_schedule(progress["batch_schedule"], stages, cfg,
                                           ddp_world_size)
                if kept is not None:
                    if master_process and [s.batch for s in kept] != [s.batch for s in stages]:
                        print(f"[resume] keeping the batches this run started with "
                              f"(they split over {ddp_world_size} GPU(s)), not this "
                              f"GPU count's rounding")
                    stages = kept
                elif master_process:
                    print(f"[WARN] the batches this run started with no longer apply "
                          f"(the batch schedule config changed, or a batch does not "
                          f"split over {ddp_world_size} GPU(s)); the rest of the run "
                          f"follows the schedule above")
            else:
                # A checkpoint from a trainer without a batch schedule: assume
                # steps [0, start_step) ran this schedule, and go that many
                # tokens into the stream.
                _, tokens_seen = walk_schedule(stages, float("inf"), start_step)
                train_loader.skip_tokens(tokens_seen)
                if master_process:
                    print(f"[resume] checkpoint has no data position; replayed this "
                          f"schedule to step {start_step} ({tokens_seen:,} tokens)")
            if master_process:
                print(f"[resume] {tokens_seen:,} tokens consumed", flush=True)
            # Steps left = what the batches from here need to reach
            # total_tokens. The plan printed above assumed this GPU count's
            # rounding from step 0; when the run continues with other batches
            # (kept from the checkpoint, or changed) this keeps it, and the LR
            # decay, ending at total_tokens.
            rest, planned_tokens = walk_schedule(stages, cfg.total_tokens,
                                                 start_tokens=tokens_seen)
            if args.max_steps is None and start_step + sum(n for _, n in rest) != cfg.max_steps:
                cfg.max_steps = start_step + sum(n for _, n in rest)
                cfg.stable_steps = int(cfg.max_steps * cfg.stable_frac) - cfg.warmup_steps
                if master_process:
                    print(f"[resume] with the batches this run continues with, it ends "
                          f"at step {cfg.max_steps} ({planned_tokens:,} tokens; WSD "
                          f"stable until step {cfg.warmup_steps + cfg.stable_steps}):")
                    print_batch_schedule(stages, rest, "steps left",
                                         lr_ref_batch(cfg, stages))
            run_cfg["batch_schedule"] = [[s.start_tokens, s.target, s.batch] for s in stages]

    # Fixed after resume, which may have swapped in the batches the run started
    # with, so a restarted run scales its LR off the same first-stage batch.
    ref_batch = lr_ref_batch(cfg, stages)

    # ---- 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,tokens,global_batch\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 = []
    steady_tokens = 0
    stage = None      # the batch stage the previous step ran
    t_train_start = time.perf_counter()
    ema_loss = None   # EMA-smoothed train loss, checked against target_losses
    # Milestones already crossed, restored from the checkpoint on resume so a
    # restart doesn't re-discover them all at once (see CheckpointManager.save).
    milestones = dict(resumed_milestones)   # target loss -> (step, elapsed_hours, tokens)
    # Remaining targets, highest loss first (the order the EMA passes them in
    # as it decreases) — excludes anything already in `milestones`.
    pending_targets = [t for t in sorted(cfg.target_losses, reverse=True)
                       if t not in milestones]

    # ---- 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)

        # ---------- batch stage ----------
        # Chosen by tokens consumed, the rule walk_schedule() planned max_steps
        # with. tokens_seen is a function of the step and the schedule, so every
        # rank switches on the same step. A switch changes how many micro-steps
        # of what size the next steps read, and rescales the LR by sqrt(batch)
        # (lr_scale, applied from this step on); momentum and the data stream
        # carry straight on.
        new_stage = stage_at(stages, tokens_seen)
        if stage is None or (new_stage.batch, new_stage.micro_B) != (stage.batch, stage.micro_B):
            train_loader.set_micro_batch(new_stage.micro_B)
            grad_accum = new_stage.grad_accum
            tokens_per_step = new_stage.batch
            lr_scale = lr_batch_scale(new_stage.batch, ref_batch)
            if master_process:
                prev = f"{stage.batch:,} -> " if stage is not None else ""
                print(f"[batch] step {step} | {tokens_seen:,} tokens | batch {prev}"
                      f"{new_stage.batch:,} tokens/step ({new_stage.micro_B} x "
                      f"{new_stage.grad_accum} micro-steps per rank) | LR x "
                      f"{lr_scale:.4f} (sqrt of batch / {ref_batch:,})", flush=True)
        stage = new_stage

        # ---------- 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}")

        # EMA-smoothed train loss vs. the pending milestones. loss_val is
        # already all-reduced across ranks above, so ema_loss (and the
        # milestones it crosses) comes out identical on every rank without an
        # extra collective.
        beta = cfg.target_loss_ema_beta
        ema_loss = loss_val if ema_loss is None else beta * ema_loss + (1 - beta) * loss_val
        elapsed_hr = (time.perf_counter() - t_train_start) / 3600.0
        target_reached = False
        while pending_targets and ema_loss <= pending_targets[0]:
            t = pending_targets.pop(0)
            # tokens_seen, not +this step: E = B*S in the fixed-batch arms.
            milestones[t] = (step, elapsed_hr, tokens_seen)
            target_reached = True
            if master_process:
                print(f"\n[target] loss <= {t:.6f} reached at step {step} "
                      f"({elapsed_hr:.3f} hrs, {tokens_seen:,} tokens)\n", flush=True)

        # 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) * lr_scale                        # AdamW group
        muon_lr = get_lr(step, cfg, cfg.muon_lr, cfg.muon_min_lr) * lr_scale
        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)
            steady_tokens += tokens_per_step

        # ---------- 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},{tokens_seen},{tokens_per_step}\n")
                log_file.flush()

        # This step's tokens are consumed. Counted here, after the report and
        # before the checkpoint, so a checkpoint labeled step+1 records the
        # tokens consumed by steps [0, step+1).
        tokens_seen += tokens_per_step

        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

        # ---------- checkpoint on target-loss crossing, every ckpt_freq steps,
        #            and at end-of-run / 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.
        #
        # Interval saves land on round labels (step+1 = 1000, 2000, ...) and
        # bound how much a crash can cost; the target-loss saves are the ones
        # the ablation compares, and CheckpointManager.save never prunes them.
        # Both triggers are functions of `step` and the all-reduced loss_val,
        # so they agree on every rank without a collective of their own.
        interval_due = cfg.ckpt_freq > 0 and (step + 1) % cfg.ckpt_freq == 0
        if cfg.checkpointing and (target_reached
                                  or interval_due
                                  or step == cfg.max_steps - 1
                                  or interrupted):
            if master_process:
                try:
                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, run_cfg,
                              milestones, {"tokens_seen": tokens_seen,
                                           "data_pos": train_loader.state()})
                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()

    # ---- target-loss milestone summary ----
    if master_process:
        print("\n" + "=" * 68)
        print("  TARGET LOSS MILESTONES")
        print("=" * 68)
        print(f"  {'target loss':>12}   {'step':>10}   {'hours':>10}   {'tokens':>15}")
        for t in sorted(cfg.target_losses, reverse=True):
            if t in milestones:
                step_hit, hrs, tok_hit = milestones[t]
                print(f"  {t:>12.6f}   {step_hit:>10}   {hrs:>10.3f}   {tok_hit:>15,}")
            else:
                print(f"  {t:>12.6f}   {'not reached':>10}   {'-':>10}   {'-':>15}")
        print("=" * 68)

    # ---- 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.
    # Step time mixes batch stages; throughput is over every step measured.
    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         : {steady_tokens / sum(xs):10.1f} tok/s")
        print(f"  tokens / step      : {tokens_per_step:10d}  (last stage)")
        print(f"  tokens consumed    : {tokens_seen:,}")
        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()