extras / train_390M_7e4_mbs8.py
Copy files from models/jayasriii/390M_7e4_mbs8
f93e631 verified 28 days ago
90.1 kB
#!/usr/bin/env python3
# =============================================================================
# train_muon.py — the MUON arm of train.py. Same model, same data order, same
# schedule shape, same seed: the ONLY thing that changes is how a gradient
# becomes a weight update.
#
# Fork discipline — everything outside section 5 is byte-identical to train.py,
# so `diff train.py train_muon.py` is exactly the experiment being run. If you
# edit a kernel or the model in one file, port it to the other or the comparison
# stops meaning anything.
#
# The optimizer itself is torch.optim.Muon — stock, not reimplemented. This file
# supplies only the two things it does not do for you: the ROUTING (which
# parameters it may have) and the fp32 MASTER WEIGHTS (it updates the master,
# not the bf16 param). Requires a torch new enough to ship it; main() checks and
# says so.
#
# What Muon does: AdamW rescales each weight ENTRY by its own gradient history.
# Muon rescales each weight MATRIX by its singular values — it takes the
# momentum buffer M and replaces it with the nearest semi-orthogonal matrix
# (all singular values driven to ~1) via a 5-step Newton-Schulz iteration, then
# steps along that. Directions the gradient has barely explored get the same
# step size as the dominant ones, which is where the reported speedup comes
# from. Only 2D hidden weights can be orthogonalized, so this is a HYBRID:
#
#     Muon   : every 2D Linear weight inside the blocks   (75.5M params, 66%)
#     AdamW  : the token embedding / lm_head              (38.6M params)
#              + RMSNorm weights + QK-norm gammas (1D)    (22.3k params)
#
# The embedding is deliberately NOT given to Muon: its rows are per-token, not a
# linear map between two feature spaces, and orthogonalizing across the vocab
# axis mixes unrelated tokens. That split is the standard Muon recipe, and it
# keeps the AdamW arm's exact update on 34% of the parameters — so a difference
# in the loss curve is attributable to the hidden matrices alone.
#
# Everything train_ddp.py / train_zero1.py import from blu_common / blu_data /
# blu_kernels / blu_model is inlined here, in that order, so this file runs on
# its own:
#     section 1  fused QK-norm + RoPE (Triton)  <- blu_kernels.py
#     section 2  config + WSD schedule          <- blu_common.py
#     section 3  data loader                    <- blu_data.py
#     section 4  model                          <- blu_model.py
#     section 5  optimizer / clipping / checkpointing
#     section 6  training
#
# Architecture:
#   token embedding
#   -> N x { RMSNorm -> [fused QK-norm + RoPE] -> causal GQA (flash) -> Wo -> +x
#            RMSNorm -> SwiGLU FFN (gate_up -> silu(gate)*up -> down)  -> +x }
#   -> final RMSNorm
#   -> lm_head (weight-tied with the token embedding)
#   -> cross-entropy
#   -> backward -> AdamW (decoupled weight decay) on a WSD (warmup / stable /
#      cosine-decay) schedule
#
# Kept from the AdamW arm, unchanged, because they are confounds if they move:
#   * fp32 master weights for every non-fp32 param — Muon updates the MASTER,
#     reading the same fp32 landing buffer AdamW reads, so both optimizers see
#     bit-identical gradients;
#   * decoupled weight decay 0.1 on dim>=2 params, applied to the master;
#   * ONE global-norm gradient clip over every parameter, BEFORE the optimizer.
#     (Muon's orthogonalization makes the update scale-invariant, so clipping is
#     nearly a no-op for its group — but it still gates the AdamW group, and
#     removing it would change two things at once.)
#
# Memory, at the 114M config this repo actually runs (printed at startup):
#
#     parameters (bf16 Linear + fp32 embed/norms)  0.28 GiB
#     gradients (same dtypes as the params)        0.28 GiB
#     fp32 master weights (the bf16 Linears)       0.28 GiB
#     fp32 gradient landing buffers                0.28 GiB
#     Muon momentum   (75.5M params, fp32 x1)      0.28 GiB
#     AdamW m + v     (38.7M params, fp32 x2)      0.29 GiB
#     -------------------------------------------------------
#     static total                                 1.70 GiB   (AdamW arm: 1.98)
#
# Muon carries ONE state tensor per param where Adam carries two, so the saving
# scales with how much of the model is hidden matrices: 0.28 GiB here (14% of
# static), because this config spends a third of its parameters on a 50304-row
# embedding that stays on AdamW. At the 1.73B config in the original header the
# same split saves ~5.5 GiB.
#
# Activations dominate either way (63.96 GiB measured at B=16 x T=4096), so the
# OOM escape hatches are unchanged: --B 8, --loss-chunks 4, --activation-checkpoint.
#
# Fidelity notes (why a run here is statistically but not bit-wise identical):
#   * init draws come from torch's Philox, not the C++ single advancing
#     std::mt19937, so the two runs start from different samples of the same
#     distributions;
#   * QK-norm reads the projection output in bf16 here (autocast) where the C++
#     casts it up to fp32 first; the reduction and the rotation are fp32 in both.
#
# Determinism: two runs of the same launch line on the same GPU produce
# bit-identical losses. Seeding is only half of that — the default CUDA kernels
# for the embedding backward and the flash/mem-efficient attention backward
# accumulate partial sums with atomicAdd, so the order they land in (and with it
# the fp32 rounding) follows the scheduler, not the program. --determinism
# strict (the default) forces the ordered variants of those kernels; it costs a
# few percent of step time, which --determinism off buys back.
#
# Env overrides (the bluscript_single.cpp names, so an existing launch line keeps
# working; the equivalent CLI flag wins if both are given):
#   BLU_DEVICE=<idx>          CUDA device to train on            (default 0)
#   BLU_MAX_STEPS=<n>         short benchmark run (rescales WSD) (default off)
#   BLU_B=<n>                 micro-batch size                   (default 4)
#   BLU_GLOBAL_BATCH=<tokens> tokens per optimizer step          (default 524288)
#   BLU_DIAG=1                per-micro-batch loss + depth probes
#
# Multi-GPU: DDP, one process per GPU, via torchrun. RANK / LOCAL_RANK /
# WORLD_SIZE (set by torchrun, not this file) switch the bootstrap in main();
# their absence is the single-process path above, unchanged. --device /
# BLU_DEVICE are ignored under torchrun — LOCAL_RANK picks each rank's GPU.
# --global-batch is still the TOTAL tokens per optimizer step across every
# rank (grad_accum = global_batch / (B*T*world_size)), so the same
# --global-batch reproduces the same schedule at any world size. Data is
# rank-sharded (DataLoaderLite reads a disjoint, rank-offset slice of the
# shards); gradients and the reported train/val loss are averaged across
# ranks; only rank 0 writes the CSV and checkpoints. Same launch line as
# single-GPU otherwise:
#   torchrun --standalone --nproc_per_node=8 train_muon.py
#
# Learning rate — the one knob that decides whether the comparison is fair.
#
# Muon's update is spectrally normalized, so its natural scale has nothing to do
# with AdamW's. torch.optim.Muon ships both published conventions behind its
# `adjust_lr_fn`; --muon-adjust-lr passes the string straight through:
#
#   match_rms_adamw      Moonshot's: gamma <- 0.2 * gamma * sqrt(max(A,B)). The
#     (DEFAULT HERE)     Newton-Schulz output has singular values ~1, so its
#                        entries have RMS 1/sqrt(max(A,B)); this lifts them to
#                        RMS 0.2, which is what a converged AdamW update looks
#                        like. The consequence: --muon-lr is on the SAME SCALE
#                        as --max-lr, and the default (2e-4, equal to the AdamW
#                        arm's max_lr) gives both arms the same step size on the
#                        first run — only the update DIRECTION differs. That is
#                        the controlled experiment; run it before tuning.
#   original             Keller's modded-nanogpt scaling, sqrt(max(1, A/B)).
#                        Entries stay at RMS 1/sqrt(max(A,B)), so an equivalent
#                        LR is ~5.5x larger at d_model=768 (0.2*sqrt(768)).
#                        Use it to reproduce published Muon numbers.
#
# NOTE: torch's own default for adjust_lr_fn is `original`. This file defaults
# to `match_rms_adamw` instead, because an LR-matched first run is the point of
# having two arms. If you cite a number from here, cite the flag with it.
#
# Muon usually tolerates 2-4x more LR than AdamW once it is working, so the
# honest protocol is: baseline run, matched run, then a short --max-steps 200
# sweep over --muon-lr before claiming a win.
#
# Launch (the two arms; separate CSV + checkpoint dirs so they never collide):
#   python3 train.py                                  # AdamW baseline
#   python3 train_muon.py                             # Muon, LR-matched
#   python3 train_muon.py --muon-lr 6e-4              # after a sweep
#   python3 train_muon.py --max-steps 200 --muon-lr 4e-4   # short sweep point
# =============================================================================

import argparse
import glob
import math
import os
import random
import re
import signal
import time
from contextlib import nullcontext
from dataclasses import dataclass

# cuBLAS reads CUBLAS_WORKSPACE_CONFIG once, when it creates its handle on the
# first GEMM, and torch refuses to run fp32 cuBLAS GEMMs under
# use_deterministic_algorithms() unless it is one of the two allowed values. It
# has to be in the environment before torch touches CUDA, hence: before the
# import, not inside main(). An explicit setting from the launch line wins.
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.checkpoint import checkpoint

try:
    import triton
    import triton.language as tl
    _HAS_TRITON = True
except ImportError:                                   # eager fallback below
    _HAS_TRITON = False


# =============================================================================
# 1. Fused QK-norm (RMSNorm over head_dim) + RoPE, in Triton.
#
# PyTorch counterpart of the C++ `autograd::fused_rms_rope`. One kernel does,
# per (batch, head, position) row:
#
#     s = rsqrt(mean(x^2) + eps)          # RMS over head_dim
#     y = x * s * gamma                   # learnable per-head_dim scale
#     o = rope_neox(y, cos[t], sin[t])    # NeoX / Llama half-split rotation
#
# Attention is deliberately NOT fused in here — it stays a separate op
# (F.scaled_dot_product_attention on the flash backend), mirroring the C++ split
# between `fused_rms_rope` and `gqa_fused_flash_attn_forward/backward`.
#
# Why fuse: the eager formulation (`rms_rope_eager`) materialises an fp32 copy of
# x, the squared tensor, the normalised tensor, the gamma-scaled tensor and a
# `cat` of the two rotated halves — roughly 6 extra [B,H,T,D] buffers held for
# backward. The fused version saves only x (which the projection's backward
# already holds) plus one fp32 scalar per row.
#
# Layout notes
#   - x is [B,H,T,D] and may be a non-contiguous transpose of a [B,T,H,D] view;
#     arbitrary (b,h,t) strides are honoured, D must be stride-1.
#   - D must be a power of two (the C++ kernel requires 64 or 128).
#   - cos/sin are [T_ctx, D//2] fp32, as produced by `build_rope_cache`.
#   - the row index is r = (b*H + h)*T + t, so the output is contiguous
#     [B,H,T,D] in that same order.
# =============================================================================
if _HAS_TRITON:

    @triton.jit
    def _fused_rms_rope_fwd(
        X, GAMMA, COS, SIN, OUT, RSTD,
        sx_b, sx_h, sx_t,                   # input strides, in elements
        T, H, N,                            # seq len, n_heads, total rows
        eps,
        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
    ):
        pid = tl.program_id(0)
        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)          # [R]
        mask_r = rows < N

        # r = (b*H + h)*T + t  ->  recover t (for RoPE) and the (b, h) base.
        t = rows % T
        bh = rows // T
        b = bh // H
        h = bh % H
        base = b * sx_b + h * sx_h + t * sx_t                  # [R]

        dh = tl.arange(0, HALF)                               # [HALF]
        p1 = X + base[:, None] + dh[None, :]
        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)

        # RMS over the FULL head_dim (both halves), fp32 accumulation.
        ss = tl.sum(x1 * x1, 1) + tl.sum(x2 * x2, 1)          # [R]
        rstd = 1.0 / tl.sqrt(ss / D + eps)                    # [R]

        g1 = tl.load(GAMMA + dh).to(tl.float32)
        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
        y1 = x1 * rstd[:, None] * g1[None, :]
        y2 = x2 * rstd[:, None] * g2[None, :]

        cp = t[:, None] * HALF + dh[None, :]
        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)

        # NeoX half-split rotation.
        o1 = y1 * c - y2 * s
        o2 = y2 * c + y1 * s

        op = OUT + rows[:, None] * D + dh[None, :]
        tl.store(op, o1.to(OUT.dtype.element_ty), mask=mask_r[:, None])
        tl.store(op + HALF, o2.to(OUT.dtype.element_ty), mask=mask_r[:, None])
        tl.store(RSTD + rows, rstd, mask=mask_r)

    # -------------------------------------------------------------------------
    # backward
    #
    #   RoPE is an orthogonal rotation, so its vjp is the inverse rotation:
    #       dy1 = do1*cos + do2*sin ;  dy2 = do2*cos - do1*sin
    #   then, with n = x * rstd (the normalised row),
    #       dgamma += dy * n
    #       dn      = dy * gamma
    #       dx      = rstd * (dn - n * mean(dn * n))
    # -------------------------------------------------------------------------
    @triton.jit
    def _fused_rms_rope_bwd(
        X, GAMMA, COS, SIN, DOUT, RSTD, DX, DGP,
        sx_b, sx_h, sx_t,                   # input strides
        sd_b, sd_h, sd_t,                   # dx strides
        T, H, N,
        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
    ):
        pid = tl.program_id(0)
        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
        mask_r = rows < N

        t = rows % T
        bh = rows // T
        b = bh // H
        h = bh % H
        base = b * sx_b + h * sx_h + t * sx_t

        dh = tl.arange(0, HALF)
        p1 = X + base[:, None] + dh[None, :]
        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)

        rstd = tl.load(RSTD + rows, mask=mask_r, other=1.0)
        n1 = x1 * rstd[:, None]
        n2 = x2 * rstd[:, None]

        dp = DOUT + rows[:, None] * D + dh[None, :]
        do1 = tl.load(dp, mask=mask_r[:, None], other=0.0).to(tl.float32)
        do2 = tl.load(dp + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)

        cp = t[:, None] * HALF + dh[None, :]
        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)

        # inverse rotation
        dy1 = do1 * c + do2 * s
        dy2 = do2 * c - do1 * s

        # per-program partial dgamma; summed by the caller (deterministic, no atomics)
        tl.store(DGP + pid * D + dh, tl.sum(dy1 * n1, 0))
        tl.store(DGP + pid * D + HALF + dh, tl.sum(dy2 * n2, 0))

        g1 = tl.load(GAMMA + dh).to(tl.float32)
        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
        dn1 = dy1 * g1[None, :]
        dn2 = dy2 * g2[None, :]

        dot = (tl.sum(dn1 * n1, 1) + tl.sum(dn2 * n2, 1)) / D      # [R]
        dx1 = rstd[:, None] * (dn1 - n1 * dot[:, None])
        dx2 = rstd[:, None] * (dn2 - n2 * dot[:, None])

        dbase = b * sd_b + h * sd_h + t * sd_t
        q1 = DX + dbase[:, None] + dh[None, :]
        tl.store(q1, dx1.to(DX.dtype.element_ty), mask=mask_r[:, None])
        tl.store(q1 + HALF, dx2.to(DX.dtype.element_ty), mask=mask_r[:, None])


_BLOCK_R = 16
_NUM_WARPS = 4


def _launch_fwd(x, gamma, cos, sin, eps):
    B, H, T, D = x.shape
    N = B * H * T
    out = torch.empty((B, H, T, D), dtype=x.dtype, device=x.device)
    rstd = torch.empty(N, dtype=torch.float32, device=x.device)
    grid = (triton.cdiv(N, _BLOCK_R),)
    _fused_rms_rope_fwd[grid](
        x, gamma, cos, sin, out, rstd,
        x.stride(0), x.stride(1), x.stride(2),
        T, H, N, eps,
        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
    )
    return out, rstd


def _launch_bwd(x, gamma, cos, sin, dout, rstd):
    B, H, T, D = x.shape
    N = B * H * T
    dout = dout.contiguous()
    dx = torch.empty_like(x)
    nprog = triton.cdiv(N, _BLOCK_R)
    dgp = torch.empty((nprog, D), dtype=torch.float32, device=x.device)
    _fused_rms_rope_bwd[(nprog,)](
        x, gamma, cos, sin, dout, rstd, dx, dgp,
        x.stride(0), x.stride(1), x.stride(2),
        dx.stride(0), dx.stride(1), dx.stride(2),
        T, H, N,
        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
    )
    return dx, dgp.sum(0).to(gamma.dtype)


class _FusedRMSRoPE(torch.autograd.Function):
    """QK-norm + RoPE for q and k in one autograd node (two kernel launches).
    Mirrors the C++ `fused_rms_rope(q, k, q_gamma, k_gamma, cache)`: q/k keep
    their incoming dtype (bf16 under autocast), the reduction and the rotation
    run in fp32 internally, gamma and the cache stay fp32.
    """

    @staticmethod
    def forward(ctx, q, k, q_gamma, k_gamma, cos, sin, eps):
        qo, q_rstd = _launch_fwd(q, q_gamma, cos, sin, eps)
        ko, k_rstd = _launch_fwd(k, k_gamma, cos, sin, eps)
        ctx.save_for_backward(q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd)
        return qo, ko

    @staticmethod
    def backward(ctx, dqo, dko):
        q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd = ctx.saved_tensors
        dq, dqg = _launch_bwd(q, q_gamma, cos, sin, dqo, q_rstd)
        dk, dkg = _launch_bwd(k, k_gamma, cos, sin, dko, k_rstd)
        return dq, dk, dqg, dkg, None, None, None


def fused_rms_rope(q, k, q_gamma, k_gamma, cos, sin, eps):
    """q: [B,Hq,T,D], k: [B,Hkv,T,D] -> same shapes, contiguous.
    `cos`/`sin` are [>=T, D//2] fp32; only the first T rows are read.
    """
    assert q.shape[-1] == k.shape[-1], "q and k must share head_dim"
    D = q.shape[-1]
    assert D % 2 == 0 and (D & (D - 1)) == 0, f"head_dim must be a power of two, got {D}"
    assert q.stride(-1) == 1 and k.stride(-1) == 1, "head_dim must be stride-1"
    T = q.shape[2]
    return _FusedRMSRoPE.apply(q, k, q_gamma, k_gamma,
                               cos[:T].contiguous(), sin[:T].contiguous(), eps)


def rms_rope_eager(x, gamma, cos, sin, eps):
    """Unfused QK-norm then NeoX RoPE, fp32 throughout — the reference the fused
    kernel is measured against (`--rms-rope eager`), and the fallback when
    Triton is unavailable."""
    T, D = x.shape[2], x.shape[-1]
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * gamma
    x1, x2 = xf[..., : D // 2], xf[..., D // 2:]
    c = cos[:T][None, None, :, :]
    s = sin[:T][None, None, :, :]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)


def build_rope_cache(seq_len, head_dim, theta, device):
    """NeoX / Llama RoPE cache: cos, sin of shape [seq_len, head_dim//2], fp32."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
                                             device=device) / head_dim))
    pos = torch.arange(seq_len, dtype=torch.float32, device=device)
    freqs = torch.outer(pos, inv_freq)                       # [T, D/2]
    return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous()


# =============================================================================
# 2. Configuration — 1:1 with ModelConfig in bluscript_single.cpp — and the WSD
#    learning-rate schedule.
# =============================================================================
@dataclass
class ModelConfig:
    # ---- model ----
    d_model: int = 960
    n_layers: int = 30
    q_heads: int = 15
    kv_heads: int = 5                     # grouped-query: G = q/kv = 4
    head_dim: int = 64                   # the fused kernels need 64 or 128
    ffn_hidden: int = 2560                # ~8/3 * d_model, rounded to /256
    context_length: int = 4096            # RoPE cache length; must be >= T
    rope_theta: float = 500000.0
    vocab_size: int = 50304               # 50257 (tiktoken) padded for the GPU
    weight_tying: bool = False

    rms_eps: float = 1e-5                 # block RMSNorm epsilon
    qk_eps: float = 1e-6                  # QK-norm (fused) epsilon

    # ---- optimization / data ----
    B: int = 4                            # micro-batch
    T: int = 4096                         # sequence length (<= context_length)
    global_batch: int = 524288            # tokens per optimizer step
    grad_clip: float = 1.0
    weight_decay: float = 0.1

    # WSD: linear warmup -> constant at max_lr -> cosine decay to min_lr.
    # Requires warmup_steps + stable_steps < max_steps - 1 (the remainder is
    # the decay tail).
    max_lr: float = 0.00071
    min_lr: float = 0.000071                # 10% of max_lr
    
    muon_lr: float = 0.00071
    muon_min_lr: float = 0.000071
    
    max_steps: int = 26703

    warmup_steps: int = 500
    stable_steps: int = int(max_steps * 0.9) - warmup_steps           # steps held at max_lr after warmup
     

    # ---- Muon (hidden 2D weights); the AdamW fields above still drive the
    #      embedding / lm_head / 1D group, unchanged from the AdamW arm ----
    # Default == max_lr: under adjust_lr_fn=match_rms_adamw the two are the same unit,
    # so out of the box both arms take the same size step and differ only in
    # direction. Same 10% floor, same WSD shape, so the schedules stay aligned.
   
    # The rest are torch.optim.Muon's own arguments, passed through unchanged.
    # momentum/nesterov/ns_steps are torch's defaults; adjust_lr_fn is not (see
    # the header — torch defaults to "original").
    muon_momentum: float = 0.95           # beta on the momentum buffer
    muon_ns_steps: int = 5                # Newton-Schulz iterations
    muon_adjust_lr: str = "match_rms_adamw"   # "match_rms_adamw" | "original"

    val_freq: int = 250
    val_loss_steps: int = 20

    # ---- checkpointing ----
    checkpointing: bool = True
    ckpt_freq: int = 1000                 # save every 1800 optimizer steps
    ckpt_keep: int = 40                    # retain the most-recent N checkpoints
    # Separate from the AdamW arm's dir: the two checkpoint formats differ (Muon
    # momentum vs Adam m/v) and auto-resume picks the newest file it finds.
    ckpt_dir: str = "ckpt_blu_muon"
    ckpt_prefix: str = "blumodel_muon"

    # ---- runtime knobs (not in the C++ struct) ----
    fused_rms_rope: bool = True
    activation_checkpoint: bool = False
    loss_chunks: int = 1
    # Depth probe: which layers --diag reports gradient / weight statistics for.
    probe_layers: tuple = (0, 8, 17, 26, 35)


def get_lr(step, cfg, max_lr=None, min_lr=None):
    """WSD: linear Warmup -> Stable (constant max_lr) -> cosine Decay to min_lr.
    Byte-for-byte the C++ get_lr().
    `max_lr`/`min_lr` override the config pair so the Muon group can ride the
    SAME schedule shape — same warmup length, same stable window, same cosine
    tail — at its own amplitude. Called with no override it is identical to the
    AdamW arm's get_lr(), which is what the embedding/norm group uses.
    """
    max_lr = cfg.max_lr if max_lr is None else max_lr
    min_lr = cfg.min_lr if min_lr is None else min_lr
    if step < cfg.warmup_steps:
        return max_lr * (step + 1) / cfg.warmup_steps
    if step < cfg.stable_steps + cfg.warmup_steps:
        return max_lr
    tail = cfg.max_steps - (cfg.stable_steps + cfg.warmup_steps) - 1
    decay_ratio = (step - (cfg.stable_steps + cfg.warmup_steps)) / max(tail, 1)
    decay_ratio = min(max(decay_ratio, 0.0), 1.0)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))       # 1 -> 0
    return min_lr + coeff * (max_lr - min_lr)


# =============================================================================
# 3. Data — port of the C++ DataLoaderLite (Tensor-Implementations/include/
#    DataLoader.h), including its pinned-host / side-stream prefetch ring.
#
# Shards are HEADERLESS flat uint16 token files. A shard is "val" if "val"
# appears anywhere in its basename, "train" otherwise — this covers both an
# explicit flux_train_*/flux_val_* naming and an unmarked-train naming like
# flux_20b_shard_*.bin alongside a separately named *val* shard. For a batch
# at position p:
#     x = tokens[p : p+B*T]      y = tokens[p+1 : p+B*T+1]
# p strides by B*T (world_size == 1 here). A shard tail that cannot fit a full
# stride is dropped and the loader wraps to the next shard.
#
# The prefetch mirrors the C++ ring: NBUF pinned staging buffers and NBUF device
# tensors, H2D issued on a side stream, CUDA events ordering producer against
# consumer. Each slot also carries the (shard, token offset) it was filled from,
# which is what --diag prints to name the bytes behind a loss spike.
# =============================================================================
class DataLoaderLite:
    NBUF = 4

    def __init__(self, B, T, split, data_root, device, master=True,
                 max_tokens_per_shard=100_000_000, rank=0, world_size=1):
        assert split in ("train", "val")
        self.B, self.T = B, T
        self.split = split
        self.max_tokens = max_tokens_per_shard
        self.device = device
        # DDP: every rank reads the SAME shard list but a DISJOINT, interleaved
        # slice of it (offset by rank, strided by world_size), so a training
        # step's world_size*grad_accum micro-batches never repeat a window.
        self.rank = rank
        self.world_size = world_size

        # Shard naming isn't consistent across datasets: some mark the split
        # explicitly in both names (flux_train_*.bin / flux_val_*.bin), others
        # only mark val and leave train shards unmarked (flux_20b_shard_*.bin
        # + a separate *val* file). Treat "contains val" as the val split and
        # everything else as train so both conventions resolve correctly.
        self.shards = sorted(
            p for p in glob.glob(os.path.join(data_root, "*.bin"))
            if ("val" in os.path.basename(p).lower()) == (split == "val")
        )
        if not self.shards:
            raise RuntimeError(f"no .bin shards found for split {split} in {data_root}")
        if master:
            print(f"found {len(self.shards)} shards for split {split}")

        BT = B * T
        self.pin_x = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
        self.pin_y = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
        self.dev_x = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
        self.dev_y = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
        self.copy_stream = torch.cuda.Stream(device=device)
        self.h2d_done = [torch.cuda.Event() for _ in range(self.NBUF)]
        self.h2d_recorded = [False] * self.NBUF
        self.where = [(0, 0)] * self.NBUF        # (shard_index, token_offset)
        self.reset()

    # -- shard plumbing ------------------------------------------------------
    def _load_shard(self, i):
        self.cur = i
        arr = np.memmap(self.shards[i], dtype=np.uint16, mode="r")
        self.tok = arr
        self.ntok = min(len(arr), self.max_tokens)

    def _advance_shard(self):
        self._load_shard((self.cur + 1) % len(self.shards))
        self.pos = self.B * self.T * self.rank

    def reset(self):
        self.copy_stream.synchronize()
        self._load_shard(0)
        self.pos = self.B * self.T * self.rank
        self.primed = False
        self.buf = 0

    # -- prefetch ------------------------------------------------------------
    def _stage(self, idx):
        BT = self.B * self.T
        stride = BT * self.world_size
        if self.ntok < BT + 1:
            raise RuntimeError(f"shard too small for one batch: {self.shards[self.cur]}")
        if self.pos + BT + 1 > self.ntok:
            self._advance_shard()

        # Host-side WAR guard: the pinned buffer may still be the source of the
        # H2D issued the last time this slot was filled (that copy is gated on
        # the device, not the host, so a host running ahead of the GPU would
        # otherwise overwrite bytes still in flight and hand the model a batch
        # spliced from two different windows).
        if self.h2d_recorded[idx]:
            self.h2d_done[idx].synchronize()

        self.where[idx] = (self.cur, self.pos)
        buf = np.asarray(self.tok[self.pos: self.pos + BT + 1]).astype(np.int64)
        self.pin_x[idx].copy_(torch.from_numpy(buf[:-1]))
        self.pin_y[idx].copy_(torch.from_numpy(buf[1:]))

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

        self.pos += stride
        if self.pos + BT + 1 > self.ntok:
            self._advance_shard()

    def next_batch(self):
        """-> (x, y, shard_index, token_offset), x/y int64 [B,T] on device."""
        if not self.primed:
            for i in range(self.NBUF):
                self._stage(i)
            self.primed = True
            self.buf = 0
        else:
            # Refill the slot returned LAST call. Its consumer's kernels are
            # already queued on the compute stream, so the producer must wait
            # for them before overwriting the device tensor.
            self.copy_stream.wait_stream(torch.cuda.current_stream())
            self._stage((self.buf - 1) % self.NBUF)

        torch.cuda.current_stream().wait_event(self.h2d_done[self.buf])
        i = self.buf
        self.buf = (self.buf + 1) % self.NBUF
        return self.dev_x[i], self.dev_y[i], self.where[i][0], self.where[i][1]

    def skip_batches(self, n):
        """Fast-forward `n` batches without touching the GPU (resume path)."""
        self.copy_stream.synchronize()
        self.primed = False
        BT = self.B * self.T
        stride = BT * self.world_size
        for _ in range(n):
            if self.pos + BT + 1 > self.ntok:
                self._advance_shard()
            self.pos += stride
            if self.pos + BT + 1 > self.ntok:
                self._advance_shard()


# =============================================================================
# 4. Model.
#
# Precision (bf16 mixed precision; both --precision modes run under autocast):
#   bf16-params : Linear weights stored bf16, embedding / RMSNorm / QK-norm
#                 scales stay fp32, fp32 master weights live in the optimizer
#                 — the exact scheme bluscript_single.cpp uses (default).
#   bf16        : all params fp32, autocast casts per-op (standard PyTorch AMP).
# In both cases: matmuls + attention bf16, RMSNorm / QK-norm reductions, RoPE,
# softmax and the loss in fp32. No GradScaler (bf16 has fp32's exponent range).
# =============================================================================
_FLASH = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]


class RMSNorm(nn.Module):
    """RMSNorm over d_model, fp32 weight. F.rms_norm is a fused ATen op and is
    not autocast-eligible, so it runs in the input dtype — fp32 on the residual
    stream, matching the C++ nn::RMSNorm."""

    def __init__(self, dim, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
        self.dim = (dim,)

    def forward(self, x):
        return F.rms_norm(x, self.dim, self.weight, self.eps)


class CausalGQA(nn.Module):
    """RMSNorm -> fused(QK-norm + RoPE) -> causal grouped-query attention -> Wo -> +x"""

    def __init__(self, cfg, resid_scale):
        super().__init__()
        self.cfg = cfg
        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.wQ = nn.Linear(cfg.d_model, cfg.q_heads * cfg.head_dim, bias=False)
        self.wK = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
        self.wV = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
        self.wO = nn.Linear(cfg.q_heads * cfg.head_dim, cfg.d_model, bias=False)
        # QK-norm scales, initialised to 1.0 (identity RMS scale), fp32.
        self.q_gamma = nn.Parameter(torch.ones(cfg.head_dim))
        self.k_gamma = nn.Parameter(torch.ones(cfg.head_dim))
        # GPT-2/Llama init: projections std=0.02, residual proj scaled by
        # (2*n_layers)^-0.5 to keep residual-stream variance ~constant.
        for lin in (self.wQ, self.wK, self.wV):
            nn.init.normal_(lin.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.wO.weight, mean=0.0, std=0.02 * resid_scale)

    def forward(self, x, cos, sin):
        cfg = self.cfg
        B, T, _ = x.shape
        h = self.norm(x)

        # [B,T,H*D] -> [B,H,T,D]. The transpose stays a view; the fused kernel
        # honours arbitrary (b,h,t) strides, so no .contiguous() is paid here.
        q = self.wQ(h).view(B, T, cfg.q_heads, cfg.head_dim).transpose(1, 2)
        k = self.wK(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)
        v = self.wV(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)

        if cfg.fused_rms_rope:
            q, k = fused_rms_rope(q, k, self.q_gamma, self.k_gamma, cos, sin, cfg.qk_eps)
        else:
            q = rms_rope_eager(q, self.q_gamma, cos, sin, cfg.qk_eps)
            k = rms_rope_eager(k, self.k_gamma, cos, sin, cfg.qk_eps)

        # Separate fused flash kernel. enable_gqa maps query head i to kv head
        # i // (q_heads/kv_heads), the same grouping as the C++ kernel.
        with sdpa_kernel(_FLASH):
            o = F.scaled_dot_product_attention(
                q.to(v.dtype), k.to(v.dtype), v, is_causal=True, enable_gqa=True)

        o = o.transpose(1, 2).reshape(B, T, cfg.q_heads * cfg.head_dim)
        return x + self.wO(o)


class SwiGLUMLP(nn.Module):
    """RMSNorm -> gate_up -> silu(gate)*up -> down -> +x   ([gate | up] layout)"""

    def __init__(self, cfg, resid_scale):
        super().__init__()
        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.gate_up = nn.Linear(cfg.d_model, 2 * cfg.ffn_hidden, bias=False)
        self.down = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
        nn.init.normal_(self.gate_up.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.down.weight, mean=0.0, std=0.02 * resid_scale)

    def forward(self, x):
        gate, up = self.gate_up(self.norm(x)).chunk(2, dim=-1)
        return x + self.down(F.silu(gate) * up)


class Block(nn.Module):
    def __init__(self, cfg, resid_scale):
        super().__init__()
        self.attn = CausalGQA(cfg, resid_scale)
        self.mlp = SwiGLUMLP(cfg, resid_scale)

    def forward(self, x, cos, sin):
        return self.mlp(self.attn(x, cos, sin))


class BLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        resid_scale = 1.0 / math.sqrt(2.0 * cfg.n_layers)

        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
        nn.init.normal_(self.wte.weight, mean=0.0, std=0.02)

        self.blocks = nn.ModuleList([Block(cfg, resid_scale) for _ in range(cfg.n_layers)])
        self.norm_f = RMSNorm(cfg.d_model, cfg.rms_eps)
        if not cfg.weight_tying:
            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
            nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)

        cos, sin = build_rope_cache(cfg.context_length, cfg.head_dim, cfg.rope_theta, "cpu")
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def apply_param_dtype(self, mode):
        """`bf16-params`: store Linear weights in bf16 (embedding, RMSNorm and
        QK-norm scales stay fp32) — the two-dtype-group layout of the C++ model.
        Under autocast this is transparent: the cast to bf16 becomes a no-op."""
        if mode != "bf16-params":
            return
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.weight.data = m.weight.data.to(torch.bfloat16)

    def forward(self, idx, targets=None):
        cfg = self.cfg
        T = idx.shape[1]
        cos, sin = self.rope_cos[:T], self.rope_sin[:T]

        x = self.wte(idx)                                    # fp32 [B,T,d_model]
        for blk in self.blocks:
            if cfg.activation_checkpoint and self.training:
                x = checkpoint(blk, x, cos, sin, use_reentrant=False)
            else:
                x = blk(x, cos, sin)
        x = self.norm_f(x)

        head_w = self.wte.weight if cfg.weight_tying else self.lm_head.weight
        if targets is None:
            return F.linear(x, head_w), None
        if cfg.loss_chunks > 1:
            return None, _chunked_ce(x, head_w, targets, cfg.loss_chunks)
        # Matches the C++: bf16 matmul, logits promoted to fp32 for the loss
        # (cross_entropy is an autocast-fp32 op, so the promotion is implicit).
        logits = F.linear(x, head_w)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss


def _chunked_ce(x, weight, targets, chunks):
    """lm_head + cross-entropy in row chunks, logits recomputed in backward.
    Peak logit memory drops from [B*T, V] to [B*T/chunks, V] at the cost of one
    extra lm_head matmul. Off by default (`--loss-chunks 1`) so the memory
    profile matches the C++ script, which materialises the full fp32 logits.
    """
    xf = x.reshape(-1, x.shape[-1])
    tf = targets.reshape(-1)
    n = xf.shape[0]
    size = (n + chunks - 1) // chunks
    total = xf.new_zeros((), dtype=torch.float32)
    for lo in range(0, n, size):
        hi = min(lo + size, n)
        part = checkpoint(_ce_chunk, xf[lo:hi], weight, tf[lo:hi], use_reentrant=False)
        total = total + part * (hi - lo)
    return total / n


def _ce_chunk(x, weight, targets):
    return F.cross_entropy(F.linear(x, weight).float(), targets)


# =============================================================================
# 5. Optimizer (Muon on the hidden matrices, AdamW on the rest), gradient
#    clipping, checkpointing.
# =============================================================================
def _is_muon_param(name, p):
    """Muon orthogonalizes a matrix, so it needs the parameter to BE a linear
    map between two feature spaces:
      * dim >= 2 — RMSNorm weights and QK-norm gammas are 1D and have no
        singular values to equalise;
      * not the embedding / lm_head — its rows index tokens, not features, so
        the left singular vectors mix unrelated vocabulary entries. Weight
        tying makes this one tensor serving both ends; either way it is AdamW's.
    Everything else — wQ/wK/wV/wO, gate_up, down — is Muon's: 75.5M of the
    114.2M parameters at the default config.
    """
    return p.dim() >= 2 and not name.startswith(("wte.", "lm_head."))


def _state_bytes(opt):
    """Bytes held in a torch optimizer's per-parameter state — Adam's m/v, or
    Muon's momentum buffer. Reads 0 before the first step(): both allocate
    lazily."""
    return sum(v.numel() * v.element_size()
               for st in opt.state.values()
               for v in st.values() if torch.is_tensor(v))


class MasterMuon:
    """torch.optim.Muon on the hidden matrices + torch.optim.AdamW on the rest,
    both driving fp32 master weights.
    Neither optimizer is reimplemented here. What this class owns is the part
    torch cannot know about:
      * ROUTING — torch.optim.Muon accepts whatever you hand it and the docs are
        explicit that only 2D hidden weights belong ("bias, and embedding,
        should be optimized by a standard method such as AdamW"). The split is
        `_is_muon_param`.
      * MASTER WEIGHTS — the parameters are bf16 under --precision bf16-params,
        and a bf16 tensor cannot absorb an lr*update of 1e-4 relative size. So
        both optimizers are constructed over the fp32 MASTERS, fed from the fp32
        landing buffers, and the masters are cast back over the params after the
        step. Unchanged from the AdamW arm — that machinery is what makes the
        two arms comparable, since both optimizers then read a bit-identical
        fp32 gradient and write a bit-identical bf16 parameter.
    Weight decay is decoupled in both (torch.optim.Muon applies it before the
    update, as AdamW does) and is applied to the master for the same reason: a
    decay written to the bf16 param is overwritten by the master write-back
    microseconds later and silently no-ops.
    State cost: Muon holds one momentum buffer per param, against Adam's two.
    """

    def __init__(self, named_params, cfg):
        self.cfg = cfg
        self.triples = []                 # (param, master, landing_grad), ALL params
        self.muon_masters = []            # masters routed to Muon
        decay, nodecay = [], []
        self.n_muon = self.n_decay = self.n_nodecay = 0
        for name, p in named_params:
            if p.dtype == torch.float32:
                m, g = p, None
            else:
                m = p.detach().float().clone()
                g = torch.zeros_like(m)
                m.grad = g
            self.triples.append((p, m, g))
            if _is_muon_param(name, p):
                self.muon_masters.append(m)
                self.n_muon += m.numel()
            elif p.dim() >= 2:
                decay.append(m)
                self.n_decay += m.numel()
            else:
                nodecay.append(m)
                self.n_nodecay += m.numel()
        self.opt = torch.optim.AdamW(
            [{"params": decay, "weight_decay": cfg.weight_decay},
             {"params": nodecay, "weight_decay": 0.0}],
            lr=cfg.max_lr, betas=(0.9, 0.95), eps=1e-8, fused=True)
        # Everything below is a torch.optim.Muon default except adjust_lr_fn
        # (torch: "original") and weight_decay, which is pinned to the AdamW
        # arm's 0.1 so the decay is identical across both groups and both arms.
        self.muon = torch.optim.Muon(
            self.muon_masters,
            lr=cfg.muon_lr,
            weight_decay=cfg.weight_decay,
            momentum=cfg.muon_momentum,
            nesterov=True,
            ns_steps=cfg.muon_ns_steps,
            adjust_lr_fn=cfg.muon_adjust_lr)

    @torch.no_grad()
    def step(self, lr, muon_lr):
        for p, m, g in self.triples:
            if g is not None:
                g.copy_(p.grad)               # bf16 grad -> fp32 landing buffer
        # AdamW group (embedding / lm_head / norms / gammas) — identical to the
        # baseline's update, on identical gradients.
        for grp in self.opt.param_groups:
            grp["lr"] = lr
        self.opt.step()
        # Muon group. Its nesterov path may consume the gradient in place; that
        # is harmless here because grad_norm and the --diag probes are read
        # before step() and the landing buffers are overwritten wholesale on the
        # next one.
        for grp in self.muon.param_groups:
            grp["lr"] = muon_lr
        self.muon.step()
        for p, m, g in self.triples:
            if g is not None:
                p.data.copy_(m)               # fp32 master -> bf16 param

    def state_breakdown(self):
        """-> (master_bytes, landing_grad_bytes, adam_mv_bytes, muon_mom_bytes).
        The first two are zero under --precision bf16, where the master IS the
        parameter and the gradient is already fp32. Both optimizers allocate
        their state lazily on the first step(), so the last two read 0 until
        then — the startup report uses the expected_* predictors instead.
        """
        master = landing = 0
        for p, m, g in self.triples:
            if m is not p:
                master += m.numel() * m.element_size()
                landing += g.numel() * g.element_size()
        mv = _state_bytes(self.opt)
        mom = _state_bytes(self.muon)
        return master, landing, mv, mom

    def state_bytes(self):
        return sum(self.state_breakdown())

    def expected_mv_bytes(self):
        """Adam m/v are lazily allocated, so the startup report has to predict
        them rather than measure them. Muon's params are excluded — they carry a
        momentum buffer instead, reported separately."""
        muon_ids = {id(m) for m in self.muon_masters}
        return sum(m.numel() * m.element_size() * 2
                   for _, m, _ in self.triples if id(m) not in muon_ids)

    def expected_muon_mom_bytes(self):
        """One fp32 momentum buffer per Muon parameter — the prediction the
        startup report prints before step 0 has allocated it."""
        return sum(m.numel() * m.element_size() for m in self.muon_masters)

    # -- checkpointing -------------------------------------------------------
    def state_dict(self):
        # `optimizer` tags the arm: resuming a Muon run from an AdamW checkpoint
        # would restore the weights and then silently drop the momentum, so
        # CheckpointManager refuses on a mismatch rather than producing a run
        # that looks fine and is not comparable to either arm.
        return {"optimizer": "muon",
                "adam": self.opt.state_dict(),
                "muon": self.muon.state_dict(),
                "masters": [None if m is p else m for p, m, g in self.triples]}

    @torch.no_grad()
    def load_state_dict(self, blob):
        """-> True if the fp32 masters were restored from the checkpoint."""
        if blob.get("optimizer") != "muon":
            raise SystemExit(
                "ERROR: this checkpoint was written by the AdamW arm (train.py). "
                "Point --ckpt-dir somewhere else, or resume it with train.py.")
        self.opt.load_state_dict(blob["adam"])
        if blob.get("muon") is not None:
            self.muon.load_state_dict(blob["muon"])
        masters = blob.get("masters")
        if masters is None or len(masters) != len(self.triples):
            return False
        for (p, m, g), saved in zip(self.triples, masters):
            if saved is not None and m is not p:
                m.copy_(saved.to(m.device))
        return True

    @torch.no_grad()
    def resync_masters(self):
        """Re-seed every master from its parameter.
        Needed after loading a checkpoint that carries no masters: the masters
        were built in __init__ from the RANDOM INIT, so without this the first
        step() would write them straight back over the restored weights and
        silently throw the run away.
        """
        for p, m, g in self.triples:
            if m is not p:
                m.copy_(p.detach().float())


def clip_grads_(params, max_norm):
    """ONE global L2 norm over every parameter, then one rescale — the textbook
    definition. (It differs from the ZeRO script, which clips each dtype group
    against its own norm and can therefore let the true global norm reach
    sqrt(n_groups) * grad_clip.) The per-tensor norms are promoted to fp32
    before they are combined, or a reduction over ~250 bf16 norms would lose
    ~3 digits."""
    grads = [p.grad for p in params if p.grad is not None]
    norms = torch._foreach_norm(grads)
    total = torch.linalg.vector_norm(torch.stack([n.float() for n in norms]))
    if max_norm and max_norm > 0:
        clip = (max_norm / (total + 1e-6)).clamp(max=1.0)
        torch._foreach_mul_(grads, clip)
    return total


class CheckpointManager:
    """Synchronous save / keep-N / auto-resume, mirroring the C++
    CheckpointManager. A checkpoint holds the model parameters, Adam's m/v, the
    fp32 masters, and the step it should resume at."""

    def __init__(self, ckpt_dir, prefix, keep):
        self.dir, self.prefix, self.keep = ckpt_dir, prefix, keep
        os.makedirs(ckpt_dir, exist_ok=True)

    def _path(self, step):
        return os.path.join(self.dir, f"{self.prefix}_step{step:08d}.pt")

    def _existing(self):
        out = []
        for p in glob.glob(os.path.join(self.dir, f"{self.prefix}_step*.pt")):
            m = re.search(r"_step(\d+)\.pt$", os.path.basename(p))
            if m:
                out.append((int(m.group(1)), p))
        return sorted(out)

    def save(self, step, model, opt, loss, last_val, cfg):
        tmp = self._path(step) + ".tmp"
        torch.save({"step": step, "loss": loss, "last_val": last_val,
                    # B/T/global_batch pin the data stride this checkpoint's
                    # step count assumes. world_size is deliberately NOT
                    # recorded: grad_accum = global_batch // (B*T*world_size)
                    # already makes tokens-per-step (and skip_batches'
                    # target offset) independent of how many ranks are
                    # running, so a GPU-count change alone is not a mismatch.
                    "run_cfg": {"B": cfg.B, "T": cfg.T,
                                "global_batch": cfg.global_batch},
                    "model": model.state_dict(),
                    "opt": opt.state_dict()}, tmp)
        # Rename last: a crash mid-write leaves a .tmp, never a truncated
        # checkpoint that the next load_latest would pick as the newest.
        os.replace(tmp, self._path(step))
        if self.keep > 0:
            for _, stale in self._existing()[:-self.keep]:
                os.remove(stale)

    def load_latest(self, model, opt, cfg, master=True):
        """-> (start_step, loss, last_val). (0, 0.0, nan) when there is
        nothing to resume."""
        existing = self._existing()
        if not existing:
            return 0, 0.0, float("nan")
        _, path = existing[-1]
        return self.load(path, model, opt, cfg, master)

    def load(self, path, model, opt, cfg, master=True):
        """Load a specific checkpoint file (e.g. from --resume-from), rather
        than auto-discovering the latest one in self.dir.
        -> (start_step, loss, last_val)."""
        blob = torch.load(path, map_location="cpu", weights_only=True)
        model.load_state_dict(blob["model"])
        if not opt.load_state_dict(blob["opt"]):
            opt.resync_masters()
        run_cfg = blob.get("run_cfg")
        if master and run_cfg is not None:
            mismatches = [f"{k} checkpoint={v!r} current={getattr(cfg, k)!r}"
                         for k, v in run_cfg.items() if getattr(cfg, k) != v]
            if mismatches:
                print("[WARN] resume config differs from the checkpoint's — "
                      "the data offset this step count assumes will be wrong:")
                for m in mismatches:
                    print(f"         {m}")
        print(f"[resume] loaded {path}")
        return (int(blob["step"]), float(blob["loss"]),
                float(blob.get("last_val", float("nan"))))


# =============================================================================
# 6. Training.
# =============================================================================
def _env_int(name, default=0):
    v = os.environ.get(name)
    if v in (None, ""):
        return default
    try:
        return int(v)
    except ValueError:
        print(f"[WARN] ignoring non-integer {name}={v!r}")
        return default


def seed_everything(seed):
    """Every RNG the run can draw from, off one seed.
    torch's CPU generator is the one that matters here (the model is built on
    CPU and moved), but `random` / numpy are seeded too so that anything added
    later — a shuffled shard order, a sampled probe — is reproducible by
    default rather than by accident.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def apply_determinism(mode):
    """Pin the run-to-run drift that seeding does not cover.
    Three sources, all on the backward pass:
      * nn.Embedding's dense backward scatters one atomicAdd per token into
        [vocab, d_model]; with weight tying that is also the lm_head gradient,
        so it is the largest reduction in the step;
      * the flash / mem-efficient attention backward splits dq over the key
        blocks and atomicAdds them together;
      * cuDNN picks kernels by wall-clock timing when `benchmark` is on, so the
        chosen algorithm can differ between runs of the same shape.
    Each is a float reduction in an order the scheduler chooses, so two runs
    diverge in the last mantissa bits of the first step and, through the
    optimizer, visibly by a few hundred steps.
    `strict` raises if some op has no deterministic kernel — better than
    reporting a reproducible run that is not one. `warn` downgrades that to a
    warning and keeps training; `off` restores the default fast kernels.
    """
    if mode == "off":
        return
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    # Routes embedding-backward and SDPA-backward to their ordered variants, and
    # enforces the cuBLAS workspace setting made at the top of this file.
    torch.use_deterministic_algorithms(True, warn_only=(mode == "warn"))


def build_argparser():
    ap = argparse.ArgumentParser(
        description="BLU single-GPU training — PyTorch port of bluscript_single.cpp")
    ap.add_argument("--data-root",
                    default=("/root/.cache/huggingface/hub/datasets--jayasriii--flux_20b/"
                             "/snapshots/ddc6cc7e43365a6b9096c30053df0fb001c94f4c"))
    ap.add_argument("--device", type=int, default=None,
                    help="CUDA device index (env: BLU_DEVICE)")
    ap.add_argument("--precision", choices=["bf16-params", "bf16"], default="bf16-params",
                    help="bf16-params (default): Linear weights stored bf16 with "
                         "fp32 master weights in the optimizer — the C++ scheme. "
                         "bf16: fp32 params + autocast (standard PyTorch AMP).")
    ap.add_argument("--rms-rope", choices=["fused", "eager"], default="fused",
                    help="fused = Triton QK-norm+RoPE kernel (default), "
                         "eager = unfused reference, for A/B on the kernel")
    ap.add_argument("--compile", action="store_true", help="torch.compile each block")
    ap.add_argument("--activation-checkpoint", action="store_true",
                    help="recompute each block in backward (trades ~30%% time for memory)")
    ap.add_argument("--loss-chunks", type=int, default=1,
                    help="chunk lm_head+CE into N pieces with recompute; 1 (default) "
                         "materialises the full logits, as the C++ does")
    # bluscript_single.cpp defaults unless overridden.
    ap.add_argument("--B", type=int, default=None, help="micro-batch (env: BLU_B)")
    ap.add_argument("--T", type=int, default=None)
    ap.add_argument("--d-model", type=int, default=None, help="smoke-test override")
    ap.add_argument("--n-layers", type=int, default=None, help="smoke-test override")
    ap.add_argument("--q-heads", type=int, default=None, help="smoke-test override")
    ap.add_argument("--kv-heads", type=int, default=None, help="smoke-test override")
    ap.add_argument("--ffn-hidden", type=int, default=None, help="smoke-test override")
    ap.add_argument("--vocab-size", type=int, default=None, help="smoke-test override")
    ap.add_argument("--global-batch", type=int, default=None,
                    help="tokens per optimizer step (env: BLU_GLOBAL_BATCH)")
    ap.add_argument("--max-steps", type=int, default=None,
                    help="caps the run for benchmarking; warmup/stable shrink "
                         "proportionally so the WSD decay tail stays non-empty "
                         "(env: BLU_MAX_STEPS)")
    # ---- Muon ----
    ap.add_argument("--muon-lr", type=float, default=None,
                    help="peak LR for the hidden 2D weights. Under the default "
                         "--muon-adjust-lr match_rms_adamw this is the same unit as "
                         "--max-lr, and defaults to the same value (2e-4) so "
                         "both arms take a matched first run; sweep it upward "
                         "from there (Muon usually takes 2-4x AdamW's LR).")
    ap.add_argument("--muon-min-lr", type=float, default=None,
                    help="floor of the Muon cosine tail (default 10%% of --muon-lr)")
    ap.add_argument("--muon-momentum", type=float, default=None,
                    help="torch.optim.Muon momentum (default 0.95, torch's own)")
    ap.add_argument("--muon-ns-steps", type=int, default=None,
                    help="torch.optim.Muon ns_steps, Newton-Schulz iterations "
                         "(default 5, torch's own); fewer = looser "
                         "orthogonalization, cheaper")
    ap.add_argument("--muon-adjust-lr", choices=["match_rms_adamw", "original"],
                    default=None,
                    help="torch.optim.Muon adjust_lr_fn. match_rms_adamw "
                         "(default HERE): 0.2*sqrt(max(A,B)), putting --muon-lr "
                         "on the same scale as --max-lr. original (default in "
                         "torch): sqrt(max(1,A/B)), needs a ~5.5x larger "
                         "--muon-lr at d_model=768.")
    ap.add_argument("--val-freq", type=int, default=None)
    ap.add_argument("--val-loss-steps", type=int, default=None)
    ap.add_argument("--diag", action="store_true",
                    help="per-micro-batch loss + depth probes (env: BLU_DIAG=1)")
    ap.add_argument("--warmup-measure", type=int, default=3,
                    help="steps excluded from the steady-state throughput summary")
    ap.add_argument("--csv", default="blu_muon_log.csv",
                    help="CSV log path; appended (not truncated) on resume. The "
                         "first 8 columns match the AdamW arm's CSV exactly, so "
                         "the two files plot against each other directly; "
                         "muon_lr is appended as a 9th.")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--determinism", choices=["strict", "warn", "off"], default="strict",
                    help="strict (default): deterministic CUDA kernels, so a rerun "
                         "of the same launch line reproduces the loss curve bit for "
                         "bit; raises if an op has no deterministic implementation. "
                         "warn: same, but non-deterministic ops warn instead of "
                         "raising. off: fastest kernels, results drift between runs.")
    # checkpointing
    ap.add_argument("--no-checkpointing", action="store_true")
    ap.add_argument("--ckpt-dir", default=None)
    ap.add_argument("--ckpt-prefix", default=None)
    ap.add_argument("--ckpt-freq", type=int, default=None)
    ap.add_argument("--ckpt-keep", type=int, default=None)
    ap.add_argument("--resume-from", default=None,
                    help="path to a specific *_step*.pt checkpoint to resume "
                         "from, overriding auto-discovery of the latest "
                         "checkpoint in --ckpt-dir")
    return ap


def config_from_args(args):
    cfg = ModelConfig()
    for f in ("B", "T", "d_model", "n_layers", "q_heads", "kv_heads", "ffn_hidden",
              "vocab_size", "global_batch", "val_freq", "val_loss_steps",
              "ckpt_dir", "ckpt_prefix", "ckpt_freq", "ckpt_keep",
              "muon_lr", "muon_min_lr", "muon_momentum", "muon_ns_steps",
              "muon_adjust_lr"):
        v = getattr(args, f, None)
        if v is not None:
            setattr(cfg, f, v)
    # A --muon-lr given without a floor keeps the 10% ratio the WSD tail assumes,
    # instead of silently decaying to the default 2e-5 from wherever it starts.
    if args.muon_lr is not None and args.muon_min_lr is None:
        cfg.muon_min_lr = 0.1 * cfg.muon_lr
    cfg.fused_rms_rope = (args.rms_rope == "fused")
    cfg.activation_checkpoint = args.activation_checkpoint
    cfg.loss_chunks = args.loss_chunks
    cfg.checkpointing = not args.no_checkpointing
    cfg.context_length = max(cfg.context_length, cfg.T)

    # Mirrors the BLU_MAX_STEPS branch in bluscript_single.cpp: shrink warmup and
    # stable with the run, or a short benchmark leaves no decay tail (the C++
    # asserts on it; get_lr here would clamp, which is worse — it would silently
    # report a schedule the config does not describe).
    if args.max_steps is not None:
        ms = args.max_steps
        cfg.max_steps = ms
        cfg.warmup_steps = max(1, ms // 10)
        cfg.stable_steps = max(1, ms // 2)
        if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
            cfg.stable_steps = max(1, cfg.max_steps - cfg.warmup_steps - 2)
        print(f"[bench] max_steps={ms} (warmup={cfg.warmup_steps} "
              f"stable={cfg.stable_steps})")
    return cfg


def validate_config(cfg, world_size=1, master=True):
    """Cheap checks; each one is a silent-corruption or out-of-bounds failure
    mode if it does not hold."""
    bad = []
    if cfg.T > cfg.context_length:
        bad.append(f"T ({cfg.T}) > context_length ({cfg.context_length}): "
                   "the RoPE cache is too short")
    if cfg.q_heads % cfg.kv_heads != 0:
        bad.append(f"q_heads ({cfg.q_heads}) not divisible by kv_heads ({cfg.kv_heads})")
    if cfg.head_dim not in (64, 128):
        bad.append(f"the fused kernels support head_dim 64 or 128, got {cfg.head_dim}")
    # Under DDP each rank runs B*T tokens per micro-step, so the world's tokens
    # per micro-step are B*T*world_size; global_batch must land on a whole
    # number of those or grad_accum below truncates and quietly under-shoots
    # the configured optimizer-step token count.
    if cfg.global_batch % (cfg.B * cfg.T * world_size) != 0:
        bad.append(f"global_batch ({cfg.global_batch}) not divisible by "
                   f"B*T*world_size ({cfg.B * cfg.T * world_size})")
    if cfg.warmup_steps < 1:
        bad.append("warmup_steps must be >= 1 (get_lr divides by it)")
    if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
        bad.append(f"warmup_steps + stable_steps "
                   f"({cfg.warmup_steps + cfg.stable_steps}) leaves no decay tail "
                   f"before max_steps ({cfg.max_steps})")
    if not hasattr(torch.optim, "Muon"):
        bad.append(f"torch {torch.__version__} has no torch.optim.Muon — it first "
                   "appears in the 2.9 docs. Upgrade torch, or run the AdamW "
                   "arm (train.py).")
    if cfg.muon_ns_steps < 1:
        bad.append("muon_ns_steps must be >= 1 (0 leaves the update unorthogonalized, "
                   "which is not Muon)")
    if cfg.muon_momentum <= 0.0 or cfg.muon_momentum >= 1.0:
        bad.append(f"muon_momentum must be in (0,1), got {cfg.muon_momentum}")
    if cfg.q_heads * cfg.head_dim != cfg.d_model and master:
        print(f"WARNING: q_heads*head_dim ({cfg.q_heads * cfg.head_dim}) "
              f"!= d_model ({cfg.d_model})")
    if bad:
        raise SystemExit("ERROR: " + "\nERROR: ".join(bad))


def main():
    args = build_argparser().parse_args()

    # Env overrides, for parity with the C++ launch line. An explicit CLI flag
    # wins over the environment.
    if args.B is None:
        args.B = _env_int("BLU_B", 0) or None
    if args.global_batch is None:
        args.global_batch = _env_int("BLU_GLOBAL_BATCH", 0) or None
    if args.max_steps is None:
        args.max_steps = _env_int("BLU_MAX_STEPS", 0) or None
    diag = args.diag or _env_int("BLU_DIAG", 0) != 0

    if not torch.cuda.is_available():
        raise SystemExit("ERROR: no CUDA device available")

    # ---- DDP bootstrap ----
    # torchrun sets RANK/LOCAL_RANK/WORLD_SIZE; their absence means a plain
    # single-process launch, where --device / BLU_DEVICE pick the GPU instead.
    # torch.cuda.set_device MUST happen before init_process_group: NCCL creates
    # its context on whatever device is current at that call, so doing it after
    # would leave every rank on this node contending for device 0.
    ddp = int(os.environ.get("RANK", -1)) != -1
    if ddp:
        ddp_rank = int(os.environ["RANK"])
        ddp_local_rank = int(os.environ["LOCAL_RANK"])
        ddp_world_size = int(os.environ["WORLD_SIZE"])
        master_process = ddp_rank == 0
        if ddp_local_rank >= torch.cuda.device_count():
            raise SystemExit(f"ERROR: LOCAL_RANK={ddp_local_rank} but only "
                             f"{torch.cuda.device_count()} CUDA device(s) visible")
        torch.cuda.set_device(ddp_local_rank)
        device = torch.device("cuda", ddp_local_rank)
        dist.init_process_group(backend="nccl")
        if master_process and (args.device is not None
                               or os.environ.get("BLU_DEVICE") not in (None, "")):
            print("[WARN] --device / BLU_DEVICE is ignored under torchrun — "
                  "LOCAL_RANK selects each rank's GPU.")
    else:
        ddp_rank, ddp_local_rank, ddp_world_size = 0, 0, 1
        master_process = True
        if args.device is None:
            args.device = _env_int("BLU_DEVICE", 0)
        if args.device >= torch.cuda.device_count():
            raise SystemExit(f"ERROR: BLU_DEVICE={args.device} but only "
                             f"{torch.cuda.device_count()} CUDA device(s) visible")
        torch.cuda.set_device(args.device)
        device = torch.device("cuda", args.device)

    seed_everything(args.seed)
    apply_determinism(args.determinism)
    # TF32 loses mantissa bits but is deterministic, so it is orthogonal to the
    # setting above and stays on.
    torch.set_float32_matmul_precision("high")

    cfg = config_from_args(args)
    if cfg.fused_rms_rope and not _HAS_TRITON:
        if master_process:
            print("[WARN] Triton not available — falling back to the eager QK-norm+RoPE")
        cfg.fused_rms_rope = False
    validate_config(cfg, ddp_world_size, master_process)

    # Tokens consumed per optimizer step = B * T * grad_accum * world_size.
    grad_accum = cfg.global_batch // (cfg.B * cfg.T * ddp_world_size)
    tokens_per_step = cfg.B * cfg.T * grad_accum * ddp_world_size

    if master_process:
        print("================ Configuration ================")
        mode_str = (f"DDP x{ddp_world_size} (torchrun, NCCL)" if ddp
                   else "SINGLE GPU (no torchrun / DDP / ZeRO)")
        print(f"  mode            : {mode_str}")
        print(f"  cuda device     : {ddp_local_rank if ddp else args.device} "
              f"({torch.cuda.get_device_name(device)})")
        print(f"  d_model         : {cfg.d_model}")
        print(f"  n_layers        : {cfg.n_layers}")
        print(f"  q_heads/kv_heads: {cfg.q_heads} / {cfg.kv_heads} "
              f"(G={cfg.q_heads // cfg.kv_heads})")
        print(f"  head_dim        : {cfg.head_dim}")
        print(f"  ffn_hidden      : {cfg.ffn_hidden}  (SwiGLU)")
        print(f"  context_length  : {cfg.context_length}")
        print(f"  rope_theta      : {cfg.rope_theta:g}  (NeoX half-split)")
        print(f"  vocab_size      : {cfg.vocab_size}")
        print(f"  weight_tying    : {'ON' if cfg.weight_tying else 'OFF'}")
        print(f"  B x T           : {cfg.B} x {cfg.T}")
        print(f"  global_batch    : {cfg.global_batch}  (grad_accum={grad_accum}, "
              f"world_size={ddp_world_size})")
        print("  optimizer       : MUON (hidden 2D) + AdamW (embed/head/1D)")
        print(f"  max_lr / min_lr : {cfg.max_lr:g} / {cfg.min_lr:g}   (AdamW group)")
        print(f"  muon lr / min   : {cfg.muon_lr:g} / {cfg.muon_min_lr:g}   "
              f"(adjust_lr_fn={cfg.muon_adjust_lr})")
        print(f"  muon momentum   : {cfg.muon_momentum:g}  | NS steps {cfg.muon_ns_steps}")
        print(f"  warmup / steps  : {cfg.warmup_steps} / {cfg.max_steps}")
        print(f"  WSD stable_steps: {cfg.stable_steps}  "
              f"(decay tail={cfg.max_steps - cfg.warmup_steps - cfg.stable_steps})")
        print(f"  precision       : {args.precision}")
        print(f"  seed / determ.  : {args.seed} / {args.determinism}")
        print(f"  rms+rope        : {'fused' if cfg.fused_rms_rope else 'eager'}"
              f" | compile {args.compile} | act-ckpt {cfg.activation_checkpoint}"
              f" | loss_chunks {cfg.loss_chunks}")
        print("===============================================", flush=True)

    # ---- model ----
    # Built and seeded identically on every rank (same seed, deterministic CPU
    # init); DDP's constructor broadcasts rank 0's parameters/buffers anyway,
    # so this is belt-and-suspenders against any future source of divergence.
    if master_process:
        print(f"Initializing model on CUDA device {device.index}...", flush=True)
    raw_model = BLU(cfg).to(device)
    raw_model.apply_param_dtype(args.precision)
    params = list(raw_model.parameters())
    n_params = sum(p.numel() for p in params)
    if master_process:
        print(f"Number of parameters: {n_params}")
    if args.compile:
        for blk in raw_model.blocks:
            blk.compile()

    # DDP wraps the (already dtype-cast, already compiled) raw module. Buffer
    # broadcasting is off: rope_cos/rope_sin are deterministic functions of
    # cfg, identical on every rank already, so syncing them every forward
    # would just be wasted communication.
    # if ddp:
    #     model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
    # else:
    #     model = raw_model

    if ddp:
        model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
        torch.cuda.synchronize()   # let DDP's constructor-time NCCL broadcast fully land
                                # before DataLoaderLite's independent copy_stream starts
                                # issuing its own async H2D work — see [crash investigation]
    else:
        model = raw_model


    # ---- optimizer ----
    # named_parameters(), not parameters(): the routing is by name (the
    # embedding and lm_head stay on AdamW). Duplicates are dropped by default,
    # so a weight-tied head is visited once, as `wte.weight`. Built from
    # raw_model so names never carry DDP's "module." prefix.
    opt = MasterMuon(raw_model.named_parameters(), cfg)
    if master_process:
        print(f"muon params: {opt.n_muon} | adamw decayed: {opt.n_decay}, "
              f"non-decayed: {opt.n_nodecay}")
    # Static footprint. Nothing here is sharded — plain DDP replicates the
    # full optimizer state on every rank (the whole cost of not running
    # ZeRO-1), so this is the per-GPU charge before a single activation is
    # allocated.
    gb = 2 ** 30
    param_bytes = sum(p.numel() * p.element_size() for p in params)
    grad_bytes = param_bytes                       # grads share the param dtype
    master_bytes, landing_bytes, _, _ = opt.state_breakdown()
    # Both optimizers allocate their state on the first step(), so at this point
    # the measured values are 0 and the report has to predict them.
    mv_bytes = opt.expected_mv_bytes()
    mom_bytes = opt.expected_muon_mom_bytes()
    static_bytes = (param_bytes + grad_bytes + master_bytes + landing_bytes
                    + mv_bytes + mom_bytes)
    if master_process:
        print(f"[{'ddp' if ddp else 'single'}] static footprint  (per GPU: "
              f"params + grads + optimizer state)")
        print(f"    parameters              : {param_bytes / gb:7.2f} GiB")
        print(f"    gradients               : {grad_bytes / gb:7.2f} GiB")
        print(f"    fp32 master weights     : {master_bytes / gb:7.2f} GiB")
        print(f"    fp32 grad landing bufs  : {landing_bytes / gb:7.2f} GiB")
        print(f"    Muon momentum           : {mom_bytes / gb:7.2f} GiB")
        print(f"    Adam exp_avg + exp_avg_sq: {mv_bytes / gb:6.2f} GiB")
        print(f"    static total            : {static_bytes / gb:7.2f} GiB", flush=True)

    # ---- data ----
    # Sharded by rank: each rank reads a disjoint, interleaved slice of the
    # same shard list (see DataLoaderLite), so DDP's cross-rank gradient
    # average is over world_size DIFFERENT windows, not world_size copies of
    # the same one.
    train_loader = DataLoaderLite(cfg.B, cfg.T, "train", args.data_root, device,
                                  master=master_process, rank=ddp_rank,
                                  world_size=ddp_world_size)
    # rank=0, world_size=1 UNCONDITIONALLY, on every rank: val_loss_steps
    # windows at fixed token offsets 0, BT, 2BT, ... Sharding this by
    # ddp_rank/ddp_world_size like the train loader would make the reported
    # val_loss a function of how many GPUs happen to be running (world_size=2
    # averages 2x as many, differently-offset windows as world_size=1) —
    # the two numbers would not be comparable across a GPU-count change, which
    # is what produced the apparent val_loss jump on the 2-GPU -> 1-GPU
    # resume. Every rank now redundantly evaluates the identical fixed
    # windows, so the metric means the same thing regardless of world_size.
    val_loader = DataLoaderLite(cfg.B, cfg.T, "val", args.data_root, device,
                                master=master_process, rank=0, world_size=1)

    # ---- checkpointing: manager + auto-resume ----
    # Every rank independently loads the same checkpoint file from the same
    # (assumed shared) filesystem — a read, so no race — which keeps model,
    # optimizer state, and start_step identical across ranks without an
    # explicit broadcast. Only rank 0 ever WRITES one (below).
    start_step, ckpt = 0, None
    last_val = float("nan")
    if cfg.checkpointing or args.resume_from:
        ckpt = CheckpointManager(cfg.ckpt_dir, cfg.ckpt_prefix, cfg.ckpt_keep)
        if args.resume_from:
            if master_process:
                print(f"[resume] loading explicit checkpoint: {args.resume_from}")
            start_step, resume_loss, last_val = ckpt.load(
                args.resume_from, raw_model, opt, cfg, master_process)
        else:
            start_step, resume_loss, last_val = ckpt.load_latest(
                raw_model, opt, cfg, master_process)
        if start_step > 0:
            if master_process:
                print(f"[resume] at step {start_step} (loss {resume_loss:.6f}, "
                      f"last val {last_val:.6f})")
            # Skip the batches already consumed by steps [0, start_step).
            # global_batch // (B*T*world_size) makes tokens-per-step (and so
            # this target offset) independent of world_size, so a GPU-count
            # change alone lands on the correct byte offset as long as B, T,
            # and global_batch are unchanged (checked above).
            train_loader.skip_batches(start_step * grad_accum)

    # ---- CSV log ----
    # Appended (not truncated) when resuming so a restart does not erase the
    # history of the run it is continuing. Rank 0 only — every rank would
    # otherwise open and write the same path.
    log_file = None
    if args.csv and master_process:
        resuming = start_step > 0 and os.path.exists(args.csv)
        log_file = open(args.csv, "a" if resuming else "w")
        if not resuming:
            log_file.write("step,train_loss,val_loss,tok_per_sec,grad_norm,dt_ms,lr,"
                           "elapsed_min,muon_lr\n")
        log_file.flush()

    amp = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    probes = [L for L in cfg.probe_layers if L < cfg.n_layers]
    steady_dts = []
    t_train_start = time.perf_counter()

    # ---- Ctrl-C: finish the in-flight step, save a checkpoint, exit clean ----
    # Registered here (not earlier) so a Ctrl-C during setup/data-loading still
    # gets Python's default immediate-abort behaviour — only the training loop
    # below gets the graceful path. A second Ctrl-C re-arms the default handler
    # and re-raises, for when the checkpoint save itself needs interrupting.
    interrupted = False

    def _handle_sigint(signum, frame):
        nonlocal interrupted
        if interrupted:
            signal.signal(signal.SIGINT, signal.SIG_DFL)
            raise KeyboardInterrupt
        interrupted = True
        if master_process:
            print("\n[ctrl-c] finishing this step, then checkpointing and exiting "
                  "(press again to abort immediately)...", flush=True)

    # A terminal Ctrl-C is delivered to every process in the foreground
    # process group, so under torchrun every rank sets `interrupted` at
    # roughly the same time and the ranks stay in lockstep for the collective
    # ops (DDP's grad all-reduce, the checkpoint barrier below) still ahead.
    signal.signal(signal.SIGINT, _handle_sigint)

    if master_process:
        print("\nStarting training" +
              (f" (resumed at step {start_step})" if start_step else "") + "...\n", flush=True)

    for step in range(start_step, cfg.max_steps):
        # ---------- periodic validation ----------
        if cfg.val_loss_steps > 0 and (step % cfg.val_freq == 0 or step == cfg.max_steps - 1):
            model.eval()
            val_loader.reset()
            acc = 0.0
            with torch.no_grad(), amp:
                for _ in range(cfg.val_loss_steps):
                    x, y, _, _ = val_loader.next_batch()
                    acc += model(x, y)[1].item() / cfg.val_loss_steps
            model.train()
            # val_loader is rank=0/world_size=1 on every rank (see its
            # construction above), so every rank just computed the identical
            # fixed windows redundantly — no cross-rank reduction needed, and
            # val_loss no longer depends on how many ranks are running.
            last_val = acc
            if master_process:
                print(f"  [val] step {step} | val_loss {last_val:.6f}", flush=True)

        # ---------- training step ----------
        t0 = time.perf_counter()
        model.zero_grad(set_to_none=True)
        # Fresh zeros rather than an in-place scale of last step's accumulator:
        # multiplying by zero does NOT clear a NaN, so a single non-finite
        # micro-batch would poison the accumulator for the rest of the run.
        loss_accum = torch.zeros((), device=device)
        micro_loss, micro_where = [], []

        for micro in range(grad_accum):
            x, y, shard, offset = train_loader.next_batch()
            # DDP allreduces gradients on every .backward(); for grad_accum > 1
            # only the LAST micro-batch of the step should pay for that — the
            # others accumulate purely locally under no_sync().
            sync_ctx = (nullcontext() if (not ddp or micro == grad_accum - 1)
                       else model.no_sync())
            with sync_ctx:
                with amp:
                    _, loss = model(x, y)
                loss_accum += loss.detach()
                # Per-micro-batch loss (--diag). The reported train loss is the mean
                # over grad_accum micro-steps, so a single pathological micro-batch
                # is diluted while its gradient dominates the accumulated norm
                # almost undiluted — which is how a modest loss bump can sit next to
                # a large norm spike. `micro_where` names the shard and token offset
                # so the responsible bytes can be decoded offline.
                #
                # Costs a device sync per micro-step, so it is off by default.
                if diag:
                    micro_loss.append(loss.item())
                    micro_where.append((shard, offset))
                (loss / grad_accum).backward()

        if ddp:
            # Average each rank's accumulated (summed) loss BEFORE dividing by
            # grad_accum, so loss_val is the mean over world_size*grad_accum
            # micro-batches — the true loss behind this optimizer step, not
            # just rank 0's slice. Also doubles as a sync point: a NaN on any
            # rank poisons the average for every rank, so the check below
            # raises everywhere instead of only where it originated (which
            # would otherwise hang the survivors on the next collective).
            dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
        loss_val = loss_accum.item() / grad_accum
        if math.isnan(loss_val) or math.isinf(loss_val):
            if log_file:
                log_file.close()
            # Every rank reaches this identically (the all_reduce above is a
            # sync point and NaN propagates through AVG), so every rank exits
            # here together rather than the survivors hanging on the next
            # collective.
            if ddp:
                dist.destroy_process_group()
            raise SystemExit(f"ERROR: NaN/Inf loss at step {step}")

        # Depth probe (--diag), read BEFORE the update, so these are the
        # pre-update gradients. Compare the profile on a spiking step against a
        # quiet one: a blow-up that begins mid-stack and grows toward the loss
        # points at the forward, while a flat multiplicative lift across all
        # depths points at the optimizer/update.
        depth = []
        if diag:
            for L in probes:
                a = raw_model.blocks[L].attn
                gn = a.wQ.weight.grad
                depth.append((L,
                              gn.float().norm().item() if gn is not None else -1.0,
                              a.q_gamma.abs().max().item(),
                              a.wQ.weight.abs().max().float().item()))

        grad_norm = clip_grads_(params, cfg.grad_clip).item()
        lr = get_lr(step, cfg)                                   # AdamW group
        muon_lr = get_lr(step, cfg, cfg.muon_lr, cfg.muon_min_lr)
        opt.step(lr, muon_lr)

        torch.cuda.synchronize()   # make the per-step timing reflect completed GPU work
        dt = time.perf_counter() - t0
        tok_per_sec = tokens_per_step / dt
        elapsed_min = (time.perf_counter() - t_train_start) / 60.0
        if step - start_step >= args.warmup_measure:
            steady_dts.append(dt)

        # ---------- per-step report (rank 0 only — every rank's numbers agree
        # on loss/grad_norm post-allreduce, so printing from every rank would
        # only duplicate lines, not add information) ----------
        if master_process:
            vs = "     -  " if last_val != last_val else f"{last_val:.6f}"
            print(f"step {step:6d} | train {loss_val:.6f} | val {vs} "
                  f"| tok/s {tok_per_sec:10.2f} | norm {grad_norm:.4f} "
                  f"| dt {dt * 1e3:8.2f} ms | lr {lr:.4e} | mlr {muon_lr:.4e} "
                  f"| mem {torch.cuda.max_memory_allocated() / gb:5.2f}/"
                  f"{torch.cuda.max_memory_reserved() / gb:5.2f} GiB", flush=True)

            if diag and micro_loss:
                # Flag micro-batches that stand out against this step's own median.
                med = sorted(micro_loss)[len(micro_loss) // 2]
                marks = " ".join(("*" if m > med + 0.5 else "") + f"{m:.3f}"
                                 for m in micro_loss)
                print(f"  [micro] med {med:.4f} | {marks}")
                for m, (ml, (shard, offset)) in enumerate(zip(micro_loss, micro_where)):
                    if ml > med + 0.5:
                        print(f"  [micro] OUTLIER m={m} loss {ml:.4f} (median {med:.4f}) "
                              f"from shard {shard} token offset {offset}")
            if diag and depth:
                print("  [depth] ||g(wQ)|| by layer:"
                      + "".join(f" L{L} {g:.3e}" for L, g, _, _ in depth)
                      + " | max|q_gamma|:"
                      + "".join(f" L{L} {gm:.4f}" for L, _, gm, _ in depth)
                      + " | max|wQ|:"
                      + "".join(f" L{L} {w:.4f}" for L, _, _, w in depth), flush=True)

            if log_file:
                log_file.write(f"{step},{loss_val:.6f},{last_val:.6f},{tok_per_sec:.6f},"
                               f"{grad_norm:.6f},{dt * 1e3:.6f},{lr:.6f},{elapsed_min:.6f},"
                               f"{muon_lr:.6f}\n")
                log_file.flush()

        if ddp:
            # `interrupted` is set by this process's own SIGINT handler, so in
            # principle one rank could see it before another (a signal isn't
            # guaranteed simultaneous delivery, even though in practice a
            # terminal Ctrl-C hits every rank in the same process group at
            # once). Left asymmetric, it would desync the checkpoint barrier
            # and the `break` below against ranks still issuing DDP's regular
            # per-step collectives — a hang, not just a wrong answer. This
            # runs unconditionally every step (never gated on a per-rank
            # value), so it can never itself mismatch across ranks.
            flag = torch.tensor(1.0 if interrupted else 0.0, device=device)
            dist.all_reduce(flag, op=dist.ReduceOp.MAX)
            interrupted = flag.item() > 0

        # ---------- periodic checkpoint (also fires once on Ctrl-C) ----------
        # Labeled step+1 (the next iteration to run) so load_latest resumes
        # exactly where this save left off — no duplicated or skipped iteration.
        # Captures weights + Adam m/v AFTER this step's update. Rank 0 only —
        # every rank's raw_model/opt state is identical post-allreduce, so a
        # second writer would only race the first on the same path. The
        # barrier holds every rank at this point until the save (or its
        # failure) is done, so a slow write can't let other ranks race ahead
        # into the next step's collectives before the checkpoint is durable.
        if cfg.checkpointing and ((step + 1) % cfg.ckpt_freq == 0
                                  or step == cfg.max_steps - 1
                                  or interrupted):
            if master_process:
                try:
                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, cfg)
                except Exception as e:                   # noqa: BLE001 — advisory
                    print(f"[WARN] checkpoint save at step {step + 1} failed: {e}")
            if ddp:
                dist.barrier()

        if interrupted:
            if not cfg.checkpointing and master_process:
                print(f"[ctrl-c] checkpointing is disabled (--no-checkpointing); "
                      f"stopping at step {step} without saving.", flush=True)
            break

        # Reclaim the extra allocator high-water mark from the first couple of
        # steps' one-time warmup allocations (autograd graph caches, etc.).
        if step in (start_step, start_step + 1):
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()

    torch.cuda.synchronize()
    if log_file:
        log_file.close()

    # ---- steady-state summary ----
    # dt/throughput are this rank's own timing; under DDP all ranks run in
    # lockstep (grad sync + the checkpoint barrier), so rank 0's numbers are
    # representative and `tokens_per_step` already counts every rank's tokens.
    if steady_dts and master_process:
        xs = sorted(steady_dts)
        med = xs[len(xs) // 2]
        mean = sum(xs) / len(xs)
        alloc = torch.cuda.max_memory_allocated() / gb
        reserved = torch.cuda.max_memory_reserved() / gb
        # Measured now that Adam's m/v exist, rather than the startup estimate.
        static = (param_bytes + grad_bytes + opt.state_bytes()) / gb
        mode_str = f"DDP x{ddp_world_size}" if ddp else "single GPU"
        print("\n" + "=" * 78)
        print(f"  STEADY-STATE SUMMARY — {mode_str}  optimizer=muon+adamw  "
              f"precision={args.precision}  "
              f"rms+rope={'fused' if cfg.fused_rms_rope else 'eager'}")
        print(f"  torch.optim.Muon  lr={cfg.muon_lr:g} "
              f"adjust_lr_fn={cfg.muon_adjust_lr} "
              f"momentum={cfg.muon_momentum:g} ns_steps={cfg.muon_ns_steps}")
        print(f"  ({len(xs)} steps, first {args.warmup_measure} excluded)")
        print("=" * 78)
        print(f"  step time          : {med * 1e3:10.2f} ms  (mean {mean * 1e3:.2f})")
        print(f"  throughput         : {tokens_per_step / med:10.1f} tok/s")
        print(f"  tokens / step      : {tokens_per_step:10d}")
        print(f"  static (p+g+opt)   : {static:10.2f} GiB  (per GPU)")
        print(f"  activations (peak) : {max(alloc - static, 0.0):10.2f} GiB")
        print(f"  peak mem allocated : {alloc:10.2f} GiB")
        print(f"  peak mem reserved  : {reserved:10.2f} GiB")
        print("=" * 78)

    if ddp:
        dist.destroy_process_group()
    if master_process:
        print("\n=== Training complete ===")


if __name__ == "__main__":
    main()