#!/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= CUDA device to train on (default 0) # BLU_MAX_STEPS= short benchmark run (rescales WSD) (default off) # BLU_B= micro-batch size (default 4) # BLU_GLOBAL_BATCH= 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 = 2048 n_layers: int = 36 q_heads: int = 16 kv_heads: int = 4 # grouped-query: G = q/kv = 4 head_dim: int = 128 # the fused kernels need 64 or 128 ffn_hidden: int = 5632 # ~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 = True rms_eps: float = 1e-5 # block RMSNorm epsilon qk_eps: float = 1e-6 # QK-norm (fused) epsilon # ---- optimization / data ---- B: int = 2 # 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 = 2e-3 min_lr: float = 2e-4 # 10% of max_lr max_steps: int = 2000 warmup_steps: int = 200 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. muon_lr: float = 2e-3 muon_min_lr: float = 2e-4 # 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 = int(max_steps / 20) # 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 whose basename contains the # split ("train" / "val"). 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 self.shards = sorted( p for p in glob.glob(os.path.join(data_root, "*.bin")) if split in os.path.basename(p) ) 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--navingv--flux/" "snapshots/5232b3d133d2eb71905725f288ff6b56dc0df773")) 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 # ---- 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()