code/run_full_d24_w8.py

# decoderstack_medium_pt-sft-fable.py
#
# Single-file d24 pre-training pipeline with a handwritten forward/backward and a
# written-out optimizer: no autograd, no torch.optim, no param groups, no nn.Module.
#
# (From Chris -- Core design decisions):
# - No nn.Module, no m.to, no state_dict / load_state_dict.
#   - Every tensor is created directly on the device, at its final dtype.
#   - No accommodations for "prior checkpoints", we're starting from scratch.
# - No torch.optim or autograd, we're doing everything manually.
# - Use globals -- global cfg, global m -- don't pass things around.
#   - The model is a plain class used as a namespace of plain torch.Tensors.
#     nn.Parameter does nothing for us: Parameter exists for autograd leaf
#     bookkeeping and Module registration, neither of which we use. Plain
#     tensors are directly usable in the math (m.W_in, not m.W_in.weight),
#     accept attached state (.grad32, .mantissa, ...) just like Parameters,
#     and default to requires_grad=False -- which is what we want everywhere,
#     because we implement grad.
# - Dtypes are hardcoded everywhere -- stated at creation, never inferred by
#   matching another tensor's dtype. (No fp64 parity tier in this file.)
# - Hardcoded to the d24 config; none of nanochat's auto-scaling by model size.
# - Multi-GPU shards the optimizer, not the model (nanochat's scheme): every
#   rank holds the full bf16 live weights and full grad accumulators, optimizer
#   state is allocated at shard sizes, and optimizer_step wraps the same update
#   kernels in reduce-scatter -> owned-shard update -> live all-gather.
# - We're not doing FP8 yet.
#
# The "§" technique defines the code sections in here.
#
# The model/training code comes from the nanochat repo, branch fwd-bwd
# (nanochat/train_step.py, nanochat/gpt.py). That branch's d24 run is the
# reference implementation we want to match -- we're refactoring and dropping
# baggage, not changing the math:
# C:\Users\chris\Documents\GitHub\agent-ops\nanochat\2026-07-29_0833am_d24-throughput-gap\NOTES.md
#
# The code below the seam (marked near the bottom) comes from the 'stacks' repo,
# pulled mainly for the pre-tokenized data + distributed loader and CORE eval.
#
# One-off derived quantities (parameter counts, flops/token, the training
# horizon, the LR/WD batch corrections, cu_seqlens sizing) are HARDCODED in
# this script; `scaling.py` (kept alongside it) recomputes and documents them.


# --------------------------------------------------------------------------------
# § Setup
# --------------------------------------------------------------------------------

import os
import sys
import time as _time
run_wall_t0 = _time.perf_counter()
del _time

with open(sys.argv[0], 'r') as f:
    code = f.read()   # the run section logs the script source to wandb

import datetime
import gc
import glob
import json
import math
import random
import threading
import time
from pathlib import Path
from types import SimpleNamespace
from typing import NamedTuple

import numpy as np
import wandb

os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
import torch
import torch._dynamo as dynamo
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from kernels import get_kernel

dynamo.config.recompile_limit = 64

# ==== Distributed setup ====
# dist is always initialized (launch under torchrun, even for one process) --
# the data pipeline below the seam uses dist.barrier() and the loader shards
# by rank.
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
assert torch.cuda.is_available()
device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
torch.cuda.set_device(device)
dist.init_process_group(backend="nccl", device_id=device)
dist.barrier()
master_process = (rank == 0)

def print0(*args, console=False, **kwargs):
    if master_process:
        print(*args, **kwargs)

# ==== Wandb helpers ====

class DummyWandb:
    """No-op wandb replacement when logging is disabled."""
    def log(self, *args, **kwargs): pass
    def finish(self): pass
        

# BF16 dense peak FLOPS by GPU, for the MFU denominator. Just the GPUs this
# pipeline actually runs on; the full many-vendor table (and sources) lives in
# scaling.py. GH200 carries the same H100-class SXM die: 989 TFLOPS.
PEAK_FLOPS = {"GH200": 989e12, "H100": 989e12, "A100": 312e12}

def next_multiple_of_n(v: float | int, *, n: int):
    return next(x for x in range(n, int(v) + 1 + n, n) if x >= v)


# --------------------------------------------------------------------------------
# § Flash Attention (raw FA3 forward/backward)
# --------------------------------------------------------------------------------
# The handwritten backward calls FA3's raw _flash_attn_forward/_flash_attn_backward
# torch.library ops directly -- no autograd Function in between. The forward
# returns the softmax LSE, which the backward consumes alongside the stashed
# output. FA3 only; there is no SDPA/naive fallback in this file.

_cc_major, _ = torch.cuda.get_device_capability()
if _cc_major == 9:   # Hopper: the varunneal build gets better H100 results
    fa3 = get_kernel("varunneal/flash-attention-3").flash_attn_interface
    RAW_BWD_TAKES_BUFFERS = False   # raw backward allocates and RETURNS dq/dk/dv
else:                # Ampere sm80/86 / Ada sm89: community FA3 build
    assert _cc_major == 8, f"FA3 required (sm8x or sm90); got sm{_cc_major}x"
    _k = get_kernel("kernels-community/flash-attn3")
    # The raw ops live in flash_attn_interface; the top level only re-exports
    # the varlen/kvcache wrappers.
    fa3 = getattr(_k, "flash_attn_interface", _k)
    RAW_BWD_TAKES_BUFFERS = True    # raw backward takes pre-allocated dq/dk/dv buffers


def flash_attn_varlen_fwd_lse(q, k, v, cu_seqlens, max_seqlen, window_size):
    """Attention forward that also returns what the handwritten backward needs:
    (out, softmax_lse), with lse (H, T) fp32."""
    out, softmax_lse, *_ = fa3._flash_attn_forward(
        q, k, v,
        cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
        max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
        softmax_scale=q.shape[-1] ** -0.5, causal=True,
        window_size_left=window_size[0], window_size_right=window_size[1])
    return out, softmax_lse


def flash_attn_varlen_bwd(dout, q, k, v, out, softmax_lse, cu_seqlens, max_seqlen, window_size):
    """Attention backward for flash_attn_varlen_fwd_lse: returns (dq, dk, dv).
    The two FA3 builds' raw backward ops differ in calling convention -- the
    sm80 community build's schema takes pre-allocated dq/dk/dv buffers (grads
    come back through them), the sm90 varunneal build's allocates and returns
    them -- hence the branch on the module-level flag."""
    softmax_scale = q.shape[-1] ** -0.5
    if RAW_BWD_TAKES_BUFFERS:
        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
        fa3._flash_attn_backward(
            dout, q, k, v, out, softmax_lse,
            cu_seqlens, cu_seqlens,     # cu_seqlens_q, cu_seqlens_k
            None, None,                 # seqused_q, seqused_k
            max_seqlen, max_seqlen,
            dq, dk, dv,
            softmax_scale,
            True,                       # is_causal
            window_size[0], window_size[1],
            0.0,                        # softcap
            False,                      # deterministic
            0,                          # sm_margin
        )
    else:
        dq, dk, dv, _ = fa3._flash_attn_backward(
            dout, q, k, v, out, softmax_lse,
            cu_seqlens, cu_seqlens,     # cu_seqlens_q, cu_seqlens_k
            None, None,                 # seqused_q, seqused_k
            max_seqlen, max_seqlen,
            softmax_scale,
            True,                       # is_causal
            window_size[0], window_size[1],
            0.0,                        # softcap
            False,                      # deterministic
            0,                          # sm_margin
        )
    return dq, dk, dv


# --------------------------------------------------------------------------------
# § Model Config
# --------------------------------------------------------------------------------

# Value embeddings (ResFormer-style) live on alternating layers, last always
# included. Banked over just the VE layers; ve_index maps layer -> bank slot
# (-1 = no VE on this layer) and is read by every forward body.
#
# Note: Deriving head size or count from d_model is a bad habit that has
#       propagated through ~everyone's model code.
#       There are only three real constraints--these values must match:
#         1. Number of key and value heads
#         2. Query-key head sizes
#         3. Value-output head sizes
#
# Recommended short window size:
# -(-seq_len // 4 // 128) * 128  # ceil to FA3 tile size (2048 -> 768)
class StackConfig:

    # Model
    n_layers:   int = 24
    d_model:    int = 1536

    # Input
    d_vocab:    int = 32768 # Must arrive padded (tensor cores, sharding) -- no
                            # auto-padding in this file; asserted below.
    d_smr_gate: int = 24    # Input to smear gate is first 'd' positions of the
                            # normed input embedding.
    # Attention
    n_q_heads:  int = 12
    n_kv_heads: int = 12
    n_o_heads:  int = 12  # TODO - fold into n_qo_heads, since the code doesn't support 
                          #        a different ratio (group size) for qk vs. vo. 
    d_qk:       int = 128 # Note: FA2 requires d_qk == d_vo, FA3 does not.
    d_vo:       int = 128

    # Context and Sliding Window Attention
    seq_len:          int = 2048
    short_win_size:   int = 768
    full_ctxt_layers: list[int] = [   3,    7,    11,     15,     19,     23] # "sssL" pattern

    window_sizes:     list[tuple[int, int]]  # Derived below.

    # Attention - Value Embeddings
    d_ve_gate: int = 12  # First 'd' positions of residual stream (after x0
                         # blending and norm) are the gate input.
                         # ve gates exist per head, per layer.
    ve_layers: list[int] = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23] # 0-indexed.
    ve_index:  list[int] # Derived from ve_layers.
    num_ves:   int

    # MLP
    d_mlp:      int = 4 * 1536

    # Training batch (nanochat d24 speedrun spec). Tokens, not sequences: with
    # varlen packing a micro-batch is one packed 1-D stream, so the token count
    # is the real quantity (= 16 seqs x 2048 in nanochat's batched terms).
    # Total batch 2^20 tokens/step is nanochat's Power Lines auto-compute for
    # d24.
    micro_batch_tokens: int = 65536   # per rank, per micro-batch
    total_batch_size:   int = 2**20   # tokens per optimizer step

    # Training horizon: the d24 speedrun spec (data:param ratio 8) --
    # 8 x 729,810,624 scaling params = 5,838,484,992 tokens // 2^20 per step
    # = 5,568 steps. Derivation: scaling.py.
    num_iterations: int = 5568

    # Evaluation and logging
    val_tokens:      int = 10485760   # per val-bpb pass: 320 training-shaped micro-batches
    val_loss_every:  int = 250
    eval_buffer_tokens: int = 65536   # CORE/chat eval packing buffer. Eval is
                                      # forward-only (no stash, no grads), so a
                                      # buffer well past the training micro-batch
                                      # fits easily; the rotary cache is sized to
                                      # cover it.
    save_checkpoint: bool = True
    # Mid-run checkpoint capture, in COMPLETED optimizer steps: state is
    # written on entering these loop steps (the final state always saves).
    # 1950 = the first LR/momentum-cooldown step at the 5568-step horizon
    # (the hold ends after update 1949 = N - round(0.65*N)): the last
    # uncooled state -- the one to resume from to train the horizon longer.
    save_steps:      tuple = (1950,)
    run_id:          str = f"{str(datetime.datetime.now().strftime('%Y-%m-%d_%H%M%S'))}-d24"
    wandb_run:       str = "dummy"    # "dummy" disables wandb
    wandb_project:   str = "decoderstack"

cfg = StackConfig() # Make config a global, don't pass it around.

# Sanity: the constraints the axes above must satisfy.
assert cfg.d_vocab % 64 == 0, "vocab must arrive padded to 64 (no auto-padding here)"
assert cfg.n_o_heads == cfg.n_q_heads, "attention output consumes one slot per query head"
assert cfg.n_q_heads % cfg.n_kv_heads == 0, "GQA needs query heads to tile over kv heads"
assert cfg.d_qk % 2 == 0, "rotary splits the qk head dim in half"
assert cfg.full_ctxt_layers[-1] == cfg.n_layers - 1, "final layer recommended to have full context"

# Derived quantities:
# Map layers to VE bank slots.
cfg.ve_index = [cfg.ve_layers.index(i) if i in cfg.ve_layers else -1 for i in range(cfg.n_layers)]
cfg.num_ves = len(cfg.ve_layers)

# Per-layer window sizes for sliding window attention.
# List of (left, right) tuples for FA3's window_size parameter:
# - left: how many tokens before current position to attend to
# - right: how many tokens after current position to attend to (0 for causal)
# "Full context" is (seq_len, 0): documents are at most seq_len tokens and
# varlen attention is doc-isolated, so a seq_len window is unlimited in effect.
cfg.window_sizes = [(cfg.short_win_size, 0)] * cfg.n_layers  # All short, ...
for i in cfg.full_ctxt_layers:
    cfg.window_sizes[i] = (cfg.seq_len, 0)                   # ... then overwrite with full.

# Derived batch quantities. Fixed total => grad accum scales down as GPUs are
# added: 32 at world=1, 4 at world=8. grad_scale rides into forward_backward
# as loss_scale, replacing the loss division of an autograd loop; at world>1
# it composes with ReduceOp.AVG grad comm to give the global batch mean.
assert cfg.total_batch_size % (cfg.micro_batch_tokens * world_size) == 0, \
    "total batch must divide evenly into per-rank micro-batches"
grad_accum_steps = cfg.total_batch_size // (cfg.micro_batch_tokens * world_size)
grad_scale = 1 / grad_accum_steps


# --------------------------------------------------------------------------------
# § Shard Assignment
# --------------------------------------------------------------------------------
# Each GPU is responsible for a "shard" of the optimizer work:
# - Muon banks shard over their layer axis (dim 0).
# - AdamW params shard over the row axis of their (rows, cols) view -- vocab
#   rows for input_embeds/lm_head, flattened (ve_slot * vocab) rows for
#   value_embeds. (ve_slot alone is too small to divide across a world, and the
#   rows are interchangeable for AdamW's elementwise update.)
# - ve_gate is NOT sharded: it is tiny (~thousands of floats) and ragged
#   against world sizes, so every rank runs the full-size update instead.
# - grad32 always stays FULL size on every rank -- it is the source buffer for
#   the reduce-scatter, not a shard.
#
# No zero-padding support: every sharded axis must divide evenly (asserted
# below). d24's axes -- 24 layers, 32768 vocab rows, 393,216 ve rows -- all
# divide by the world sizes we'd run (1, 2, 4, 8).
#
# At world_size == 1 every shard IS the whole tensor: the slices below span
# their full axes and optimizer_step's collectives short-circuit. One code
# path, degenerate comm.

assert cfg.n_layers % world_size == 0, \
    f"Muon layer-sharding needs n_layers % world == 0 ({cfg.n_layers} % {world_size})"
layer_shard_size  = cfg.n_layers // world_size
layer_shard_start = rank * layer_shard_size
layer_shard_slice = slice(layer_shard_start, layer_shard_start + layer_shard_size)

assert cfg.d_vocab % world_size == 0, \
    f"AdamW row-sharding needs vocab % world == 0 ({cfg.d_vocab} % {world_size})"
vocab_shard_size  = cfg.d_vocab // world_size
vocab_shard_start = rank * vocab_shard_size
vocab_shard_slice = slice(vocab_shard_start, vocab_shard_start + vocab_shard_size)

ve_rows = cfg.num_ves * cfg.d_vocab
assert ve_rows % world_size == 0, \
    f"AdamW row-sharding needs ve_slot*vocab % world == 0 ({ve_rows} % {world_size})"
ve_row_shard_size  = ve_rows // world_size
ve_row_shard_start = rank * ve_row_shard_size
ve_row_shard_slice = slice(ve_row_shard_start, ve_row_shard_start + ve_row_shard_size)


# --------------------------------------------------------------------------------
# § Model Initialization
# --------------------------------------------------------------------------------

class Model:
    """Namespace of plain tensors -- the live weights. Each weight also carries
    its training state as attached attributes, allocated alongside it below:

      .grad32        full-size gradient accumulator (fp32; bf16 for the two
                     embedding tables), explicitly zeroed between steps
      .grad32_slices per-layer views of grad32 for the 3-D banks (see below)
      .mantissa      lower 16 bits of the fp32 master (uint16, shard-size)
      .frst_mntm     Muon first moment (fp32, shard-size)
      .scnd_mntm     Muon factored second moment (fp32, shard-size)
      .residual_dim  the weight axis that faces the residual stream (-1 or -2);
                     NorMuon's per-neuron mean-square is taken along it
      .exp_avg       AdamW first moment (fp32, shard-size)
      .exp_avg_sq    AdamW second moment (fp32, shard-size)
    """

    # Input
    input_embeds: Tensor
    smear_gate:   Tensor
    smear_lambda: Tensor

    # Attention
    W_Q: Tensor
    W_K: Tensor
    W_V: Tensor
    W_O: Tensor
    value_embeds: Tensor
    ve_gate:      Tensor

    # MLP
    W_in:  Tensor
    W_out: Tensor

    # Cross-Layer
    resid_lambdas:  Tensor  # Per-layer gain on the residual stream.
    x0_lambdas:     Tensor  # Per-layer coefficient for reading the input embedding.
    backout_lambda: Tensor  # How much of layer 16's output to remove from the stream
                            # prior to the lm head.

    # Output
    lm_head: Tensor

    # Buffers (rotary cache; not trained, not checkpointed)
    cos: Tensor
    sin: Tensor

    # The trained weights, in declaration order -- this tuple defines "every
    # trained weight". __iter__ walks them so call sites can just say
    # `for p in m` (grad zeroing); the names key the checkpoint dicts.
    weight_names = ("input_embeds", "smear_gate", "smear_lambda",
                    "W_Q", "W_K", "W_V", "W_O", "value_embeds", "ve_gate",
                    "W_in", "W_out", "resid_lambdas", "x0_lambdas",
                    "backout_lambda", "lm_head")

    def __iter__(self):
        return (getattr(self, n) for n in self.weight_names)


# ==== Tensor Creation Idioms ====
# Reduce the boilerplate for defining weights and buffers.

fp32_empty   = lambda *shape: torch.empty(*shape, dtype=torch.float32, device=device)
bf16_empty   = lambda *shape: torch.empty(*shape, dtype=torch.bfloat16, device=device)
fp32_zeros   = lambda *shape: torch.zeros(*shape, dtype=torch.float32, device=device)
bf16_zeros   = lambda *shape: torch.zeros(*shape, dtype=torch.bfloat16, device=device)
uint16_zeros = lambda *shape: torch.zeros(*shape, dtype=torch.uint16, device=device)

# We use fp32 for the "master" weights, which are what we store on disk, and for 
# avoiding rounding off small optimizer updates. 
# All forward and backward computation is done on bf16 matrices (the "live" weights). 
# Note that bf16 is just fp32 with the lower 16-bits of mantissa dropped;
# rather than hold 16-bit and 32-bit copies at once, we stash those lower
# 16 mantissa bits, and reconstruct the full 32-bit precision to update then
# resplit.
upper_bf16   = lambda w: (w.contiguous().view(torch.int32) >> 16).to(torch.int16).view(torch.bfloat16)
lower_uint16 = lambda w: (w.contiguous().view(torch.int32)      ).to(torch.int16).view(torch.uint16)

# Set the seed so that every rank gets the same initialization -- no broadcast
# from a master rank needed.
torch.manual_seed(42)
torch.cuda.manual_seed(42)

m = Model()

# Written out one tensor per line, deliberately: the shape, the dtype, and
# therefore the memory cost of every weight and every piece of optimizer state
# is readable in one place, and the axis names say which dimension is sharded.
#
# Dtype scheme (hardcoded, stated per tensor below):
# - Matrix banks + lm_head: bf16 live + uint16 mantissa (fp32 master via the
#   mantissa trick), fp32 gradients, fp32 moments.
# - Embedding tables (input_embeds, value_embeds): bf16 live + uint16 mantissa
#   (fp32 master via the mantissa trick). This deviates from nanochat, which
#   kept its embeddings plain bf16 and let AdamW update them in place -- we
#   pair them with a mantissa so the one AdamW kernel serves everything,
#   rather than carrying a second bf16-live variant. Overall our code
#   ~matches the validation loss of the original.
#   Gradients are bf16 -- these are the two biggest tensors in the model,
#   fp32 grads would double their scatter traffic and (at world>1) comm bytes,
#   and bf16 matches the autograd baseline's numerics (bf16 params -> bf16
#   .grad). Everything else accumulates gradients in fp32.
# - Scalars (resid/x0 lambdas, smear, backout): fp32 live, no mantissa, same
#   as they've always been. (Rounding them to bf16 was tried during the port
#   and cost +0.016 val bpb, so they stay fp32.)
#
# Initialization values:
#   input_embeds:   normal,  std=0.8
#   lm_head:        normal,  std=0.001
#   W_Q, W_K, W_V:  uniform, bound=sqrt(3)/sqrt(d_model)     -> std = 1/sqrt(d_model)
#   W_O:            zeros
#   W_in:           uniform, bound=0.4*sqrt(3)/sqrt(d_model)  -> std = 0.4/sqrt(d_model)
#   W_out:          zeros
#   value_embeds:   uniform, bound=sqrt(3)/sqrt(d_model)      (same as W_V)
#   ve_gate:        uniform in [0, 0.02] (slightly above neutral)
#   resid_lambdas:  1.15 -> 1.05 linear decay over depth
#   x0_lambdas:     0.20 -> 0.05 linear decay over depth
#   smear_gate:     zeros
#   smear_lambda:   zeros (smear disabled at init)
#   backout_lambda: zeros (backout disabled at init)
#   (Zeros for smear/backout is what nanochat's baselines actually trained
#   with: it intended backout_lambda=0.2 and a kaiming smear_gate, but its
#   meta-device init never ran those. Details at the Scalars block below.)

# Uniform init bound. Var(Uniform(-a, a)) = a^2/3, so std = a/sqrt(3): to hit
# a target std of 1/sqrt(d_model), the bound must be sqrt(3) times it.
matrix_init_s = (3 ** 0.5) * (cfg.d_model ** -0.5)

# ==== Input Embeddings ====
# bf16 live; draw in fp32 and let copy_ round -- drawing straight into bf16
# would quantize the distribution rather than the samples. The master upcast of
# a bf16 live is lossless, so the mantissa starts at zero.
# TODO - Leaving the zero-mantissa init for the moment (it matches the fwd-bwd
#        reference), but we'll likely switch to keeping the draw's lower 16
#        bits (the lm_head split pattern) once we have the chance to test.
m.input_embeds = bf16_empty(cfg.d_vocab, cfg.d_model)
m.input_embeds.copy_(fp32_empty(cfg.d_vocab, cfg.d_model).normal_(mean=0.0, std=0.8))

m.input_embeds.grad32     = bf16_zeros(cfg.d_vocab, cfg.d_model) # TODO - Change to `grad` since there's no colision?
m.input_embeds.mantissa   = uint16_zeros(vocab_shard_size, cfg.d_model)
m.input_embeds.exp_avg    = fp32_zeros(vocab_shard_size, cfg.d_model)
m.input_embeds.exp_avg_sq = fp32_zeros(vocab_shard_size, cfg.d_model)

# ==== Value Embeddings ====
# Same init std as W_V; same bf16-live / zero-mantissa path as input_embeds.
# AdamW state is shaped over the FLATTENED (ve_slot * vocab) row axis;
# optimizer_step passes matching 2-D views of the live bank and its grad.
# Flattening (vs a 3-D state mirroring the bank) is what lets ONE
# reduce-scatter/all-gather over dim-0 rows shard the whole bank evenly --
# per-slot vocab sharding on the 3-D layout would need a collective per VE
# slot. At world=1 a 3-D state would also work, but would have to reallocate
# the moment we go multi-GPU.
m.value_embeds = bf16_empty(cfg.num_ves, cfg.d_vocab, cfg.n_kv_heads * cfg.d_vo)
m.value_embeds.copy_(fp32_empty(cfg.num_ves, cfg.d_vocab, cfg.n_kv_heads * cfg.d_vo)
                     .uniform_(-matrix_init_s, matrix_init_s))

m.value_embeds.grad32     = bf16_zeros(cfg.num_ves, cfg.d_vocab, cfg.n_kv_heads * cfg.d_vo)
m.value_embeds.mantissa   = uint16_zeros(ve_row_shard_size, cfg.n_kv_heads * cfg.d_vo)
m.value_embeds.exp_avg    = fp32_zeros(ve_row_shard_size, cfg.n_kv_heads * cfg.d_vo)
m.value_embeds.exp_avg_sq = fp32_zeros(ve_row_shard_size, cfg.n_kv_heads * cfg.d_vo)

# ==== LM Head ====
# Drawn in fp32 and split -- unlike the embeddings, its mantissa is real from
# step zero.
lm_head_fp32 = fp32_empty(cfg.d_vocab, cfg.d_model).normal_(mean=0.0, std=0.001)

m.lm_head          = upper_bf16(lm_head_fp32)                       # Live weights - bf16
m.lm_head.mantissa = lower_uint16(lm_head_fp32[vocab_shard_slice])  # Lower 16 bits for optimizer

del lm_head_fp32

m.lm_head.grad32     = fp32_zeros(cfg.d_vocab, cfg.d_model)
m.lm_head.exp_avg    = fp32_zeros(vocab_shard_size, cfg.d_model)
m.lm_head.exp_avg_sq = fp32_zeros(vocab_shard_size, cfg.d_model)

# ==== Attention ====
# Parameter banks: the layer index is dim 0. Each slice uses F.linear's
# (out_features, in_features) convention and is consumed as `x @ w.mT`.
# Initialize in fp32 and split into bf16 live + uint16 mantissa.
W_Q_fp32 = fp32_empty(cfg.n_layers, cfg.n_q_heads  * cfg.d_qk, cfg.d_model).uniform_(-matrix_init_s, matrix_init_s)
W_K_fp32 = fp32_empty(cfg.n_layers, cfg.n_kv_heads * cfg.d_qk, cfg.d_model).uniform_(-matrix_init_s, matrix_init_s)
W_V_fp32 = fp32_empty(cfg.n_layers, cfg.n_kv_heads * cfg.d_vo, cfg.d_model).uniform_(-matrix_init_s, matrix_init_s)
W_O_fp32 = fp32_zeros(cfg.n_layers,               cfg.d_model, cfg.n_o_heads * cfg.d_vo)  # projections start at zero

m.W_Q = upper_bf16(W_Q_fp32) # Live weights - bf16
m.W_K = upper_bf16(W_K_fp32)
m.W_V = upper_bf16(W_V_fp32)
m.W_O = upper_bf16(W_O_fp32)

# For the mantissa, we only need to hold our shard of the weights.
m.W_Q.mantissa = lower_uint16(W_Q_fp32[layer_shard_slice]) # Lower 16 bits for optimizer
m.W_K.mantissa = lower_uint16(W_K_fp32[layer_shard_slice])
m.W_V.mantissa = lower_uint16(W_V_fp32[layer_shard_slice])
m.W_O.mantissa = lower_uint16(W_O_fp32[layer_shard_slice])

del W_Q_fp32, W_K_fp32, W_V_fp32, W_O_fp32

# Gradients (full size -- the reduce-scatter source, never sharded)
m.W_Q.grad32 = fp32_zeros(cfg.n_layers, cfg.n_q_heads  * cfg.d_qk, cfg.d_model)
m.W_K.grad32 = fp32_zeros(cfg.n_layers, cfg.n_kv_heads * cfg.d_qk, cfg.d_model)
m.W_V.grad32 = fp32_zeros(cfg.n_layers, cfg.n_kv_heads * cfg.d_vo, cfg.d_model)
m.W_O.grad32 = fp32_zeros(cfg.n_layers,               cfg.d_model, cfg.n_o_heads * cfg.d_vo)

# First-momentum buffers for Muon (sharded)
m.W_Q.frst_mntm = fp32_zeros(layer_shard_size, cfg.n_q_heads  * cfg.d_qk, cfg.d_model)
m.W_K.frst_mntm = fp32_zeros(layer_shard_size, cfg.n_kv_heads * cfg.d_qk, cfg.d_model)
m.W_V.frst_mntm = fp32_zeros(layer_shard_size, cfg.n_kv_heads * cfg.d_vo, cfg.d_model)
m.W_O.frst_mntm = fp32_zeros(layer_shard_size,               cfg.d_model, cfg.n_o_heads * cfg.d_vo)

# Second momentum (NorMuon variance reduction) holds a running average of each
# neuron's mean-square update, so it is a vector (per layer) rather than a
# matrix mirroring the weights. (The neuron's rms is the square root of what's
# stored; the kernel applies it as an rsqrt.)
# NorMuon is a ~no-op for square matrices: polar express produces a
# ~orthonormal matrix, so the neuron norms are already ~uniform and there is
# nothing to normalize (confirmed with experiments). It only affects attention
# when the number of heads times the head size differs from d_model.
# The original code uses a heuristic to infer the neuron dimension by assuming
# that it is the smaller of the two. While typical, it's not certain. Instead,
# we specify it directly.
# Neurons can be identified directly by their interaction with the residual
# stream--they read from it and write to it and match it in length, so the
# mean-square is taken along the residual dimension.
# Note that the attention output projection consists of heads as well, and
# they are stored transposed relative to QKV, so we calculate the mean-square
# along dim -2.
m.W_Q.residual_dim = -1
m.W_K.residual_dim = -1
m.W_V.residual_dim = -1
m.W_O.residual_dim = -2
m.W_Q.scnd_mntm = fp32_zeros(layer_shard_size, cfg.n_q_heads  * cfg.d_qk, 1)
m.W_K.scnd_mntm = fp32_zeros(layer_shard_size, cfg.n_kv_heads * cfg.d_qk, 1)
m.W_V.scnd_mntm = fp32_zeros(layer_shard_size, cfg.n_kv_heads * cfg.d_vo, 1)
m.W_O.scnd_mntm = fp32_zeros(layer_shard_size, 1, cfg.n_o_heads * cfg.d_vo)

# ==== MLPs ====
# For a transformer, 'MLP' is something of a misnomer. It's closer to a
# lookup table, containing pairs of vectors, both of length d_m.
# For a given pair (w_in, w_out), if the residual stream is positively
# aligned with w_in, then w_out is written back to it.
# But unlike a look up table, where a read-write operation is captured
# by a single row, here the model composes the operation across many
# vector pairs.

W_in_fp32  = fp32_empty(cfg.n_layers, cfg.d_mlp,   cfg.d_model).uniform_(-matrix_init_s * 0.4, matrix_init_s * 0.4)
W_out_fp32 = fp32_zeros(cfg.n_layers, cfg.d_model, cfg.d_mlp)             # projections start at zero

m.W_in  = upper_bf16(W_in_fp32) # Live weights - bf16
m.W_out = upper_bf16(W_out_fp32)

m.W_in.mantissa  = lower_uint16(W_in_fp32[layer_shard_slice]) # Lower 16 bits for optimizer
m.W_out.mantissa = lower_uint16(W_out_fp32[layer_shard_slice])

del W_in_fp32, W_out_fp32

# Gradients (full size)
m.W_in.grad32  = fp32_zeros(cfg.n_layers, cfg.d_mlp,   cfg.d_model)
m.W_out.grad32 = fp32_zeros(cfg.n_layers, cfg.d_model, cfg.d_mlp)

# First-momentum buffers for Muon (sharded)
m.W_in.frst_mntm  = fp32_zeros(layer_shard_size, cfg.d_mlp,   cfg.d_model)
m.W_out.frst_mntm = fp32_zeros(layer_shard_size, cfg.d_model, cfg.d_mlp)

# Residual dimension: W_in rows read from the residual stream, W_out columns
# write to it.
m.W_in.residual_dim  = -1
m.W_out.residual_dim = -2
m.W_in.scnd_mntm  = fp32_zeros(layer_shard_size, cfg.d_mlp, 1)
m.W_out.scnd_mntm = fp32_zeros(layer_shard_size, 1, cfg.d_mlp)

# ==== VE Gates ====
# Muon, REPLICATED: tiny and ragged against world sizes, so every rank runs the
# full-size update rather than paying comm to shard a few thousand floats.
ve_gate_fp32 = fp32_empty(cfg.num_ves, cfg.n_kv_heads, cfg.d_ve_gate).uniform_(0.0, 0.02)

m.ve_gate          = upper_bf16(ve_gate_fp32)
m.ve_gate.mantissa = lower_uint16(ve_gate_fp32)  # replicated: full-size mantissa

del ve_gate_fp32

m.ve_gate.grad32     = fp32_zeros(cfg.num_ves, cfg.n_kv_heads, cfg.d_ve_gate)
m.ve_gate.frst_mntm  = fp32_zeros(cfg.num_ves, cfg.n_kv_heads, cfg.d_ve_gate)
m.ve_gate.residual_dim = -1  # gate rows read a d_ve_gate slice of the residual stream
m.ve_gate.scnd_mntm  = fp32_zeros(cfg.num_ves, cfg.n_kv_heads, 1)

# ==== Scalars ====
# fp32-LIVE with no mantissa pair (see the dtype scheme note above). AdamW,
# replicated.
# These serve separate purposes:
# - resid_lambdas: Directly scales the residual stream at the start of each layer.
# - x0_lambdas: How strongly the input embedding is added to the residual stream.
# Per-layer scalars: linear decay over depth. Stronger residual and more
# input-embedding blending at early layers, both tapering with depth.
m.resid_lambdas = torch.linspace(1.15, 1.05, cfg.n_layers, dtype=torch.float32, device=device)
m.x0_lambdas    = torch.linspace(0.20, 0.05, cfg.n_layers, dtype=torch.float32, device=device)

# Smear/backout start disabled, zeros everywhere.
# Note: nanochat pre-flattening had a bug here--it intended backout_lambda=0.2
# and a kaiming smear_gate, but under meta-device init those never executed and
# to_empty() left zeroed storage. Zeros is what every tuned baseline actually
# trained with, so now it's explicit rather than luck.
m.smear_gate     = fp32_zeros(1, cfg.d_smr_gate)
m.smear_lambda   = fp32_zeros(1)
m.backout_lambda = fp32_zeros(1)

m.resid_lambdas.grad32  = fp32_zeros(cfg.n_layers)
m.x0_lambdas.grad32     = fp32_zeros(cfg.n_layers)
m.smear_gate.grad32     = fp32_zeros(1, cfg.d_smr_gate)
m.smear_lambda.grad32   = fp32_zeros(1)
m.backout_lambda.grad32 = fp32_zeros(1)

m.resid_lambdas.exp_avg     = fp32_zeros(cfg.n_layers)
m.resid_lambdas.exp_avg_sq  = fp32_zeros(cfg.n_layers)
m.x0_lambdas.exp_avg        = fp32_zeros(cfg.n_layers)
m.x0_lambdas.exp_avg_sq     = fp32_zeros(cfg.n_layers)
m.smear_gate.exp_avg        = fp32_zeros(1, cfg.d_smr_gate)
m.smear_gate.exp_avg_sq     = fp32_zeros(1, cfg.d_smr_gate)
m.smear_lambda.exp_avg      = fp32_zeros(1)
m.smear_lambda.exp_avg_sq   = fp32_zeros(1)
m.backout_lambda.exp_avg    = fp32_zeros(1)
m.backout_lambda.exp_avg_sq = fp32_zeros(1)

# ==== Rotary Cache ====
# Without an nn.Module these are just attributes on m -- register_buffer only
# existed for state_dict/.to() plumbing we no longer have (and these were
# persistent=False anyway). With varlen training the whole micro-batch is one
# packed sequence, so the cache spans the largest T any forward sees: the
# training micro-batch (val micro-batches match it) or the CORE/chat eval
# packing buffer, whichever is bigger. The assert in the forward bodies
# catches it if we ever exceed.
rotary_seq_len = max(cfg.micro_batch_tokens, cfg.eval_buffer_tokens)
channel_range = torch.arange(0, cfg.d_qk, 2, dtype=torch.float32, device=device)  # stride the channels
inv_freq = 1.0 / (100000 ** (channel_range / cfg.d_qk))
t_pos = torch.arange(rotary_seq_len, dtype=torch.float32, device=device)          # stride the time steps
freqs = torch.outer(t_pos, inv_freq)   # rotation frequency at each (time, channel) pair
m.cos = freqs.cos().to(torch.bfloat16)[None, :, None, :]  # add batch and head dims
m.sin = freqs.sin().to(torch.bfloat16)[None, :, None, :]  # for later broadcasting
del channel_range, inv_freq, t_pos, freqs

# ==== Bank Gradient Slice Views ====
# The 3-D banks get `grad32_slices`: per-slice VIEWS built OUTSIDE any compiled
# graph. The forward/backward bodies accumulate through these, never through
# `grad32[i]` -- an in-graph bank slice functionalizes into a whole-bank
# select_scatter copy (10-20x the cost of the slice add at these bank sizes),
# while a view created out of graph arrives as an input and mutates genuinely
# in place.
m.W_Q.grad32_slices = list(m.W_Q.grad32.unbind(0))
m.W_K.grad32_slices = list(m.W_K.grad32.unbind(0))
m.W_V.grad32_slices = list(m.W_V.grad32.unbind(0))
m.W_O.grad32_slices = list(m.W_O.grad32.unbind(0))

m.W_in.grad32_slices  = list(m.W_in.grad32.unbind(0))
m.W_out.grad32_slices = list(m.W_out.grad32.unbind(0))

m.ve_gate.grad32_slices      = list(m.ve_gate.grad32.unbind(0))
m.value_embeds.grad32_slices = list(m.value_embeds.grad32.unbind(0))


# (Grad zeroing happens as a plain loop at the training-loop call site --
# every .grad32 is zeroed after each optimizer_step, since gradients
# accumulate across a step's micro-batches AND Muon's nesterov lerp mutates
# grad32 in place.)


# --------------------------------------------------------------------------------
# § Schedules
# --------------------------------------------------------------------------------
# A run's optimizer is defined up front: every learning rate, beta and weight
# decay for every step is computed here, before training starts, into per-step
# tables of *update coefficients* -- the numbers the fused kernels actually
# multiply by. The optimizer then holds no hyperparameters of its own and the
# training loop has nothing to set per step; the kernels just gather row
# `t_step` of each table. Folding all the way down to coefficients buys:
# - The bias corrections leave the kernel (betas are per-role constants, so
#   the closed `1 - beta^t` form is exact).
# - Nothing about the schedule is left for the loop to do per step. Tables are
#   device-resident and the step counter is a device tensor, so a step involves
#   the host for nothing at all.

class AdamWTabs(NamedTuple):
    """What an AdamW step multiplies by, one (N,) table per field. eps is never
    scheduled, so it rides as a plain kernel argument instead of a table."""
    wd_mul: Tensor           # 1 - lr*wd             decoupled weight decay
    one_minus_beta1: Tensor  # 1 - beta1             exp_avg lerp weight
    one_minus_beta2: Tensor  # 1 - beta2             exp_avg_sq lerp weight
    rsqrt_bias2: Tensor      # 1/sqrt(bias2)         second-moment bias correction
    step_size: Tensor        # lr / bias1            lr schedule x first-moment bias correction


class MuonCoeffs(NamedTuple):
    """What a Muon step multiplies by. Muon's second moment is self-normalizing
    (the v_norm/v_norm_new rescale), so it needs no bias correction."""
    momentum: Tensor            # nesterov momentum
    one_minus_momentum: Tensor  # 1 - momentum        frst_mntm lerp weight
    one_minus_beta2: Tensor     # 1 - beta2           variance-reduction lerp weight
    lr: Tensor                  # lr (the per-bank aspect scale arrives separately, via lr_mul)
    lr_wd: Tensor               # lr * weight_decay   cautious decay


def build_schedules(num_iterations, batch_lr_scale=1.0, weight_decay=0.28,
                    warmup_steps=40, warmdown_ratio=0.65, final_lr_frac=0.05):
    """Named table sets with the tuned nanochat base_train hyperparameters,
    written out flat. Baked assumptions (a Ramp class used to support more):
    exactly three shaped schedules exist -- the shared LR multiplier, Muon
    momentum, and Muon weight decay; every Adam beta is a per-role CONSTANT;
    windows are warmup_steps + round(warmdown_ratio * N). Verified
    bitwise-identical to the Ramp implementation it replaced
    (sched_parity_test.py in the session folder).

    `weight_decay` arrives already batch/horizon-scaled. Returns a namespace:
    .matrix (MuonCoeffs) + one AdamWTabs per AdamW role, .adamw_eps, and
    .num_steps. The trainer binds the result to the global `sched`."""
    N = num_iterations
    C = round(warmdown_ratio * N)               # LR warmdown length
    assert warmup_steps + C <= N, f"warmup ({warmup_steps}) + warmdown ({C}) exceed the run ({N})"
    i = np.arange(N, dtype=np.float64)
    cool = slice(N - C + 1, N)                  # the hold covers i <= N - C
    f = (N - i[cool]) / C                       # ~1 -> ~0 across the warmdown

    # The one LR shape for the whole run: linear warmup from 0 (reaching the
    # peak on the warmup window's last step), hold at 1, linear warmdown to
    # final_lr_frac (arriving one step past the run's end -- nanochat's
    # convention). Each role scales it to its own peak below.
    lrm = np.ones(N)
    lrm[:warmup_steps] = (i[:warmup_steps] + 1.0) / warmup_steps
    lrm[cool] = final_lr_frac + (1.0 - final_lr_frac) * f

    # Muon momentum: 0.85 -> 0.97 over 400 steps (the clamp only lets short
    # smoke/debug runs build a valid schedule; identical for N >= ~1150),
    # hold, then cool to 0.90 across the LR warmdown.
    mW = min(400, int(N * (1 - warmdown_ratio)))
    momentum = np.full(N, 0.97)
    momentum[:mW] = 0.85 + (0.97 - 0.85) * (i[:mW] + 1.0) / mW
    momentum[cool] = 0.90 + (0.97 - 0.90) * f

    # Muon weight decay: half-cosine from the peak to zero over the whole run
    # (step 0 sits at the peak; the decay begins at step 1).
    muon_wd = np.empty(N)
    muon_wd[0] = weight_decay
    fw = (N - i[1:]) / N
    muon_wd[1:] = weight_decay * (0.5 * (1.0 + np.cos(math.pi * (1.0 - fw))))

    # Numpy arrays -> fp32 device tables: a step reads its coefficients with
    # an on-device gather, never a host-to-device copy.
    dev = lambda a: torch.tensor(a, dtype=torch.float32, device=device)
    t1 = np.arange(1, N + 1, dtype=np.float64)

    # Fold one AdamW role's schedule down to the kernel's update coefficients.
    # The folding is ONE policy shared by every role (repeating it 6x would
    # obscure edits); the per-role peaks/betas/wd stay visible at the call
    # sites below.
    def adamw(peak, beta1, beta2, wd):
        lr = lrm * peak
        return AdamWTabs(
            wd_mul          = dev(1.0 - lr * wd),
            one_minus_beta1 = dev(np.full(N, 1.0 - beta1)),
            one_minus_beta2 = dev(np.full(N, 1.0 - beta2)),
            rsqrt_bias2     = dev(1.0 / ((1.0 - beta2 ** t1) ** 0.5)),
            step_size       = dev(lr / (1.0 - beta1 ** t1)),
        )

    # Muon's coefficients fold directly from the three shaped schedules.
    # Canonical lr: NO per-bank aspect fold (see § Optimizer Step).
    matrix_lr = lrm * (0.02 * batch_lr_scale)
    matrix = MuonCoeffs(
        momentum           = dev(momentum),
        one_minus_momentum = dev(1.0 - momentum),
        one_minus_beta2    = dev(np.full(N, 1.0 - 0.9)),   # variance-reduction beta2 = 0.9
        lr                 = dev(matrix_lr),
        lr_wd              = dev(matrix_lr * muon_wd),
    )

    # Per-role peak LRs (tuned values). The AdamW peaks were tuned at d12's
    # width, so they carry the 1/sqrt(width ratio) correction to d24.
    adamw_lr_scale = batch_lr_scale * (cfg.d_model / 768) ** -0.5
    return SimpleNamespace(
        matrix       = matrix,
        lm_head      = adamw(0.008 * adamw_lr_scale,        0.8,  0.96,  0.01),
        input_embeds = adamw(0.3 * adamw_lr_scale,          0.8,  0.995, 0.001),
        value_embeds = adamw(0.3 * adamw_lr_scale * 0.5,    0.8,  0.995, 0.01),
        resid        = adamw(0.5 * batch_lr_scale * 0.01,   0.8,  0.95,  0.05),
        x0           = adamw(0.5 * batch_lr_scale,          0.96, 0.95,  0.0),
        smear        = adamw(0.2,                           0.8,  0.95,  0.0),
        adamw_eps    = 1e-10,
        lrm_table    = lrm,   # host-side copy, for logging only
        num_steps    = N,
        batch_lr_scale = batch_lr_scale,   # echoed into the wandb config
        weight_decay   = weight_decay,
    )


# --------------------------------------------------------------------------------
# § Optimizer Code
# --------------------------------------------------------------------------------

# --------------------------------------------------------------------------------
# Mantissa Trick
# Masters use the mantissa trick (Larry Dial via modded-nanogpt train_gpt.py):
# the fp32 master's bit pattern is (live_bf16_bits << 16) | mantissa_uint16.
# Update math runs in fp32 on the reconstructed master; the split back is a
# TRUNCATION (load-bearing: round-to-nearest could carry into the top bits and
# break the lossless live/mantissa pairing).
#
# The bit arithmetic runs in int32 (CUDA has no uint32 shifts as of torch 2.9);
# int32's truncating .to(int16) and the <<16 discard of sign-extension bits
# make it equivalent. Mantissa tensors are STORED uint16, viewed int16 for the
# math.

def fp32_master(live: Tensor, mantissa: Tensor) -> Tensor:
    """Reconstruct the fp32 master from bf16 live bits + stashed mantissa."""
    bits = (live.view(torch.int16).to(torch.int32) << 16) | \
           (mantissa.view(torch.int16).to(torch.int32) & 0xFFFF)
    return bits.view(torch.float32)


def writeback_master(master: Tensor, live: Tensor, mantissa: Tensor) -> None:
    """Truncation split of the updated master back into live + mantissa."""
    bits = master.view(torch.int32)
    live.view(torch.int16).copy_((bits >> 16).to(torch.int16))
    mantissa.view(torch.int16).copy_(bits.to(torch.int16))


# -----------------------------------------------------------------------------
# Fused update kernels. The schedule row is gathered ON DEVICE by `t` -- no
# host involvement per step.

# We use the first five, remainder are just for completeness.
polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


@torch.compile(dynamic=False, fullgraph=True)
def adamw_step_fused_fp32(
    p: Tensor,           # fp32 param, updated IN PLACE (live == master)
    grad: Tensor,
    exp_avg: Tensor,
    exp_avg_sq: Tensor,
    c: AdamWTabs,
    t: Tensor,           # (1,) int64 device tensor - the schedule row to read
    eps: float,
) -> None:
    """AdamW for the fp32-LIVE scalar params (resid/x0 lambdas, smear, backout
    -- ~30 floats). They are exempt from the bf16-live/mantissa scheme: see the
    dtype scheme note in § Model Initialization."""
    grad = grad.to(exp_avg.dtype)
    p.mul_(c.wd_mul[t])
    exp_avg.lerp_(grad, c.one_minus_beta1[t])
    exp_avg_sq.lerp_(grad.square(), c.one_minus_beta2[t])
    denom = exp_avg_sq.sqrt() * c.rsqrt_bias2[t] + eps
    p.sub_(c.step_size[t] * (exp_avg / denom))


@torch.compile(dynamic=False, fullgraph=True)
def adamw_step_fused(
    live: Tensor,        # bf16 live shard
    mantissa: Tensor,    # uint16, same shape
    grad: Tensor,        # gradient shard (fp32, or bf16 for the embeddings)
    exp_avg: Tensor,     # fp32 first moment
    exp_avg_sq: Tensor,  # fp32 second moment
    c: AdamWTabs,        # per-step coefficient tables, device-resident
    t: Tensor,           # (1,) int64 device tensor - the schedule row to read
    eps: float,
) -> None:
    """Fused AdamW step on the reconstructed master."""
    p = fp32_master(live, mantissa)
    grad = grad.to(exp_avg.dtype)  # embeddings hand in bf16 grads; moment math stays fp32
    p.mul_(c.wd_mul[t])
    exp_avg.lerp_(grad, c.one_minus_beta1[t])
    exp_avg_sq.lerp_(grad.square(), c.one_minus_beta2[t])
    denom = exp_avg_sq.sqrt() * c.rsqrt_bias2[t] + eps
    p.sub_(c.step_size[t] * (exp_avg / denom))
    writeback_master(p, live, mantissa)

# The update kernels take explicit per-tensor arguments rather than the model
# object, twice over: (1) at world>1 the SAME kernels run on shard views
# (p[layer_shard_slice] with the shard-size state) rather than on m.X -- an
# object-reading kernel would need a different body per world size; (2) under
# fullgraph compile,
# attribute access on an ad-hoc Python object turns into dynamo guards on
# object identity/attributes -- fragile and recompile-prone next to plain
# tensor arguments.
@torch.compile(dynamic=False, fullgraph=True)
def muon_step_fused(
    grad: Tensor,        # (K, out, in) fp32 gradient shard -- MUTATED (nesterov lerp)
    live: Tensor,        # (K, out, in) bf16 live shard
    mantissa: Tensor,    # (K, out, in) uint16
    frst_mntm: Tensor,   # (K, out, in) fp32
    scnd_mntm: Tensor,   # (K, out, 1) or (K, 1, in) fp32 - factored second moment
    c: MuonCoeffs,       # per-step coefficient tables, device-resident (UNfolded lr)
    t: Tensor,           # (1,) int64 device tensor - the schedule row to read
    ns_steps: int,       # 5 - number of Polar Express iterations
    residual_dim: int,   # -1 or -2 - residual-facing axis; per-neuron mean-square is taken along it
    lr_mul: Tensor,      # (K, 1, 1) fp32 per-slice LR multiplier (aspect scale today)
    wd_mul: Tensor,      # (K, 1, 1) fp32 per-slice WD multiplier
) -> None:
    """Fused Muon step: momentum -> polar_express -> variance_reduction ->
    cautious update on the reconstructed master. The sqrt(fan_out/fan_in)
    aspect scale is NOT in `c` -- it arrives through lr_mul/wd_mul, per slice,
    so the one coefficient table stays valid for every bank."""
    dtype = grad.dtype

    # Nesterov momentum
    frst_mntm.lerp_(grad, c.one_minus_momentum[t].to(dtype))
    g = grad.lerp_(frst_mntm, c.momentum[t].to(dtype))

    # Polar express (orthogonalization)
    X = g.bfloat16()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-6)
    if g.size(-2) > g.size(-1): # Tall matrix
        for a, b, c_ns in polar_express_coeffs[:ns_steps]:
            A = X.mT @ X
            B = b * A + c_ns * (A @ A)
            X = a * X + X @ B
    else: # Wide matrix (original math)
        for a, b, c_ns in polar_express_coeffs[:ns_steps]:
            A = X @ X.mT
            B = b * A + c_ns * (A @ A)
            X = a * X + B @ X
    g = X

    # Variance reduction (NorMuon). The lerp weight stays fp32.
    v_mean = g.float().square().mean(dim=residual_dim, keepdim=True)
    residual_dim_size = g.size(residual_dim)
    v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * residual_dim_size
    v_norm = v_norm_sq.sqrt()
    scnd_mntm.lerp_(v_mean.to(dtype=scnd_mntm.dtype),
                    c.one_minus_beta2[t].to(scnd_mntm.dtype))
    step_size = scnd_mntm.clamp_min(1e-10).rsqrt()
    scaled_sq_sum = (v_mean * residual_dim_size) * step_size.float().square()
    v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt()
    final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10))
    g = g * final_scale.to(g.dtype)

    # Cautious weight decay + master update + truncation split back to live
    p = fp32_master(live, mantissa)
    mask = (g * p) >= 0
    lr = (c.lr[t] * lr_mul).to(g.dtype)
    lr_wd = (c.lr_wd[t] * wd_mul).to(g.dtype)
    p.sub_(lr * g + lr_wd * p * mask)
    writeback_master(p, live, mantissa)


# --------------------------------------------------------------------------------
# § Model Code (Forward/Backward)
# --------------------------------------------------------------------------------
# Handwritten training step: explicit forward + backward (no autograd),
# accumulating into the fp32/bf16 `.grad32` buffers.
#
# Design notes:
# - Attention runs through the raw FA3 ops above, stashing out + LSE.
# - rms_norms: we stash the norm OUTPUT plus the per-vector 1/rms `r`. In
#   output space the backward is dx = r*(dy - y*mean(y*dy)) for ANY eps, so the
#   pre-norm input is never needed. Cheap norms (the MLP-side xm) are
#   recomputed from the stashed pre-norm x1 instead of stashed.
# - Weight-grad matmuls run in bf16, then accumulate upcast into grad32 -- the
#   same numerics autograd produces for a bf16 matmul.
# - loss_scale (1/grad_accum_steps) replaces the loss division of an autograd
#   loop; the returned loss is the plain (unscaled) mean CE for logging.

# Cast shorthands for the bodies below: the fp32 scalars/gates need explicit
# bf16 casts at their use sites (see forward_backward's docstring), and the
# scalar-parameter grad sums accumulate in fp32.
bf16  = lambda x: x.to(torch.bfloat16)
sum32 = lambda x: x.sum(dtype=torch.float32)

# -----------------------------------------------------------------------------
# rms_norm forward/backward in output space
# TODO - Inline at call site. And can one not derive the other?
def _rms_fwd(x):
    """rms_norm over the last dim plus the per-vector 1/rms its backward
    needs, sharing one mean-square. r is fp32 with eps = 2^-23 (fp32 machine
    eps -- the same number compiled F.rms_norm's decomposition uses); y is
    x * r cast back to bf16. Verified bitwise-identical to the F.rms_norm
    form under torch.compile, and the same speed (bench_rms.log; eager ATen
    differs in last-ulp on ~6/1M elements, but every call site is compiled)."""
    r = (x.float().square().mean(dim=-1, keepdim=True) + 2.0 ** -23).rsqrt()
    y = bf16(x.float() * r)
    return y, r

# TODO - Inline.
def _rms_bwd(dy, y, r):
    """dx = r*(dy - y*mean(y*dy)): exact for any eps because r is the forward's
    actual 1/rms and y the actual output (substitute x = y/r in the usual
    form). Math in fp32, result back to bf16."""
    yf, dyf = y.float(), dy.float()
    dx = r * (dyf - yf * (yf * dyf).mean(dim=-1, keepdim=True))
    return bf16(dx)

# TODO - Inline.
def _rms_bwd_scaled(dy, ys, r, s):
    """Backward through ys = s * rms_norm(x), given the SCALED output ys --
    which is exactly what the attention kernel consumed, so it stashes directly
    with no recompute pass. Substituting y = ys/s into _rms_bwd's form:
    dx = r*(s*dy - ys*mean(ys*dy)/s). Exact algebra."""
    yf, dyf = ys.float(), dy.float()
    dx = r * (s * dyf - yf * ((yf * dyf).mean(dim=-1, keepdim=True) / s))
    return bf16(dx)


# -----------------------------------------------------------------------------
# forward_backward

@torch.no_grad()
def forward_backward(idx, targets, cu_seqlens, loss_scale=1.0):
    """One micro-batch: forward, stash, explicit backward into `.grad32`.
    Returns the detached mean CE loss (unscaled; grads carry loss_scale).

    Wrap in torch.compile -- the CE block below is written for inductor's
    fusion; run eager it materializes full (T, d_vocab) fp32 temporaries.

    Activations are bf16 throughout. The live weights are already bf16, so no
    per-use casts; the fp32 scalars need care: indexing a 1-D fp32 bank gives a
    0-dim tensor, which does NOT promote a bf16 tensor (resid/x0 lambdas ride
    as-is), but the (1,)-shaped smear/backout scalars and the smear_gate matrix
    WOULD promote to fp32, so those are cast explicitly."""

    assert idx.ndim == 1
    T = idx.size(0)
    nl = cfg.n_layers
    nh, nkv = cfg.n_q_heads, cfg.n_kv_heads
    dqk, dvo = cfg.d_qk, cfg.d_vo
    half = dqk // 2
    gch = cfg.d_ve_gate

    assert T > 1, "Training forward pass should have T > 1"
    assert T <= m.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {m.cos.size(1)}"
    cos, sin = m.cos[0, :T], m.sin[0, :T]  # (T, 1, half)

    # ==== forward half (mirrors forward() -- keep the two visibly line-parallel) ====
    x = F.embedding(idx, m.input_embeds)         # bf16
    xe, r_e = _rms_fwd(x)                     # post-norm embedding, pre-smear

    # Smear: mix the previous token's embedding into the current position.
    gate = bf16(m.smear_lambda) * torch.sigmoid(
        xe[1:, :cfg.d_smr_gate] @ bf16(m.smear_gate).mT)
    x = torch.cat([xe[:1], xe[1:] + gate * xe[:-1]], dim=0)

    x0 = x
    backout_layer = nl // 2
    x_backout = None
    stash = []
    for i in range(nl):
        x_in = x
        b = m.resid_lambdas[i] * x_in + m.x0_lambdas[i] * x0
        xn, r_xn = _rms_fwd(b)
        q = (xn @ m.W_Q[i].mT).view(T, nh, dqk)
        k = (xn @ m.W_K[i].mT).view(T, nkv, dqk)
        v = (xn @ m.W_V[i].mT).view(T, nkv, dvo)
        j = cfg.ve_index[i]
        if j >= 0:
            ve = F.embedding(idx, m.value_embeds[j]).view(T, nkv, dvo)
            g = 3 * torch.sigmoid(xn[..., :gch] @ m.ve_gate[j].mT)
            v = v + g.unsqueeze(-1) * ve         # ve/g recomputed in backward, not stashed
        q1, q2 = q[..., :half], q[..., half:]
        k1, k2 = k[..., :half], k[..., half:]
        q = torch.cat([q1 * cos + q2 * sin, q1 * (-sin) + q2 * cos], dim=-1)
        k = torch.cat([k1 * cos + k2 * sin, k1 * (-sin) + k2 * cos], dim=-1)
        qn, r_q = _rms_fwd(q)
        kn, r_k = _rms_fwd(k)
        qf = qn * 1.2                            # stash the SCALED q/k (the kernel's inputs);
        kf = kn * 1.2                            # backward folds the 1.2 via _rms_bwd_scaled
        y, lse = flash_attn_varlen_fwd_lse(qf, kf, v, cu_seqlens, cfg.seq_len, cfg.window_sizes[i])
        y = y.contiguous()
        x1 = b + y.view(T, -1) @ m.W_O[i].mT
        xm, _ = _rms_fwd(x1)                  # xm recomputed in backward from stashed x1
        a = F.relu(xm @ m.W_in[i].mT)
        x = x1 + a.square() @ m.W_out[i].mT
        if i == backout_layer:
            x_backout = x
        stash.append(dict(x_in=x_in, xn=xn, r_xn=r_xn, qf=qf, kf=kf, r_q=r_q, r_k=r_k,
                          v=v, y=y, lse=lse, x1=x1, a=a))

    x_pre = x - bf16(m.backout_lambda) * x_backout
    xf, r_f = _rms_fwd(x_pre)

    # lm_head + softcap + CE loss + dlogits, written for inductor's fusion:
    # tcap is an explicit CSE target (materialize once, no tanh recompute in
    # the dz pass), and the onehot is a broadcast compare (a scatter_add here
    # forces an extra full pass over the buffer). Vocab is unpadded by
    # construction, so there is no [:V] cropping anywhere. No pad/ignore
    # machinery either: every target is a real token by construction (the
    # loader packs whole documents; at a doc seam the target is the next
    # doc's BOS), so the mean runs over all T positions and the dz scale is
    # the compile-time constant loss_scale/T rather than a device n_valid.
    softcap = 15.0
    logits = xf @ m.lm_head.mT                   # (T, d_vocab) bf16
    tcap = torch.tanh(logits.float() / softcap)
    cap = softcap * tcap
    tgt = targets.unsqueeze(1)
    cap_y = cap.gather(1, tgt).squeeze(1)
    cmax = cap.amax(dim=1, keepdim=True)
    e = (cap - cmax).exp()
    ssum = e.sum(dim=1, keepdim=True)
    lse_ce = (ssum.log() + cmax).squeeze(1)
    loss = (lse_ce - cap_y).mean()
    onehot = torch.arange(cfg.d_vocab, device=targets.device).unsqueeze(0) == tgt
    dz = bf16((e / ssum - onehot.float()) * (1.0 - tcap * tcap) * (loss_scale / T))
    del logits
    m.lm_head.grad32.add_((dz.mT @ xf).float())
    dxf = dz @ m.lm_head
    del dz

    # ==== backward half ====
    # Bank wgrads add directly into grad32_slices views; only the per-layer
    # scalar sums are collected and landed stacked at the end.
    g_resid = []; g_x0 = []

    d_pre = _rms_bwd(dxf, xf, r_f)
    m.backout_lambda.grad32.add_(-sum32(d_pre * x_backout))
    d_stream = d_pre                             # grad wrt layer nl-1's output
    d_x0 = torch.zeros_like(x0)
    for i in reversed(range(nl)):
        st = stash[i]
        if i == backout_layer:
            # TRAP: x_backout gets an EXTRA contribution when the sweep passes nl//2
            d_stream = d_stream - bf16(m.backout_lambda) * d_pre
        # --- MLP backward (relu^2: dh = 2*a*du, self-masking since a = relu(h)) ---
        x1, a = st["x1"], st["a"]
        d_u = d_stream @ m.W_out[i]
        m.W_out.grad32_slices[i].add_(d_stream.mT @ a.square())
        d_h = 2.0 * a * d_u
        xm, r_xm = _rms_fwd(x1)               # cheap recompute (bitwise: same input)
        m.W_in.grad32_slices[i].add_(d_h.mT @ xm)
        d_xm = d_h @ m.W_in[i]
        d_x1 = d_stream + _rms_bwd(d_xm, xm, r_xm)
        # --- attention backward ---
        xn, y = st["xn"], st["y"]
        m.W_O.grad32_slices[i].add_(d_x1.mT @ y.view(T, -1))
        d_y = (d_x1 @ m.W_O[i]).view(T, nh, dvo)
        dqf, dkf, dv = flash_attn_varlen_bwd(
            d_y, st["qf"], st["kf"], st["v"], y, st["lse"], cu_seqlens, cfg.seq_len,
            cfg.window_sizes[i])
        # per-(token, head) norm backward with the 1.2 scale folded in
        d_qr = _rms_bwd_scaled(dqf, st["qf"], st["r_q"], 1.2)
        d_kr = _rms_bwd_scaled(dkf, st["kf"], st["r_k"], 1.2)
        # rotary backward = rotation by -theta (transpose of the forward rotation)
        dq1, dq2 = d_qr[..., :half], d_qr[..., half:]
        d_q0 = torch.cat([dq1 * cos - dq2 * sin, dq1 * sin + dq2 * cos], dim=-1)
        dk1, dk2 = d_kr[..., :half], d_kr[..., half:]
        d_k0 = torch.cat([dk1 * cos - dk2 * sin, dk1 * sin + dk2 * cos], dim=-1)
        # --- VE gate backward (ve/g recomputed) ---
        j = cfg.ve_index[i]
        d_xn_ve = None
        if j >= 0:
            ve = F.embedding(idx, m.value_embeds[j]).view(T, nkv, dvo)
            sg = torch.sigmoid(xn[..., :gch] @ m.ve_gate[j].mT)
            d_g = (dv * ve).sum(dim=-1)          # (T, n_kv_heads)
            d_zg = d_g * (3 * sg * (1 - sg))
            m.ve_gate.grad32_slices[j].add_(d_zg.mT @ xn[..., :gch])
            d_ve = (dv * (3 * sg).unsqueeze(-1)).reshape(T, nkv * dvo)
            # embedding_dense_backward (autograd's own lowering) beats raw
            # index_add_ atomics ~2x at these shapes -- see the GH200 trace hunt
            m.value_embeds.grad32_slices[j].add_(
                torch.ops.aten.embedding_dense_backward(d_ve, idx, cfg.d_vocab, -1, False))
            d_xn_ve = d_zg @ m.ve_gate[j]
        # dv passes through the VE add unchanged: v = v0 + g*ve
        d_q0 = d_q0.view(T, nh * dqk)
        d_k0 = d_k0.view(T, nkv * dqk)
        d_v0 = dv.reshape(T, nkv * dvo)
        m.W_Q.grad32_slices[i].add_(d_q0.mT @ xn)
        m.W_K.grad32_slices[i].add_(d_k0.mT @ xn)
        m.W_V.grad32_slices[i].add_(d_v0.mT @ xn)
        d_xn = d_q0 @ m.W_Q[i] + d_k0 @ m.W_K[i] + d_v0 @ m.W_V[i]
        if d_xn_ve is not None:
            d_xn[:, :gch] += d_xn_ve
        d_b = d_x1 + _rms_bwd(d_xn, xn, st["r_xn"])
        # --- blend backward: b = resid_lambdas[i]*x_in + x0_lambdas[i]*x0 ---
        g_resid.append(sum32(d_b * st["x_in"]))
        g_x0.append(sum32(d_b * x0))
        d_x0 = d_x0 + m.x0_lambdas[i] * d_b  # TRAP: x0 feeds every layer, accumulate
        d_stream = m.resid_lambdas[i] * d_b
        stash[i] = None                          # free this layer's stash as we go

    # Land the per-layer resid/x0 scalar sums (collected in REVERSED layer
    # order) as one stacked add each.
    m.resid_lambdas.grad32.add_(torch.stack(g_resid[::-1]))
    m.x0_lambdas.grad32.add_(torch.stack(g_x0[::-1]))

    # d_stream is now the grad through layer 0's input, which IS x0 (same tensor)
    d_xs = d_x0 + d_stream                       # grad wrt the smeared embedding
    # --- smear backward: xs = cat([xe[:1], xe[1:] + gate*xe[:-1]]) ---
    sg = torch.sigmoid(xe[1:, :cfg.d_smr_gate] @ bf16(m.smear_gate).mT)  # (T-1, 1), recomputed
    gate = bf16(m.smear_lambda) * sg
    d_xe = d_xs.clone()
    d_xe[:-1] += gate * d_xs[1:]                 # TRAP: shifted scatter -- p's grad reaches p-1
    d_gate = (d_xs[1:] * xe[:-1]).sum(dim=-1, keepdim=True)        # (T-1, 1)
    m.smear_lambda.grad32.add_(sum32(d_gate * sg))
    d_zs = d_gate * bf16(m.smear_lambda) * sg * (1 - sg)
    m.smear_gate.grad32.add_((d_zs.mT @ xe[1:, :cfg.d_smr_gate]).float())
    d_xe[1:, :cfg.d_smr_gate] += d_zs @ bf16(m.smear_gate)
    # --- embedding norm + token embedding scatter ---
    d_emb = _rms_bwd(d_xe, xe, r_e)
    m.input_embeds.grad32.add_(
        torch.ops.aten.embedding_dense_backward(d_emb, idx, cfg.d_vocab, -1, False))

    return loss


# --------------------------------------------------------------------------------
# § Forward-Only
# --------------------------------------------------------------------------------

# Compiled by the trainer: § Main Loop rebinds this name through torch.compile
# (one specialization per shape/targets combination -- val loss and CORE logits).
@torch.no_grad()
def forward(idx, cu_seqlens, targets=None, loss_reduction='mean'):
    """Scoring forward for validation loss and CORE eval: one packed 1D
    sequence of documents with per-document attention isolation via varlen
    flash attention. idx/targets are (T,) and activations stay (T, ...)
    throughout -- the layout the varlen kernel wants. Returns the loss if
    targets are given, else the (softcapped, fp32) logits (T, d_vocab).

    Mirrors forward_backward's forward half line for line -- keep them that
    way; diff them when either changes."""
    assert idx.ndim == 1
    T = idx.size(0)
    D = cfg.d_model
    half = cfg.d_qk // 2

    assert T > 1, "Scoring forward pass should have T > 1 (smear needs a previous token)"
    assert T <= m.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {m.cos.size(1)}"
    cos, sin = m.cos[0, :T], m.sin[0, :T]  # (T, 1, half)

    # Embed the tokens
    x = F.embedding(idx, m.input_embeds)         # bf16
    x = F.rms_norm(x, (D,))

    # Smear: mix the previous token's embedding into the current position.
    gate = bf16(m.smear_lambda) * torch.sigmoid(
        x[1:, :cfg.d_smr_gate] @ bf16(m.smear_gate).mT)
    x = torch.cat([x[:1], x[1:] + gate * x[:-1]], dim=0)

    # Forward the trunk of the Transformer
    x0 = x
    backout_layer = cfg.n_layers // 2
    x_backout = None
    for i in range(cfg.n_layers):
        x = m.resid_lambdas[i] * x + m.x0_lambdas[i] * x0
        # --- attention ---
        xn = F.rms_norm(x, (D,))
        # (T, H, D) - the varlen kernel's native layout, no transpose needed
        q = (xn @ m.W_Q[i].mT).view(T, cfg.n_q_heads,  cfg.d_qk)
        k = (xn @ m.W_K[i].mT).view(T, cfg.n_kv_heads, cfg.d_qk)
        v = (xn @ m.W_V[i].mT).view(T, cfg.n_kv_heads, cfg.d_vo)
        # Value residual (ResFormer): value embedding mixed in via an
        # input-dependent per-head gate, range (0, 3)
        j = cfg.ve_index[i]
        if j >= 0:
            ve = F.embedding(idx, m.value_embeds[j]).view(T, cfg.n_kv_heads, cfg.d_vo)
            g = 3 * torch.sigmoid(xn[..., :cfg.d_ve_gate] @ m.ve_gate[j].mT)
            v = v + g.unsqueeze(-1) * ve
        # Rotary embeddings (relative positional encoding)
        q1, q2 = q[..., :half], q[..., half:]
        k1, k2 = k[..., :half], k[..., half:]
        q = torch.cat([q1 * cos + q2 * sin, q1 * (-sin) + q2 * cos], dim=-1)
        k = torch.cat([k1 * cos + k2 * sin, k1 * (-sin) + k2 * cos], dim=-1)
        # QK norm, then sharper attention (the 1.2 splits the scale between Q and K)
        q = F.rms_norm(q, (cfg.d_qk,)) * 1.2
        k = F.rms_norm(k, (cfg.d_qk,)) * 1.2
        y, _ = flash_attn_varlen_fwd_lse(q, k, v, cu_seqlens, cfg.seq_len, cfg.window_sizes[i])
        x = x + y.contiguous().view(T, -1) @ m.W_O[i].mT
        # --- MLP (relu^2) ---
        x = x + F.relu(F.rms_norm(x, (D,)) @ m.W_in[i].mT).square() @ m.W_out[i].mT
        if i == backout_layer:
            x_backout = x
    # Subtract mid-layer residual to remove low-level features before logit projection
    x = x - bf16(m.backout_lambda) * x_backout
    x = F.rms_norm(x, (D,))

    # lm_head + softcap
    logits = (x @ m.lm_head.mT).float()          # (T, d_vocab)
    logits = 15.0 * torch.tanh(logits / 15.0)    # smoothly cap to [-15, 15]

    if targets is not None:
        # No ignore_index: targets here only ever come from the training/val
        # loader, which never emits pad (see forward_backward's CE note).
        return F.cross_entropy(logits, targets, reduction=loss_reduction)
    return logits


# --------------------------------------------------------------------------------
# § Optimizer Step
# --------------------------------------------------------------------------------
# The written-out step: one fused-kernel call per named tensor, policy at the
# call site, wrapped in the 3-phase comm flow (nanochat train_step.py):
#
#   1. Launch an async grad reduction for every sharded tensor: the full
#      grad32 reduce-scatters into a fresh shard-size buffer, in the grad's
#      dtype (bf16 for the two embedding tables, fp32 for everything else).
#      ReduceOp.AVG across ranks composes with loss_scale=1/grad_accum_steps
#      to make every reduced grad the global-batch mean.
#   2. In launch order (the comm stream completes reduces in that order):
#      wait for the tensor's reduced grad, run its update kernel on the owned
#      shard, then launch the async all-gather that writes the updated bf16
#      live shard back into every rank's full tensor. The gather is IN PLACE
#      -- our slice of the live tensor is the gather source, NCCL's
#      sanctioned in-place form; even divisibility (§ Shard Assignment) means
#      no padded staging buffer and no crop afterwards. Each gather overlaps
#      the updates that follow it. Replicated params (ve_gate, the fp32
#      scalars) ride along inline: plain all_reduce, then the identical
#      full-size update on every rank.
#   3. Wait out the gathers.
#
# Waits are stream waits, not host syncs -- the whole step stays async on the
# host, and t_step still advances on-device. At world_size == 1 every
# collective short-circuits and every shard view is the whole tensor: one
# code path, degenerate comm, numerics identical to the validated single-GPU
# step.
#
# NOTE: the world>1 path has not run yet (the reference's comm code never ran
# at world>1 either) -- it awaits an 8-GPU validation pass.

ns_steps = 5  # Polar Express iterations per Muon step

# Per-slice Muon LR/WD multipliers: each bank's sqrt(max(1, fan_out/fan_in))
# aspect scale -- Muon's tall-matrix correction -- kept OUT of the shared
# matrix table so that table stays one set of numbers valid for every bank.
# At d24 only W_in is non-square, so only it gets a real multiplier (2.0).
# TODO(Chris) - I'd like to drop this eventually. If/when we drop the 2x on
#        W_in we'll probably take a hit, since everything else is tuned around
#        it. I don't think the trick is principled--in modded-nanogpt I
#        accidentally flipped it to 2x on the mlp output and it improved loss;
#        Karpathy tried that on nanochat and it didn't help. I think the model
#        mostly adapts to it, so it's not worth the hassle. Get things working
#        as-is first, though.
mul_unit    = torch.full((cfg.n_layers, 1, 1), 1.0, dtype=torch.float32, device=device)  # W_Q/W_K/W_V/W_O (square), W_out (wide -> clamped)
mul_W_in    = torch.full((cfg.n_layers, 1, 1), (cfg.d_mlp / cfg.d_model) ** 0.5,
                         dtype=torch.float32, device=device)                             # 2.0 (4x expansion, tall)
mul_ve_unit = torch.full((cfg.num_ves,  1, 1), 1.0, dtype=torch.float32, device=device)  # ve_gate (square)

# THE schedule position: one (1,) int64 device tensor, advanced on-device at
# the end of optimizer_step -- the host never syncs on it.
t_step = torch.zeros(1, dtype=torch.int64, device=device)


@torch.no_grad()
def optimizer_step():
    """One explicit optimizer step, written out per named tensor. Reads the
    global `sched` (bind build_schedules' result to `sched` before training).
    Muon MUTATES the grad it is handed (nesterov lerp) -- grad32 itself at
    world=1, the reduce-scattered shard at world>1 -- so zero every grad32
    afterwards either way (the loop in § Main Loop does).
    no_grad is load-bearing for the fp32 scalar kernel's in-place leaf updates
    (the mantissa kernels only dodge autograd's leaf check via their int
    views)."""
    eps = sched.adamw_eps

    # ---- Phase 1: launch every async grad reduction --------------------------
    # Fresh shard buffers each step (the caching allocator makes this free);
    # the state tensors already carry the shard geometry, so empty_like is the
    # whole allocation story.
    reduced = {}   # tensor -> (async work handle, shard-size reduced grad)
    if world_size > 1:
        for p in (m.W_Q, m.W_K, m.W_V, m.W_O, m.W_in, m.W_out):
            g_shard = torch.empty_like(p.frst_mntm)                        # (layer shard, out, in) fp32
            reduced[p] = (dist.reduce_scatter_tensor(g_shard, p.grad32, op=dist.ReduceOp.AVG, async_op=True), g_shard)
        for p in (m.lm_head, m.input_embeds, m.value_embeds):
            g_shard = torch.empty_like(p.exp_avg, dtype=p.grad32.dtype)    # (row shard, cols) in the grad's dtype
            reduced[p] = (dist.reduce_scatter_tensor(g_shard, p.grad32.view(-1, p.shape[-1]), op=dist.ReduceOp.AVG, async_op=True), g_shard)

    # ---- Phase 2: wait -> owned-shard update -> gather the live shard --------
    gathers = []

    # Muon banks, sharded over layers
    for p, mul in ((m.W_Q, mul_unit), (m.W_K, mul_unit), (m.W_V, mul_unit),
                   (m.W_O, mul_unit), (m.W_in, mul_W_in), (m.W_out, mul_unit)):
        if world_size > 1:
            work, grad = reduced[p]
            work.wait()
        else:
            grad = p.grad32
        muon_step_fused(grad, p[layer_shard_slice], p.mantissa, p.frst_mntm, p.scnd_mntm,
                        sched.matrix, t_step, ns_steps, p.residual_dim,
                        mul[layer_shard_slice], mul[layer_shard_slice])
        if world_size > 1:
            gathers.append(dist.all_gather_into_tensor(p, p[layer_shard_slice], async_op=True))

    # Muon replicated: ve_gate is tiny, every rank updates all of it
    if world_size > 1:
        dist.all_reduce(m.ve_gate.grad32, op=dist.ReduceOp.AVG)
    muon_step_fused(m.ve_gate.grad32, m.ve_gate, m.ve_gate.mantissa, m.ve_gate.frst_mntm, m.ve_gate.scnd_mntm, sched.matrix, t_step, ns_steps, m.ve_gate.residual_dim, mul_ve_unit, mul_ve_unit)

    # AdamW, sharded over vocab rows. value_embeds' state is shaped over the
    # flattened (ve_slot * vocab) row axis, so live/grad pass 2-D views
    # throughout (a no-op reshape for the two already-2-D tables).
    # The roles differ only in their tables (peaks/betas: build_schedules):
    # lm_head runs the coolest peak (~40x below the embeddings); input_embeds
    # the hottest, with the heaviest second-moment smoothing (beta2 .995);
    # value_embeds rides the embedding schedule at half peak and 10x the decay.
    for p, table, row_shard in ((m.lm_head,      sched.lm_head,      vocab_shard_slice),
                                (m.input_embeds, sched.input_embeds, vocab_shard_slice),
                                (m.value_embeds, sched.value_embeds, ve_row_shard_slice)):
        rows = p.view(-1, p.shape[-1])
        if world_size > 1:
            work, grad = reduced[p]
            work.wait()
        else:
            grad = p.grad32.view(-1, p.shape[-1])
        adamw_step_fused(rows[row_shard], p.mantissa, grad, p.exp_avg, p.exp_avg_sq, table, t_step, eps)
        if world_size > 1:
            gathers.append(dist.all_gather_into_tensor(rows, rows[row_shard], async_op=True))

    # AdamW replicated scalars (fp32-live, no mantissa). Three schedule
    # flavors: resid -- the gentlest peak and the only decayed scalars (wd
    # .05); x0 -- the hottest peak with a slow first moment (beta1 .96);
    # smear -- one flat middling peak shared by all three smear/backout
    # scalars, no decay. (Peaks/betas: build_schedules.)
    if world_size > 1:
        for p in (m.resid_lambdas, m.x0_lambdas, m.smear_gate, m.smear_lambda, m.backout_lambda):
            dist.all_reduce(p.grad32, op=dist.ReduceOp.AVG)
    adamw_step_fused_fp32(m.resid_lambdas,  m.resid_lambdas.grad32,  m.resid_lambdas.exp_avg,  m.resid_lambdas.exp_avg_sq,  sched.resid, t_step, eps)
    adamw_step_fused_fp32(m.x0_lambdas,     m.x0_lambdas.grad32,     m.x0_lambdas.exp_avg,     m.x0_lambdas.exp_avg_sq,     sched.x0,    t_step, eps)
    adamw_step_fused_fp32(m.smear_gate,     m.smear_gate.grad32,     m.smear_gate.exp_avg,     m.smear_gate.exp_avg_sq,     sched.smear, t_step, eps)
    adamw_step_fused_fp32(m.smear_lambda,   m.smear_lambda.grad32,   m.smear_lambda.exp_avg,   m.smear_lambda.exp_avg_sq,   sched.smear, t_step, eps)
    adamw_step_fused_fp32(m.backout_lambda, m.backout_lambda.grad32, m.backout_lambda.exp_avg, m.backout_lambda.exp_avg_sq, sched.smear, t_step, eps)

    # ---- Phase 3: wait out the live all-gathers ------------------------------
    for work in gathers:
        work.wait()

    t_step.add_(1)  # advance the schedule on-device

# Model + optimizer state is CAPTURED to disk at cfg.save_steps and at the
# final step (write_checkpoint, below the seam): live weights, masters via
# mantissa, both optimizers' moments, and the step counter -- world-agnostic.
# There is still deliberately no LOAD path (runs start from scratch, see the
# design decisions at the top); resume arrives with the load half when first
# needed.

##########################################################################################
#  Code below comes from the 'stacks' repo
#  I pulled it mainly for:
#    - Pre-tokenized data, and the distributed data loader
#    - Simplified (maybe?) CORE eval code
# 
##########################################################################################

# --------------------------------------------------------------------------------
# § Dataset Download
# --------------------------------------------------------------------------------

NUM_TRAIN_SHARDS = 80   # full 5,568-step horizon: 70 (downloads shards 1-69,
                        # 6.9B raw ~= 6.1B usable after seq_len truncation --
                        # see the token-floor assert below the seam; 91 shards
                        # of 100M raw tokens are on the hub)
#DATASET_NAME = "fineweb_edu_32k_8_370"
DATASET_NAME = "climbmix_32k_8_170"
# Subdir for PT train/val .bin shards
#PT_DATA_SUBDIR = "fineweb_edu" 
PT_DATA_SUBDIR = "climbmix" 
HF_REPO_ID = f"ChrisMcCormick/{DATASET_NAME}"
_data_path = os.environ.get("DATA_PATH", ".")
DATASET_DIR = os.path.join(_data_path, f"data/{DATASET_NAME}")
_config_path = os.path.join(DATASET_DIR, "config.json")
train_files = os.path.join(DATASET_DIR, f"{PT_DATA_SUBDIR}/train_*.bin")
val_files   = os.path.join(DATASET_DIR, f"{PT_DATA_SUBDIR}/val_*.bin")

if master_process:
    from huggingface_hub import HfApi, hf_hub_download, login
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)
    os.makedirs(DATASET_DIR, exist_ok=True)
    api = HfApi()
    train_prefix = f"{PT_DATA_SUBDIR}/train_"
    to_download = []
    for fname in api.list_repo_files(repo_id=HF_REPO_ID, repo_type="dataset"):
        if fname.startswith(train_prefix) and int(fname[len(train_prefix):].split(".")[0]) >= NUM_TRAIN_SHARDS:
            continue
        if not os.path.exists(os.path.join(DATASET_DIR, fname)):
            to_download.append(fname)
    if to_download:
        print(f"=== Downloading {len(to_download)} files from {HF_REPO_ID} ===")
        for fname in to_download:
            hf_hub_download(repo_id=HF_REPO_ID, filename=fname, repo_type="dataset", local_dir=DATASET_DIR)
        print("  Done.")
dist.barrier()

# Load vocab config
with open(_config_path) as f:
    _vocab_config = json.load(f)
VOCAB_SIZE = _vocab_config["vocab_size"]
BOS_ID = _vocab_config["bos_id"]
assert VOCAB_SIZE == cfg.d_vocab, \
    f"dataset vocab ({VOCAB_SIZE}) != model d_vocab ({cfg.d_vocab}) -- wrong dataset for this hardcoded model"

# --------------------------------------------------------------------------------
# § Distributed Data Loader
# --------------------------------------------------------------------------------
# Based on the dataloader from modded-nanogpt.
# - Designed for use with flashattention_varlen_func, meaning it returns a packed token
#   buffer of sequences and their lengths via cu_seqlens.
# - Hardcoded for single-epoch training.
# - Compared to `modded`, it does not support changing batch size mid-training.

def _load_data_shard(file: Path):
    header = torch.from_file(str(file), False, 256, dtype=torch.int32) # header is 256 int32
    assert header[0] == 20240520, "magic number mismatch in the data .bin file"
    assert header[1] == 1, "unsupported version"
    num_tokens = int(header[2]) # number of tokens (claimed)
    with file.open("rb", buffering=0) as f:
        tokens = torch.empty(num_tokens, dtype=torch.uint16, pin_memory=True) # avoid pin_memory copy by @YouJiacheng
        f.seek(256 * 4)
        nbytes = f.readinto(tokens.numpy()) # avoid bytes->array copy by @YouJiacheng
        assert nbytes == 2 * num_tokens, "number of tokens read does not match header"
    return tokens

class Shard:
    def __init__(self, tokens: Tensor, world_size: int = 1):
        self.tokens = tokens
        self.size = tokens.numel()
        self.world_size = world_size
        self.i = 0

        # Partial index now, full index async
        self.bos_idx = (tokens[:6_000_000] == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self._full_idx = None
        self._loader_thread = None
        self._ready = threading.Event()
        self._loader_thread = threading.Thread(target=self._scan)
        self._loader_thread.start()

    def _scan(self):
        self._full_idx = (self.tokens == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self._ready.set()

    def _maybe_switch(self):
        # Switch to full index as soon as async scan completes
        if self.bos_idx is not self._full_idx and self._ready.is_set():
            self._loader_thread.join()
            self.bos_idx = self._full_idx

    def next_batch(self, num_tokens_local: int, max_seq_len: int):
        """Returns (starts, ends) per rank, or None if this shard is exhausted."""
        self._maybe_switch()
        n = len(self.bos_idx)
        starts = [[] for _ in range(self.world_size)]
        ends = [[] for _ in range(self.world_size)]

        idx = self.i
        for r in range(self.world_size):
            cur_len = 0
            while cur_len <= num_tokens_local:
                if idx >= n:
                    return None
                cur = self.bos_idx[idx]
                starts[r].append(cur)
                end = min(self.bos_idx[idx + 1] if idx + 1 < n else self.size,
                          cur + max_seq_len,
                          cur + num_tokens_local - cur_len + 1)
                ends[r].append(end)
                cur_len += end - cur
                idx += 1

            assert cur_len == num_tokens_local + 1
        self.i = idx
        return starts, ends

    @staticmethod
    def load_async(file: Path, world_size: int = 1):
        """Returns getter function for async shard loading"""
        result = {}
        ready = threading.Event()
        def load():
            tokens = _load_data_shard(file)
            result['shard'] = Shard(tokens, world_size)
            ready.set()
        thread = threading.Thread(target=load)
        thread.start()
        def get():
            ready.wait()
            thread.join()
            return result['shard']
        return get

def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_len: int, grad_accum_steps: int = 1):
    """
    Generator (i.e., yields rather than returns) of the token ids for a
    micro-batch: num_tokens / grad_accum_steps / world_size tokens per yield
    (32,768 for the d24 spec: total batch 2^20, grad accum 32 at world=1).
    Provides both the input and target ids.
    Sequences are BOS-aligned and only returned from their beginning; tokens
    past max_seq_len are discarded (the next sequence starts at the next BOS).
    Also used for validation batches.
    Args:
        filename_pattern: pattern to match the dataset .bin shard files
        num_tokens:       tokens per full batch (2^20 for training)
        max_seq_len:      2048
        grad_accum_steps: micro-batches per full batch
    """
    # This GPU's rank and total GPU count.
    rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    
    # Confirm it all divides evenly, then calculate the per-GPU micro-batch size.
    assert num_tokens % (world_size * grad_accum_steps) == 0, "Batch size must be divisible by world size"    
    num_tokens_local = num_tokens // grad_accum_steps // world_size

    # cu_seqlens is FIXED SIZE (the compiled graph needs one shape), and ghost
    # entries cost real FA3 varlen overhead, so it is sized to the DATA rather
    # than a rounded guess: the densest run of climbmix docs packs 82 into one
    # 32,768-token micro-batch (measured -- scan_max_docs.py; an upper bound,
    # since batches can only start where the previous one ended). 96 gives
    # ~17% headroom (nanochat's own estimate for these shapes also lands on
    # 96), and the overflow assert below fails loudly rather than corrupt if
    # the data ever changes.
    max_num_docs = 192
    
    # Get the list of shard files and wrap in an iterator.
    files = [Path(file) for file in sorted(glob.glob(filename_pattern))]
    if not files:
        raise FileNotFoundError(f"No files found for pattern: {filename_pattern}")
    file_iter = iter(files)

    # Load the first shard.
    tokens = _load_data_shard(next(file_iter))
    
    shard = Shard(tokens, world_size)
    remaining_files = list(file_iter)
    next_shard_idx = 0
    next_shard_getter = Shard.load_async(remaining_files[0], world_size) if remaining_files else None

    while True:
        # Get the start and end indices (within `tokens`) of the sequences to use for 
        # the current micro-batch.
        result = shard.next_batch(num_tokens_local, max_seq_len)
        
        # If this shard is exhausted,
        if result is None:
            # If there are no more shards, kill the dataloader.
            if next_shard_getter is None:
                return 

            # Load the next shard.
            shard = next_shard_getter()
            tokens = shard.tokens
            next_shard_idx += 1
            next_shard_getter = Shard.load_async(remaining_files[next_shard_idx], world_size) if next_shard_idx < len(remaining_files) else None
            
            # Re-start the loop.
            continue

        # Locations of the documents in `tokens`. Only specifies the
        # number of documents needed, not max.
        start_idxs = torch.tensor(result[0][rank])
        end_idxs = torch.tensor(result[1][rank])
        
        # `tokens` contains the entire shard. The sequences defined by the starts and ends
        # may or may not be contiguous within `tokens`, due to some sequences being
        # truncated, so we slice them and then re-concatenate into a single tensor. 
        buf = torch.cat([tokens[i:j] for i, j in zip(start_idxs, end_idxs)])
        
        # `buf` contains `num_tokens_local + 1` tokens to allow for the inputs vs.
        # targets offset.
        _inputs = buf[:-1] # All tokens minus the last
        _targets = buf[1:] # Shift the tokens to the left, so that targets contains the 
                           # next token for each input token. 

        # The final document includes an extra token that is the target of the last 
        # token in the last document. Now that we have our `_targets`, we can remove it.
        end_idxs[-1] -= 1  

        # Calculate the start indices of the documents within `_inputs`. (flashattention
        # start_idxs are relative to the `tokens` buffer, so we convert them by
        # accumulating the document lengths.  
        # cum_lengths starts with the second document, so we'll shift 
        cum_lengths = (end_idxs - start_idxs).cumsum(0)

        # One entry per doc plus the leading 0 must fit the fixed buffer.
        assert len(cum_lengths) < max_num_docs, \
            f"micro-batch packed {len(cum_lengths)} docs; cu_seqlens holds only {max_num_docs}"

        # The actual cu_seqlens array always needs to contain `max_num_docs` elements so we
        # the compiler can build a single graph.
        # We allocate that buffer here and fill it with "empty documents", i.e., setting their start index
        # to one past the end of the `_inputs` buffer.
        _cum_lengths = torch.full((max_num_docs,), num_tokens_local)
        
        # Then copy in the lengths, inserting the first document (index 0).
        _cum_lengths[0] = 0
        _cum_lengths[1:len(cum_lengths) + 1] = cum_lengths

        # Cast to int32 / int64 on the CPU before transfer to avoid dtype conversion during .to()
        _inputs = _inputs.to(dtype=torch.int32)
        _targets = _targets.to(dtype=torch.int64)
        _cum_lengths = _cum_lengths.to(dtype=torch.int32)

        yield (
            _inputs.to(device="cuda", non_blocking=True),
            _targets.to(device="cuda", non_blocking=True),
            _cum_lengths.to(device="cuda", non_blocking=True),
        )
        # Execution resumes here on the next call.

# --------------------------------------------------------------------------------
# § CORE Evaluation
# --------------------------------------------------------------------------------

# TODO - I think we can move this to a 'core_eval.py' file, I'm no longer as 
#        committed to the end-to-end single file approach.

"""
CORE evaluation using pre-tokenized benchmark data.

The CORE metric (from the DCLM paper, https://arxiv.org/abs/2406.11794) evaluates
a base model on in-context learning tasks using logit-based scoring (no generation).

Pre-tokenized .pt files are produced by data/core_dataset.py and loaded at eval time.
Sequences are packed into fixed-size 1D buffers with cu_seqlens marking boundaries,
enabling batched evaluation through the compiled varlen flash attention m.
"""

# -----------------------------------------------------------------------------
# Packed CORE evaluation: batch multiple examples into fixed-length 1D buffers

def pack_for_eval(sequences, buffer_size):
    """
    Pack pre-tokenized sequences into fixed-size 1D buffers for batched evaluation.

    Args:
        sequences: list of (tokens, start_idx, end_idx, example_idx, seq_idx_within_example)
        buffer_size: fixed buffer size (must be multiple of 16)

    Returns:
        list of dicts with keys: input_ids, cu_seqlens, metadata
    """
    assert buffer_size % 16 == 0
    # CORE eval sequences can be short (~50-200 tokens), so allow many more per buffer
    # than training's //300 estimate. Use //8 for generous headroom (memory is negligible).
    max_num_seqs = next_multiple_of_n(buffer_size // 8, n=128)

    buffers = []
    cur_tokens = []
    cur_cu = [0]
    cur_meta = []
    cur_pos = 0

    for tokens, start_idx, end_idx, example_idx, seq_idx in sequences:
        seq_len = len(tokens)
        if seq_len > buffer_size:
            continue  # should not happen after truncation

        if cur_pos + seq_len > buffer_size:
            # Finalize current buffer
            _finalize_eval_buffer(buffers, cur_tokens, cur_cu, cur_meta,
                                  buffer_size, max_num_seqs)
            cur_tokens, cur_cu, cur_meta, cur_pos = [], [0], [], 0

        # Track answer span in global buffer coordinates
        global_start = cur_pos + start_idx
        global_end = cur_pos + end_idx
        cur_meta.append((example_idx, seq_idx, global_start, global_end))
        cur_tokens.extend(tokens)
        cur_pos += seq_len
        cur_cu.append(cur_pos)

    if cur_tokens:
        _finalize_eval_buffer(buffers, cur_tokens, cur_cu, cur_meta,
                              buffer_size, max_num_seqs)

    return buffers


def _finalize_eval_buffer(buffers, cur_tokens, cur_cu, cur_meta,
                          buffer_size, max_num_seqs):
    """Pad and finalize a packed eval buffer."""
    total_packed = len(cur_tokens)
    pad_count = buffer_size - total_packed

    # Input tokens: packed sequences + BOS padding
    input_ids = torch.full((buffer_size,), BOS_ID, dtype=torch.int32)
    input_ids[:total_packed] = torch.tensor(cur_tokens, dtype=torch.int32)

    # cu_seqlens: [0, end1, end2, ..., total_packed, buffer_size, buffer_size, ...]
    if pad_count > 0:
        cur_cu.append(buffer_size)  # ghost sequence for padding region
    cu_seqlens = torch.full((max_num_seqs,), buffer_size, dtype=torch.int32)
    cu_seqlens[:len(cur_cu)] = torch.tensor(cur_cu, dtype=torch.int32)

    buffers.append({
        'input_ids': input_ids,
        'cu_seqlens': cu_seqlens,
        'metadata': cur_meta,
    })

# TODO - The FUCK is this?? Hahaha. Typical. Screenshotting for Twitter.
@torch.no_grad()
def forward_eval_packed(input_ids, cu_seqlens):
    """
    Forward a packed 1D eval buffer through the model's scoring forward.
    Returns (softcapped, fp32) logits of shape (buffer_size, vocab_size).
    """
    return forward(input_ids, cu_seqlens)


@torch.no_grad()
def evaluate_task_packed(task_data, buffer_size=cfg.eval_buffer_tokens):
    """Evaluate one task using pre-tokenized sequences and packed batched evaluation."""
    rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    task_type = task_data['task_type']
    num_examples = task_data['num_examples']
    all_sequences = task_data['sequences']
    num_seqs_per_example = task_data['num_seqs_per_example']
    gold_labels = task_data['gold_labels']

    # Step 1: Select this rank's share of pre-tokenized sequences
    rank_examples = set(range(rank, num_examples, world_size))
    sequences = [
        (s['tokens'], s['start_idx'], s['end_idx'], s['example_idx'], s['seq_idx'])
        for s in all_sequences if s['example_idx'] in rank_examples
    ]

    # Step 2: Pack into fixed-size buffers
    packed_buffers = pack_for_eval(sequences, buffer_size)

    # Step 3: Forward pass each buffer and collect per-sequence results
    seq_results = {}

    for buf in packed_buffers:
        input_ids = buf['input_ids'].to(device)
        cu_seqlens = buf['cu_seqlens'].to(device)

        logits = forward_eval_packed(input_ids, cu_seqlens)

        # Per-position losses: loss[j] = -log p(input_ids[j+1] | context up to j)
        target_ids = torch.roll(input_ids.long(), shifts=-1)
        all_losses = F.cross_entropy(logits.float(), target_ids, reduction='none')
        all_predictions = logits.argmax(dim=-1)

        for example_idx, seq_idx, gs, ge in buf['metadata']:
            # Answer span [gs, ge): logits at [gs-1, ge-1) predict tokens at [gs, ge)
            seq_results[(example_idx, seq_idx)] = {
                'losses': all_losses[gs - 1 : ge - 1],
                'predictions': all_predictions[gs - 1 : ge - 1],
                'input_ids': input_ids[gs : ge].long(),
            }

    # Step 4: Evaluate per-example correctness
    correct = torch.zeros(num_examples, dtype=torch.float32, device=device)

    for idx in range(rank, num_examples, world_size):
        if task_type == 'language_modeling':
            r = seq_results[(idx, 0)]
            is_correct = torch.all(r['predictions'] == r['input_ids']).item()
        elif task_type in ['multiple_choice', 'schema']:
            mean_losses = []
            for seq_j in range(num_seqs_per_example[idx]):
                r = seq_results[(idx, seq_j)]
                mean_losses.append(r['losses'].mean().item())
            pred_idx = mean_losses.index(min(mean_losses))
            is_correct = pred_idx == gold_labels[idx]
        else:
            raise ValueError(f"Unsupported task type: {task_type}")
        correct[idx] = float(is_correct)

    if world_size > 1:
        dist.barrier()
        dist.all_reduce(correct, op=dist.ReduceOp.SUM)
    return correct.mean().item()


@torch.no_grad()
def evaluate_chat_task_packed(task_data, buffer_size=cfg.eval_buffer_tokens):
    """Evaluate one chat categorical task using packed batched evaluation.

    Unlike CORE eval (which compares losses across multiple sequences per example),
    chat eval checks single-token logits at the answer position against letter choices.
    Each sequence ends with the prompt (including <|assistant_start|>), and we check
    what the model predicts as the next token, restricted to the valid answer letters.
    """
    rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    all_sequences = task_data['sequences']
    num_examples = task_data['num_examples']

    # Step 1: Select this rank's share and convert to pack_for_eval format.
    # We store answer_pos as start_idx (end_idx = start_idx + 1 for tuple compat)
    # and keep letter_token_ids / gold in a side table.
    sequences = []
    example_meta = {}  # example_idx -> (letter_token_ids, gold)
    for s in all_sequences:
        idx = s['example_idx']
        if idx % world_size != rank:
            continue
        answer_pos = s['answer_pos']
        sequences.append((s['tokens'], answer_pos, answer_pos + 1, idx, 0))
        example_meta[idx] = (s['letter_token_ids'], s['gold'])

    # Step 2: Pack into fixed-size buffers (reuse CORE eval packing infrastructure)
    packed_buffers = pack_for_eval(sequences, buffer_size)

    # Step 3: Forward pass each buffer and score
    correct = 0
    total = 0

    for buf in packed_buffers:
        input_ids = buf['input_ids'].to(device)
        cu_seqlens = buf['cu_seqlens'].to(device)

        logits = forward_eval_packed(input_ids, cu_seqlens)

        for example_idx, seq_idx, gs, ge in buf['metadata']:
            # gs = global position of answer_pos in the buffer.
            # logits[gs] predicts the token AFTER position gs — i.e. the assistant's answer.
            # (This differs from CORE's logits[gs-1:ge-1] convention because here the
            # answer token is NOT in the sequence — we want what the model predicts next.)
            answer_logits = logits[gs]  # (vocab_size,)
            letter_ids, gold = example_meta[example_idx]
            focus_logits = answer_logits[letter_ids]  # (num_choices,)
            pred = focus_logits.argmax().item()
            correct += int(pred == gold)
            total += 1

    # Step 4: Aggregate across ranks
    if world_size > 1:
        correct_t = torch.tensor([correct], dtype=torch.long, device=device)
        total_t = torch.tensor([total], dtype=torch.long, device=device)
        dist.all_reduce(correct_t, op=dist.ReduceOp.SUM)
        dist.all_reduce(total_t, op=dist.ReduceOp.SUM)
        correct = correct_t.item()
        total = total_t.item()

    return correct / total if total > 0 else 0.0


def evaluate_chat_categorical():
    """
    Evaluate a chat model on categorical benchmarks (MMLU, ARC-Easy, ARC-Challenge)
    using pre-tokenized data from chat_eval_dataset.py.
    Returns dict with results, centered_results, and chatcore_metric.
    """
    chat_eval_dir = os.path.join(DATASET_DIR, "chat_eval")
    config_path = os.path.join(chat_eval_dir, "config.json")

    assert os.path.exists(config_path), f"Chat eval config not found: {config_path}"

    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)

    # Evaluate each task
    results = {}
    centered_results = {}
    for task_info in config['tasks']:
        torch.cuda.synchronize()
        start_time = time.time()
        label = task_info['label']

        pt_path = os.path.join(chat_eval_dir, task_info['file'])
        assert os.path.exists(pt_path), f"Chat eval data not found: {pt_path}"
        task_data = torch.load(pt_path, weights_only=False)
        print0(f"Chat eval: {label} ({task_data['num_examples']} examples)... ", console=True)

        accuracy = evaluate_chat_task_packed(task_data)
        torch.cuda.synchronize()
        results[label] = accuracy
        random_baseline = task_data['random_baseline']
        centered_result = (accuracy - random_baseline) / (1.0 - random_baseline)
        centered_results[label] = centered_result
        elapsed = time.time() - start_time
        print0(f"accuracy: {accuracy:.4f} | centered: {centered_result:.4f} | time: {elapsed:.2f}s", console=True)

    chatcore_metric = sum(centered_results.values()) / len(centered_results)
    out = {
        "results": results,
        "centered_results": centered_results,
        "chatcore_metric": chatcore_metric,
    }
    return out


def evaluate_core():
    """
    Evaluate a base model on the CORE benchmark using pre-tokenized data.
    Returns dict with results, centered_results, and core_metric.
    """
    core_eval_dir = os.path.join(DATASET_DIR, "core_eval")
    config_path = os.path.join(core_eval_dir, "config.json")

    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)

    # Evaluate each task
    results = {}
    centered_results = {}
    for task_info in config['tasks']:
        torch.cuda.synchronize()
        start_time = time.time()
        label = task_info['label']

        task_data = torch.load(os.path.join(core_eval_dir, task_info['file']),
                               weights_only=False)
        print0(f"Evaluating: {label} ({task_data['task_type']}, "
               f"{task_data['num_examples']} examples)... ", console=True)

        accuracy = evaluate_task_packed(task_data)
        torch.cuda.synchronize()
        results[label] = accuracy
        random_baseline = task_data['random_baseline']
        centered_result = (accuracy - 0.01 * random_baseline) / (1.0 - 0.01 * random_baseline)
        centered_results[label] = centered_result
        elapsed = time.time() - start_time
        print0(f"accuracy: {accuracy:.4f} | centered: {centered_result:.4f} | time: {elapsed:.2f}s", console=True)

    core_metric = sum(centered_results.values()) / len(centered_results)
    out = {
        "results": results,
        "centered_results": centered_results,
        "core_metric": core_metric
    }
    return out

# --------------------------------------------------------------------------------
# § Main Loop
# --------------------------------------------------------------------------------
# Modeled on nanochat base_train (branch fwd-bwd) -- the flat trainer over the
# same forward_backward / optimizer_step API. No warmup-and-reset phase (that
# trick needs the state_dict save/restore this file deliberately lacks):
# compilation happens during the first real steps, and the time totals simply
# exclude the first 10 steps (the nanochat convention).

# begin logging
logfile = None
if master_process:
    run_id = cfg.run_id
    os.makedirs("logs", exist_ok=True)
    logfile = f"logs/{run_id}.txt"
    print(logfile)
def print0(s="", console=False):
    if master_process:
        with open(logfile, "a") as f:
            if console:
                print(s)
            print(s, file=f)

print0(code)
print0("="*100)
print0(f"Running Python {sys.version}")
print0(f"Running PyTorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}")

# -----------------------------------------------------------------------------
# Model stats, for MFU and the wandb config -- CONSTANTS; scaling.py recomputes
# them from the d24 shapes (params by group, 6 FLOPs per matmul-weight param
# plus the windowed attention term).

num_params          = 1_384_122_122   # every trained weight (Model.weight_names)
num_flops_per_token = 4_860_160_128   # 6 * 729,810,624 matmul params + attention

gpu_device_name = torch.cuda.get_device_name(0)
gpu_peak_flops = next((v for k, v in PEAK_FLOPS.items() if k in gpu_device_name.upper()),
                      float("inf"))
print0(f"Model parameters: {num_params:,} | FLOPs/token: {num_flops_per_token:e}", console=True)
print0(f"GPU: {gpu_device_name} | Peak FLOPS (BF16): {gpu_peak_flops:.2e}", console=True)
print0(f"Total batch size: {cfg.total_batch_size:,} tokens = {cfg.micro_batch_tokens:,} tokens/micro "
       f"x {world_size} ranks x {grad_accum_steps} grad accum", console=True)

# -----------------------------------------------------------------------------
# Schedules: every LR/beta/WD coefficient for the whole run, materialized up
# front. The two batch/horizon corrections are hardcoded (derivations in
# scaling.py):
#   batch_lr_scale = sqrt(2^20 / 2^19) = 1.4142...  -- eta ∝ sqrt(B/B_ref),
#     B_ref = 2^19 where the d12 LRs were tuned; build_schedules applies it to
#     the per-role peaks itself (do NOT also fold it into the LRs).
#   weight_decay = 0.28 * sqrt(2) * (d12/d24 scaling params) = 0.059738 -- the
#     T_epoch framework; matches nanochat's d24 printout exactly.
sched = build_schedules(cfg.num_iterations, batch_lr_scale=1.4142135623730951,
                        weight_decay=0.059738)

# -----------------------------------------------------------------------------
# Compile the training step. REQUIRED, not an optimization: the CE block in
# forward_backward is written for inductor's fusion -- run eager it
# materializes full (T, d_vocab) fp32 temporaries. fullgraph so any graph
# break errors loudly instead of silently fragmenting fusion (the FA3 raw ops
# have fake impls, so a full trace is achievable).
fb = torch.compile(forward_backward, dynamic=False, fullgraph=True)

# The eval forward is compiled too -- eager it materializes the full
# (T, d_vocab) fp32 logits chain (~13 GB of temporaries per val micro-batch).
# Rebinding the name routes every consumer (the val-loss section and
# forward_eval_packed) through it; it specializes once per shape/targets
# combination: the val path at the training micro-batch shape, the CORE
# logits path at the eval buffer shape.
forward = torch.compile(forward, dynamic=False, fullgraph=True)

# token_bytes: per-token-id byte lengths (0 for special tokens), for the
# vocab-size-independent bits-per-byte validation metric.
with open(os.path.join(DATASET_DIR, "tokenizer/token_bytes.pt"), "rb") as f:
    token_bytes = torch.load(f, map_location=device)

# Enough data for the horizon? The loader is single-epoch and TRUNCATES long
# documents at seq_len, discarding the tails: measured ~11-12% of climbmix's
# raw tokens (doc-length scan, 2026-07-31 session NOTES). 0.85 is that
# discard with margin -- a raw-token floor alone would pass configs that run
# dry ~11% before the horizon.
_shard_tokens = sum((os.path.getsize(f) - 256 * 4) // 2 for f in glob.glob(train_files))
assert _shard_tokens * 0.85 >= (cfg.num_iterations + 1) * cfg.total_batch_size, \
    f"train shards hold {_shard_tokens:,} raw tokens (~{int(_shard_tokens * 0.85):,} usable " \
    f"after seq_len truncation) < {(cfg.num_iterations + 1) * cfg.total_batch_size:,} needed " \
    f"-- raise NUM_TRAIN_SHARDS"

# --- wandb logging init ---
use_dummy_wandb = cfg.wandb_run == "dummy" or not master_process
wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(
    project=cfg.wandb_project, name=cfg.wandb_run,
    config={
        "num_params": num_params,
        "num_flops_per_token": num_flops_per_token,
        "n_layers": cfg.n_layers, "n_q_heads": cfg.n_q_heads, "d_model": cfg.d_model,
        "train_steps": cfg.num_iterations,
        "total_batch_size": cfg.total_batch_size,
        "micro_batch_tokens": cfg.micro_batch_tokens,
        "val_loss_every": cfg.val_loss_every,
        "world_size": world_size,
        "grad_accum_steps": grad_accum_steps,
        "batch_lr_scale": sched.batch_lr_scale,
        "weight_decay": sched.weight_decay,
    },
)
if not use_dummy_wandb:
    wandb.define_metric("step")
    wandb.define_metric("*", step_metric="step")

# -----------------------------------------------------------------------------
# Checkpoint capture (write only -- there is deliberately no load/resume path
# yet). Two files per capture point in logs/{run_id}/:
#   model_stepNNNNNN.pt -- {step, code, weights: {name: tensor}} -- the bf16
#     live weights + fp32 scalars, the payload the final save has always held.
#   optim_stepNNNNNN.pt -- {step, t_step, state: {"name.attr": tensor}} over
#     the five optimizer-state attrs; together with the live weights this is
#     the full fp32 masters and both optimizers' moments.
# World-agnostic: sharded state all-gathers to full size before writing, so a
# capture from an 8-GPU run loads at any world size (at world=1 the gathers
# short-circuit and this is a plain copy-out). Every rank participates in the
# gathers; only master materializes CPU copies and writes -- tensors are saved
# on CPU so the files open anywhere.

state_attrs = ("mantissa", "frst_mntm", "scnd_mntm", "exp_avg", "exp_avg_sq")

# The sharded weights -- their state gathers over dim 0; everything else is
# replicated, already full-size on every rank. Mirrors § Shard Assignment.
# A set, not a tuple: tuple membership falls through identity to elementwise
# tensor ==, while set membership stays on the identity hash.
sharded_weights = {m.W_Q, m.W_K, m.W_V, m.W_O, m.W_in, m.W_out,
                   m.lm_head, m.input_embeds, m.value_embeds}

def gather_full(t):
    """All-gather a shard-size state tensor to full size over dim 0. uint16
    (mantissa) rides as a bf16 bitcast: NCCL has no 16-bit int type, and a
    gather only moves bytes."""
    if world_size == 1:
        return t
    comm = t.view(torch.bfloat16) if t.dtype == torch.uint16 else t
    full = torch.empty(t.shape[0] * world_size, *t.shape[1:], dtype=comm.dtype, device=device)
    dist.all_gather_into_tensor(full, comm)
    return full.view(torch.uint16) if t.dtype == torch.uint16 else full

def write_checkpoint(step):
    state = {}
    for n in m.weight_names:
        p = getattr(m, n)
        for attr in state_attrs:
            if hasattr(p, attr):
                full = gather_full(getattr(p, attr)) if p in sharded_weights else getattr(p, attr)
                if master_process:
                    state[f"{n}.{attr}"] = full.cpu()
    if not master_process:
        return
    os.makedirs(f"logs/{run_id}", exist_ok=True)
    torch.save(dict(step=step, code=code,
                    weights={n: getattr(m, n).cpu() for n in m.weight_names}),
               f"logs/{run_id}/model_step{step:06d}.pt")
    torch.save(dict(step=step, t_step=int(t_step.item()), state=state),
               f"logs/{run_id}/optim_step{step:06d}.pt")

# -----------------------------------------------------------------------------
# Training and validation

train_steps = cfg.num_iterations
train_loader = distributed_data_generator(train_files, cfg.total_batch_size, cfg.seq_len, grad_accum_steps)
inputs, targets, cu_seqlens = next(train_loader)   # kick off the first batch

# Each val pass draws val_tokens through micro-batches shaped exactly like
# training's (so the rotary-cache bound holds), scored with the eager forward.
micro_world_tokens = cfg.total_batch_size // grad_accum_steps   # tokens per micro-batch across ranks
assert cfg.val_tokens % micro_world_tokens == 0
val_steps = cfg.val_tokens // micro_world_tokens

val_bpb = None
min_val_bpb = float("inf")
smooth_train_loss = 0.0
total_training_time = 0.0   # seconds; excludes the first 10 steps (compile lives there)

for step in range(train_steps + 1):
    last_step = (step == train_steps)

    # --------------- VALIDATION SECTION -----------------
    if last_step or (cfg.val_loss_every > 0 and step % cfg.val_loss_every == 0):
        torch.cuda.synchronize()
        val_t0 = time.perf_counter()
        val_loader = distributed_data_generator(val_files, cfg.total_batch_size, cfg.seq_len, grad_accum_steps)
        total_nats = torch.tensor(0.0, dtype=torch.float32, device=device)
        total_bytes = torch.tensor(0, dtype=torch.int64, device=device)
        for _ in range(val_steps):
            v_inputs, v_targets, v_cu_seqlens = next(val_loader)
            loss_flat = forward(v_inputs, v_cu_seqlens, v_targets, loss_reduction='none')
            num_bytes_flat = token_bytes[v_targets]
            total_nats += (loss_flat * (num_bytes_flat > 0)).sum()
            total_bytes += num_bytes_flat.sum()
        del val_loader
        if world_size > 1:
            dist.all_reduce(total_nats, op=dist.ReduceOp.SUM)
            dist.all_reduce(total_bytes, op=dist.ReduceOp.SUM)
        val_bpb = total_nats.item() / (math.log(2) * total_bytes.item())
        min_val_bpb = min(min_val_bpb, val_bpb)
        val_elapsed = time.perf_counter() - val_t0
        print0(f"step:{step}/{train_steps} val_bpb:{val_bpb:.6f} val_time:{val_elapsed:.2f}s", console=True)
        wandb_run.log({"step": step, "val/bpb": val_bpb, "val/eval_seconds": val_elapsed,
                       "total_training_time": total_training_time})

    # --------------- CHECKPOINT CAPTURE -----------------
    # State on entering step `step` = after `step` completed updates. Every
    # rank enters (the gathers are collectives); only master writes.
    if cfg.save_checkpoint and (last_step or step in cfg.save_steps):
        ckpt_t0 = time.perf_counter()
        write_checkpoint(step)
        print0(f"checkpoint captured at step {step} ({time.perf_counter() - ckpt_t0:.1f}s)", console=True)

    if last_step:
        # --------------- CORE EVALUATION -----------------
        if os.path.exists(os.path.join(DATASET_DIR, "core_eval/config.json")):
            core_eval_t0 = time.perf_counter()
            core_out = evaluate_core()
            core_eval_elapsed = time.perf_counter() - core_eval_t0
            print0(f"CORE metric: {core_out['core_metric']:.4f} | total CORE eval time: {core_eval_elapsed:.2f}s", console=True)
            for label, acc in core_out['results'].items():
                print0(f"  {label}: accuracy={acc:.4f} centered={core_out['centered_results'][label]:.4f}", console=True)
            wandb_run.log({
                "step": step,
                "core_metric": core_out["core_metric"],
                **{f"core/{label}/accuracy": acc for label, acc in core_out["results"].items()},
                **{f"core/{label}/centered": c for label, c in core_out["centered_results"].items()},
                "timing/core_eval_seconds": core_eval_elapsed,
            })
        else:
            print0("No core_eval/ in the dataset dir; skipping the CORE metric.", console=True)
        break

    # --------------- TRAINING SECTION -----------------
    torch.cuda.synchronize()
    step_t0 = time.perf_counter()
    for micro in range(grad_accum_steps):
        # loss_scale replaces the loss/grad_accum division of an autograd loop
        loss = fb(inputs, targets, cu_seqlens, loss_scale=grad_scale)
        inputs, targets, cu_seqlens = next(train_loader)  # prefetch while the GPU is busy
    optimizer_step()   # schedules pre-computed; advances t_step on-device
    # Zero every grad buffer: gradients accumulate across the next step's
    # micro-batches, and at world=1 Muon's nesterov lerp just MUTATED grad32
    # (at world>1 it mutates the reduce-scattered shard instead) -- this is
    # correctness, not hygiene. (`for p in m` = every trained weight, in
    # Model.weight_names order.)
    for p in m:
        p.grad32.zero_()
    train_loss = loss.item()   # the step's one host sync point
    torch.cuda.synchronize()
    dt = time.perf_counter() - step_t0

    # logging (CPU only). EMA the loss for readability; time totals exclude the
    # first 10 steps, where compilation dominates.
    ema_beta = 0.9
    smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss
    debiased_smooth_loss = smooth_train_loss / (1 - ema_beta ** (step + 1))
    if step > 10:
        total_training_time += dt
    tok_per_sec = int(cfg.total_batch_size / dt)
    mfu = 100 * num_flops_per_token * cfg.total_batch_size / dt / (gpu_peak_flops * world_size)
    steps_timed = step - 10
    if steps_timed > 0:
        eta_seconds = (train_steps - step - 1) * (total_training_time / steps_timed)
        eta_str = f" | eta: {eta_seconds/60:.1f}m"
    else:
        eta_str = ""
    pct_done = 100 * step / train_steps
    print0(f"step {step:05d}/{train_steps:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} | lrm: {sched.lrm_table[step]:.2f} | dt: {dt*1000:.2f}ms | tok/sec: {tok_per_sec:,} | bf16_mfu: {mfu:.2f} | total time: {total_training_time/60:.2f}m{eta_str}", console=True)
    wandb_run.log({
        "step": step,
        "train/loss": debiased_smooth_loss,
        "train/lrm": float(sched.lrm_table[step]),
        "train/dt": dt,
        "train/tok_per_sec": tok_per_sec,
        "train/mfu": mfu,
        "total_training_time": total_training_time,
    })

    # GC management: the collector's cycle scans cost ~500ms at random steps,
    # so collect the setup garbage once, then freeze survivors and disable.
    if step == 0:
        gc.collect()
        gc.freeze()
        gc.disable()
    elif step % 5000 == 0:
        gc.collect()

print0(f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
       f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB", console=True)
print0(f"total training time: {total_training_time/60:.2f}m", console=True)
if val_bpb is not None:
    print0(f"minimum validation bpb: {min_val_bpb:.6f}", console=True)

wandb_run.finish()
dist.destroy_process_group()