Converter now emits nanochat optimizer shards for an SFT warm-start

README.md

@@ -79,7 +79,8 @@ logs/
   run_full_d24_w8.sh                     launcher, with the config rationale
   baseline_report.md                     throughput + val bpb vs upstream nanochat
 tokenizer/                          the 32k vocab these weights were trained on
-convert_ckpt_to_nanochat.py         DecoderStack capture -> nanochat state_dict
+convert_ckpt_to_nanochat.py         DecoderStack capture -> nanochat state_dict (+optimizer)
+test_convert_ckpt_to_nanochat.py    its test, if you change the converter
 ```
 
 Two capture points: **5568** is the end of the run, **1950** is the last uncooled state —
@@ -147,11 +148,28 @@ design; a load/resume path is future work.
 
 ### Optimizer state
 
-Not converted — which matters for one downstream use and not the other.
-
 **SFT works with the model alone.** nanochat's `chat_sft` builds a fresh optimizer and only
 *optionally* warm-starts it; with no optimizer shard present it prints `starting with fresh
-optimizer (slightly worse)` and carries on. Nothing above is missing for SFT.
+optimizer (slightly worse)` and carries on.
+
+To remove that "slightly worse", add `--world-size N` — the converter then also writes the
+ZeRO-2 shards `optim_005568_rank{0..N-1}.pt` that `DistMuonAdamW` expects, matched to the
+GPU count of the SFT run:
+
+```bash
+python convert_ckpt_to_nanochat.py \
+    --model checkpoints/model_step005568.pt \
+    --optim checkpoints/optim_step005568.pt \
+    --meta  base_checkpoints/d24_decoderstack/meta_005568.json \
+    --out   ~/.cache/nanochat/base_checkpoints/d24_decoderstack \
+    --world-size 8
+```
+
+Note that `torch`'s `load_state_dict` replaces param-group dicts wholesale, so the emitted
+`lr`/`betas`/`weight_decay` become the optimizer's on load — they default to
+`setup_optimizer()`'s own values with `weight_decay=0.0` (the SFT setting, and where
+DecoderStack's cosine-to-zero Muon decay lands), and are CLI-overridable. `chat_sft`
+restores its own LRs immediately after loading regardless.
 
 **Pre-training resume is out of reach**, and it is the dataloader that closes the door, not
 the optimizer: nanochat's resume needs a `dataloader_state_dict` to put its tokenizing
@@ -176,17 +194,21 @@ the two AdamW kernels differ in whether the param carries a mantissa, not in mom
 `W_O`'s NorMuon reduction axis differs — nanochat infers it from the shape and lands on
 `-1` for a square `c_proj`, DecoderStack states `residual_dim = -2` — but at d24 it costs
 nothing: polar express returns a ~orthonormal update, and a square orthonormal matrix has
-~uniform neuron norms along either axis, so there is no variance to reduce. The explicit
-axis only earns its keep when `n_heads * d_head != d_model`.
+~uniform neuron norms along either axis, so there is no variance to reduce. The converter
+mean-fills onto nanochat's axis, and *asserts* rather than converting if the bank is not
+square, where the two axes would carry genuinely different information. The explicit axis
+only earns its keep when `n_heads * d_head != d_model` (and, for `ve_gate`, when
+`n_kv_head != d_ve_gate`; both are equalities at d24).
 
-So an SFT warm-start is transferable in principle; the remaining mapping — param ordering
-and per-rank sharding — is written out at the bottom of `convert_ckpt_to_nanochat.py`.
+**Caveat:** the optimizer path is verified against a real nanochat `GPT` and `MuonAdamW` at
+a toy config — structure, per-rank sharding, shard reassembly, and a real warm-started
+`optimizer.step()` — but has not yet been run against the full d24 capture.
 
 ## Provenance
 
 Trained by [`chrisjmccormick/stacks`](https://github.com/chrisjmccormick/stacks) — the
 converter is committed at
-[`2e199ac`](https://github.com/chrisjmccormick/stacks/commit/2e199acd743c6dcdb4d768daae9b73500c9354a8)
+[`a5e608f`](https://github.com/chrisjmccormick/stacks/commit/a5e608fac7759aefd5037ca1a985d0165a7d285c)
 (`utils/convert_ckpt_to_nanochat.py`). The training script that produced these weights is
 `code/run_full_d24_w8.py` in this repo, which is the run copy of the single-file d24
 trainer with three launcher overrides (`micro_batch_tokens` 32768→65536,

convert_ckpt_to_nanochat.py

@@ -40,22 +40,29 @@
 # and 2.4833 without (bpb 0.7252 vs 0.7254). Pass --optim when you want the
 # exact master anyway -- it is a bit-exact reconstruction, not an approximation.
 #
-# WHAT THIS DOES NOT DO
-# ---------------------
-# It does not convert optimizer state, and nanochat could not resume from it if
-# it did -- see OPTIMIZER-STATE NOTES at the bottom of this file.
+# OPTIMIZER STATE
+# ---------------
+# --world-size N additionally writes optim_NNNNNN_rank{0..N-1}.pt, the ZeRO-2
+# shards nanochat's DistMuonAdamW expects, so chat_sft can warm-start its
+# optimizer instead of printing "starting with fresh optimizer (slightly worse)".
+# It is optional: SFT is correct without it. Pre-training resume is out of reach
+# either way -- see OPTIMIZER-STATE NOTES at the bottom of this file.
 #
 # Usage:
 #   python utils/convert_ckpt_to_nanochat.py \
 #       --model checkpoints/model_step005568.pt \
 #       --optim checkpoints/optim_step005568.pt \
 #       --meta  base_checkpoints/d24_decoderstack/meta_005568.json \
-#       --out   ~/.cache/nanochat/base_checkpoints/d24_decoderstack
+#       --out   ~/.cache/nanochat/base_checkpoints/d24_decoderstack \
+#       --world-size 8            # optional: also emit the optimizer shards
 #
 # Then, in nanochat (branch fa-varlen):
 #   from nanochat.checkpoint_manager import build_model
 #   model, tokenizer, meta = build_model(checkpoint_dir, 5568, device, "eval")
 #
+# Verify a change to this file with:
+#   NANOCHAT_PATH=~/nanochat python utils/test_convert_ckpt_to_nanochat.py
+#
 # The tokenizer is NOT interchangeable with other nanochat d24 releases -- see
 # the model card. DecoderStack trained on the 32k vocab shipped with the
 # ChrisMcCormick/climbmix_32k_8_170 dataset repo; pairing these weights with a
@@ -85,6 +92,42 @@ def fp32_master(live: torch.Tensor, mantissa: torch.Tensor | None) -> torch.Tens
     return bits.view(torch.float32)
 
 
+def _adamw_groups(ve_slots: int):
+    """The AdamW half of setup_optimizer()'s group list, in its exact order.
+
+    Each entry is (lr_key, [(bank_name, bank_slot), ...], betas, eps, weight_decay).
+    bank_slot is None for a whole tensor, an int to index a bank's dim 0. The
+    betas/eps/wd here are constants in BOTH codebases -- nanochat hardcodes them
+    in setup_optimizer, DecoderStack passes the same numbers to build_schedules --
+    so they are not a guess about this run, they are the shared values.
+    """
+    return [
+        ("lm_head",      [("lm_head", None)],        (0.8,  0.96),  1e-10, 0.01),
+        ("embedding",    [("input_embeds", None)],   (0.8,  0.995), 1e-10, 0.001),
+        ("value_embeds", [("value_embeds", j) for j in range(ve_slots)],
+                                                     (0.8,  0.995), 1e-10, 0.01),
+        ("resid",        [("resid_lambdas", None)],  (0.8,  0.95),  1e-10, 0.05),
+        ("x0",           [("x0_lambdas", None)],     (0.96, 0.95),  1e-10, 0.0),
+        ("smear",        [("smear_gate", None), ("smear_lambda", None),
+                          ("backout_lambda", None)], (0.8,  0.95),  1e-10, 0.0),
+    ]
+
+
+def _matrix_params(n_layer: int, ve: list[int]):
+    """setup_optimizer()'s `matrix_params`, in list(transformer.h.parameters()) order.
+
+    Module registration order gives, per block: attn.c_q, c_k, c_v, c_proj,
+    [ve_gate], then mlp.c_fc, mlp.c_proj.
+    """
+    out = []
+    for i in range(n_layer):
+        out += [("W_Q", i), ("W_K", i), ("W_V", i), ("W_O", i)]
+        if i in ve:
+            out.append(("ve_gate", ve.index(i)))   # ve_gate banks by SLOT, not layer
+        out += [("W_in", i), ("W_out", i)]
+    return out
+
+
 def ve_layers(n_layer: int) -> list[int]:
     """Layers carrying a value embedding, in bank-slot order.
 
@@ -149,6 +192,123 @@ def convert(model_data: dict, optim_state: dict | None) -> dict:
     return sd
 
 
+def convert_optimizer(model_data: dict, optim_data: dict, world_size: int, rank: int,
+                      lrs: dict) -> dict:
+    """DecoderStack optimizer capture -> one rank's nanochat optimizer state_dict.
+
+    nanochat's state_dict is keyed by flattened param INDEX over setup_optimizer()'s
+    groups, and DistMuonAdamW shards that state per rank. Our capture all-gathered
+    everything to full size, so this is re-slicing, not reconstruction.
+
+    Returns the dict to torch.save as optim_NNNNNN_rank{rank}.pt. Call once per rank
+    rather than building them all: at d24/world=8 each shard is ~1 GB.
+    """
+    w, st = model_data["weights"], optim_data["state"]
+    t_step = optim_data["t_step"]
+    n_layer = w["W_Q"].shape[0]
+    ve = ve_layers(n_layer)
+    n_embd = w["input_embeds"].shape[1]
+    d_scale = (n_embd / 768) ** -0.5   # setup_optimizer's 1/sqrt(dmodel) AdamW LR scale
+
+    def bank(name, slot, attr):
+        """One param's full-size optimizer state. value_embeds is the odd one out:
+        its AdamW state is shaped over the FLATTENED (slot * vocab) row axis, so it
+        has to be folded back to 3-D before a slot can be indexed."""
+        t = st[f"{name}.{attr}"]
+        if slot is None:
+            return t
+        if name == "value_embeds":
+            return t.view(len(ve), -1, t.shape[-1])[slot]
+        return t[slot]
+
+    adamw_lr = {
+        "lm_head": lrs["unembedding_lr"] * d_scale,
+        "embedding": lrs["embedding_lr"] * d_scale,
+        "value_embeds": lrs["embedding_lr"] * d_scale * 0.5,
+        "resid": lrs["scalar_lr"] * 0.01,
+        "x0": lrs["scalar_lr"],
+        "smear": 0.2,   # hardcoded in setup_optimizer, not scaled
+    }
+
+    # --- Build the group plan exactly as setup_optimizer() would: AdamW groups in
+    #     a fixed order, then Muon groups keyed by `sorted({p.shape})`. ---
+    plan = []   # (kind, [(name, slot), ...], hyperparams dict)
+    for lr_key, params, betas, eps, wd in _adamw_groups(len(ve)):
+        plan.append(("adamw", params, dict(kind="adamw", lr=adamw_lr[lr_key],
+                                           betas=list(betas), eps=eps, weight_decay=wd)))
+    matrix = _matrix_params(n_layer, ve)
+    shape_of = lambda p: tuple(w[p[0]].shape[1:])
+    for shape in sorted({shape_of(p) for p in matrix}):
+        plan.append(("muon", [p for p in matrix if shape_of(p) == shape],
+                     dict(kind="muon", lr=lrs["matrix_lr"], momentum=0.95, ns_steps=5,
+                          beta2=0.9, weight_decay=lrs["weight_decay"])))
+
+    # Param indices are assigned by walking the groups in order.
+    index, idx = {}, 0
+    for _, params, _ in plan:
+        for p in params:
+            index[p] = idx
+            idx += 1
+    assert idx == 7 + len(ve) + len(matrix), f"param count {idx} does not add up"
+
+    state, groups = {}, []
+    for kind, params, hp in plan:
+        groups.append({**hp, "initial_lr": hp["lr"],
+                       "params": [index[p] for p in params]})
+        if kind == "adamw":
+            for p in params:
+                exp_avg = bank(*p, "exp_avg")
+                # ZeRO-2: params with >= 1024 elements are row-sharded over dim 0 by
+                # rank; smaller ones are replicated (nanochat batches those into an
+                # all_reduce instead of a reduce_scatter).
+                if exp_avg.numel() >= 1024:
+                    assert exp_avg.shape[0] % world_size == 0, \
+                        f"{p}: dim 0 ({exp_avg.shape[0]}) must divide world_size {world_size}"
+                    rows = exp_avg.shape[0] // world_size
+                    cut = lambda t, n=rows: t[rank * n:(rank + 1) * n].clone()
+                else:
+                    cut = lambda t: t.clone()
+                state[index[p]] = {
+                    "step": t_step,
+                    "exp_avg": cut(exp_avg),
+                    "exp_avg_sq": cut(bank(*p, "exp_avg_sq")),
+                }
+        else:
+            # Muon state is one stacked buffer per GROUP, held under the first
+            # param's entry, chunked across ranks and zero-padded when the group
+            # does not divide evenly.
+            shape = shape_of(params[0])
+            chunk = -(-len(params) // world_size)
+            start = rank * chunk
+            owned = min(chunk, max(0, len(params) - start))
+            mom = torch.zeros(chunk, *shape, dtype=torch.float32)
+            # nanochat factors the second moment along whichever axis its shape
+            # heuristic calls the neuron axis; ours is set explicitly per bank.
+            nc_shape = (shape[-2], 1) if shape[-2] >= shape[-1] else (1, shape[-1])
+            snd = torch.zeros(chunk, *nc_shape, dtype=torch.float32)
+            for k in range(owned):
+                p = params[start + k]
+                mom[k] = bank(*p, "frst_mntm")
+                ours = bank(*p, "scnd_mntm")
+                if tuple(ours.shape) == nc_shape:
+                    snd[k] = ours
+                else:
+                    # Only W_O lands here, and only because nanochat's shape
+                    # heuristic picks the other axis on a SQUARE c_proj. That is
+                    # benign: polar express returns a ~orthonormal update, whose
+                    # neuron norms are ~uniform along either axis, so their mean is
+                    # the right common value. On a non-square bank the two axes
+                    # would carry genuinely different information -- refuse.
+                    assert shape[-2] == shape[-1], (
+                        f"{p}: second-moment axis differs on a non-square bank "
+                        f"{shape} (ours {tuple(ours.shape)}, nanochat {nc_shape}); "
+                        "no faithful conversion exists")
+                    snd[k] = ours.mean()
+            state[index[params[0]]] = {"momentum_buffer": mom,
+                                       "second_momentum_buffer": snd}
+    return {"state": state, "param_groups": groups}
+
+
 def main():
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--model", required=True, help="DecoderStack model_stepNNNNNN.pt")
@@ -159,14 +319,31 @@ def main():
     p.add_argument("--meta", default=None, help="meta_NNNNNN.json to copy alongside the model")
     p.add_argument("--dump-code", action="store_true",
                    help="also write the training script embedded in the capture's `code` field")
+    p.add_argument("--world-size", type=int, default=0, metavar="N",
+                   help="also write optim_NNNNNN_rank{0..N-1}.pt for an N-GPU run "
+                        "(requires --optim). Omit to convert weights only.")
+    # Group hyperparameters for the emitted optimizer. torch's load_state_dict
+    # REPLACES param_group dicts with the saved ones, so whatever goes here becomes
+    # the optimizer's policy on load. Defaults are setup_optimizer()'s own, with
+    # weight_decay=0.0 -- both the SFT setting and where DecoderStack's cosine-to-
+    # zero Muon decay actually lands (4.8e-9 at step 5568). nanochat's chat_sft
+    # restores its own lr right after loading and schedules momentum per step, so
+    # in practice only betas/eps/weight_decay/ns_steps come from here.
+    p.add_argument("--unembedding-lr", type=float, default=0.004)
+    p.add_argument("--embedding-lr", type=float, default=0.2)
+    p.add_argument("--matrix-lr", type=float, default=0.02)
+    p.add_argument("--scalar-lr", type=float, default=0.5)
+    p.add_argument("--weight-decay", type=float, default=0.0)
     args = p.parse_args()
+    if args.world_size and not args.optim:
+        p.error("--world-size needs --optim (the optimizer state lives in that file)")
 
     # mmap so a 2.8 GB model / 11 GB optimizer file is paged, not slurped.
     model_data = torch.load(args.model, map_location="cpu", mmap=True, weights_only=True)
     step = model_data["step"]
     print(f"loaded {args.model}: step {step}, {len(model_data['weights'])} weights")
 
-    optim_state = None
+    optim_data = optim_state = None
     if args.optim:
         optim_data = torch.load(args.optim, map_location="cpu", mmap=True, weights_only=True)
         assert optim_data["step"] == step, f"optim step {optim_data['step']} != model step {step}"
@@ -200,6 +377,25 @@ def main():
         print(f"NOTE: nanochat also needs meta_{step:06d}.json in {args.out} "
               "(model_config lives there, not in the .pt)")
 
+    if args.world_size:
+        lrs = dict(unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr,
+                   matrix_lr=args.matrix_lr, scalar_lr=args.scalar_lr,
+                   weight_decay=args.weight_decay)
+        # One rank at a time -- holding all of them would cost the whole optimizer.
+        for r in range(args.world_size):
+            shard = convert_optimizer(model_data, optim_data, args.world_size, r, lrs)
+            path = os.path.join(args.out, f"optim_{step:06d}_rank{r:d}.pt")
+            torch.save(shard, path)
+            print(f"wrote {path} ({os.path.getsize(path):,} bytes)")
+            if r == 0:
+                g = shard["param_groups"]
+                print(f"  {len(g)} groups "
+                      f"({sum(1 for x in g if x['kind'] == 'adamw')} adamw / "
+                      f"{sum(1 for x in g if x['kind'] == 'muon')} muon), "
+                      f"{sum(len(x['params']) for x in g)} params, "
+                      f"step {optim_data['t_step']}")
+            del shard
+
     if args.dump_code:
         code_path = os.path.join(args.out, f"code_{step:06d}.py")
         with open(code_path, "w", encoding="utf-8", newline="\n") as f:
@@ -214,22 +410,29 @@ if __name__ == "__main__":
 # -----------------------------------------------------------------------------
 # OPTIMIZER-STATE NOTES
 # -----------------------------------------------------------------------------
-# Optimizer state is not converted -- but check which downstream use you have
-# before assuming that blocks you, because the two land very differently.
-#
 # SFT CONTINUATION WORKS WITH THE MODEL ALONE. nanochat's chat_sft builds a fresh
 # optimizer via model.setup_optimizer() and only optionally warm-starts it from
 # load_optimizer_state(); when the shard is absent it prints "optimizer
 # checkpoint not found, starting with fresh optimizer (slightly worse)" and
-# carries on. The conversion above is all SFT needs.
+# carries on. --world-size exists to remove that "slightly worse", not to unlock
+# anything.
 #
-# PRE-TRAINING RESUME is the one that is genuinely out of reach, and it is the
-# dataloader that closes the door, not the optimizer: nanochat's resume needs
+# PRE-TRAINING RESUME is genuinely out of reach, and it is the dataloader that
+# closes the door, not the optimizer: nanochat's resume needs
 # meta_data["dataloader_state_dict"] to put its tokenizing loader back in the
 # stream, and DecoderStack reads pre-tokenized binary shards through a loader
 # with no equivalent state to hand over. The data order could not be continued no
 # matter what the optimizer held.
 #
+# HYPERPARAMETERS ARE POLICY, NOT STATE. torch's Optimizer.load_state_dict
+# REPLACES each param_group dict with the saved one, keeping only 'params' -- so
+# whatever this script writes becomes the optimizer's lr/betas/wd on load. That
+# is why chat_sft saves and restores its own LRs around the call. The emitted
+# groups use setup_optimizer()'s defaults (overridable on the command line), with
+# weight_decay=0.0: both the SFT setting and where DecoderStack's cosine-to-zero
+# Muon decay actually lands (4.8e-9 at step 5568). betas/eps/adamw-wd are not a
+# guess -- they are identical constants in both codebases.
+#
 # Every buffer we keep has a nanochat counterpart, and the precisions line up on
 # everything except the two embedding tables:
 #
@@ -266,12 +469,20 @@ if __name__ == "__main__":
 # d_model, W_O looks like an MLP projection and the shape heuristic happens to
 # agree; below it the heuristic picks the wrong axis, and since W_O stores its
 # heads transposed relative to QKV, the right answer is not one a shape alone
-# can give. Every other bank agrees either way (W_in -1, W_out -2, QKV -1,
-# ve_gate -1).
+# can give. Every other bank agrees at d24 (W_in -1, W_out -2, QKV -1).
+#
+# ve_gate is the second place the axes can diverge, and it is worth knowing about
+# because it is NOT square in general. DecoderStack banks it (num_ves, n_kv_heads,
+# d_ve_gate) with residual_dim = -1, so the neurons are the n_kv_heads rows;
+# nanochat's Linear(ve_gate_channels=12, n_kv_head) hits the same heuristic and
+# agrees only when n_kv_head >= 12. At d24 n_kv_head == 12 == d_ve_gate, so the
+# bank is square and the two land together. A model with fewer than 12 KV heads
+# would disagree for real -- convert_optimizer() asserts rather than papering
+# over it, since outside the square case the two axes carry different
+# information.
 #
-# So for an SFT warm-start the state is transferable in principle. The remaining
-# mapping is mechanical but fiddly, because nanochat's state_dict is keyed by
-# flattened param INDEX and is sharded per rank:
+# --world-size implements the mapping below. It is mechanical but fiddly, because
+# nanochat's state_dict is keyed by flattened param INDEX and is sharded per rank:
 #   - Param order is setup_optimizer()'s group order: lm_head, wte,
 #     value_embeds.*, resid_lambdas, x0_lambdas, [smear_gate.weight,
 #     smear_lambda, backout_lambda], then the Muon groups in `sorted({shapes})`

test_convert_ckpt_to_nanochat.py

@@ -0,0 +1,248 @@
+# Self-contained test for convert_ckpt_to_nanochat.py. No pytest, no fixtures:
+#
+#   NANOCHAT_PATH=~/nanochat python utils/test_convert_ckpt_to_nanochat.py
+#
+# Builds a synthetic DecoderStack capture (weights + optimizer state, with real
+# mantissas) at a toy config, converts it, and checks the result against a REAL
+# nanochat GPT and a REAL MuonAdamW -- so the assertions are about nanochat's
+# actual behaviour, not a second copy of my assumptions about it.
+#
+# The toy config deliberately mirrors d24's SHAPE RELATIONSHIPS rather than just
+# being small: n_head * head_dim == n_embd (W_O square) and n_kv_head == 12 ==
+# d_ve_gate (ve_gate square). Those are exactly the conditions under which the two
+# codebases' NorMuon second-moment axes agree; a config that breaks either is
+# meant to be REFUSED, and the last check covers that.
+#
+# Requires a nanochat checkout on the fa-varlen branch (or any branch whose GPT
+# uses the modular transformer.h.N.* layout).
+import math
+import os
+import sys
+
+import torch
+
+_NC = os.environ.get("NANOCHAT_PATH")
+if not _NC or not os.path.isdir(os.path.join(os.path.expanduser(_NC), "nanochat")):
+    sys.exit("set NANOCHAT_PATH to a nanochat checkout, e.g. NANOCHAT_PATH=~/nanochat")
+sys.path.insert(0, os.path.expanduser(_NC))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# COMPUTE_DTYPE auto-detects to fp32 with no CUDA; force the bf16 a GPU box would
+# pick, so the embedding dtypes we compare against are the ones a real load wants.
+os.environ["NANOCHAT_DTYPE"] = "bfloat16"
+# nanochat's fused kernels are @torch.compile'd and inductor's CPU backend wants a
+# host compiler. Run them eager -- the bodies are plain PyTorch, so the shapes and
+# the math are still exercised, just unfused.
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+from nanochat.gpt import GPT, GPTConfig                       # noqa: E402
+from convert_ckpt_to_nanochat import (convert, convert_optimizer,  # noqa: E402
+                                      ve_layers, _matrix_params)
+
+NL, DM, NH, HD, V = 8, 192, 12, 16, 384
+KV, MLP = NH * HD, 4 * DM
+VE, D_VE_GATE, D_SMR = ve_layers(NL), 12, 24
+NVE = len(VE)
+LRS = dict(unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, scalar_lr=0.5,
+           weight_decay=0.0)
+
+upper = lambda x: (x.contiguous().view(torch.int32) >> 16).to(torch.int16).view(torch.bfloat16)
+lower = lambda x: (x.contiguous().view(torch.int32)).to(torch.int16).view(torch.uint16)
+
+torch.manual_seed(0)
+
+# ---------------------------------------------------------------- synthetic capture
+BANKS = {
+    "input_embeds": (V, DM), "value_embeds": (NVE, V, KV), "lm_head": (V, DM),
+    "W_Q": (NL, KV, DM), "W_K": (NL, KV, DM), "W_V": (NL, KV, DM), "W_O": (NL, DM, KV),
+    "W_in": (NL, MLP, DM), "W_out": (NL, DM, MLP), "ve_gate": (NVE, NH, D_VE_GATE),
+}
+fp32_src, weights, state = {}, {}, {}
+for name, shape in BANKS.items():                 # bf16 live + uint16 mantissa
+    t = torch.randn(*shape)
+    fp32_src[name], weights[name], state[f"{name}.mantissa"] = t, upper(t), lower(t)
+for name, shape in [("resid_lambdas", (NL,)), ("x0_lambdas", (NL,)),
+                    ("smear_gate", (1, D_SMR)), ("smear_lambda", (1,)),
+                    ("backout_lambda", (1,))]:    # fp32 live, no mantissa
+    weights[name] = torch.randn(*shape)
+
+# AdamW moments: param-shaped, except value_embeds, which the real capture holds
+# over the FLATTENED (slot * vocab) row axis. Every element gets a distinct value
+# so a misrouted or misaligned shard is loud, but the SCALE stays realistic --
+# with a second moment of ~1e7 the resulting update lands below fp32 epsilon and
+# the "every param moved" check below fails on the (1,) scalars for reasons that
+# have nothing to do with the conversion.
+for t_id, (name, shape) in enumerate([("input_embeds", (V, DM)), ("lm_head", (V, DM)),
+                                      ("value_embeds", (NVE * V, KV)),
+                                      ("resid_lambdas", (NL,)), ("x0_lambdas", (NL,)),
+                                      ("smear_gate", (1, D_SMR)), ("smear_lambda", (1,)),
+                                      ("backout_lambda", (1,))]):
+    for a_id, attr in enumerate(("exp_avg", "exp_avg_sq")):
+        n = math.prod(shape)
+        base = (t_id * 2 + a_id) * n
+        state[f"{name}.{attr}"] = (
+            (torch.arange(n, dtype=torch.float32) + base) * 1e-7).reshape(*shape)
+# Muon moments. The factored second moment follows DecoderStack's EXPLICIT
+# residual_dim: -1 for QKV / W_in / ve_gate, -2 for W_O / W_out.
+SND = {"W_Q": (NL, KV, 1), "W_K": (NL, KV, 1), "W_V": (NL, KV, 1), "W_O": (NL, 1, KV),
+       "W_in": (NL, MLP, 1), "W_out": (NL, 1, MLP), "ve_gate": (NVE, NH, 1)}
+for name, shape in SND.items():
+    b = BANKS[name]
+    state[f"{name}.frst_mntm"] = torch.stack(
+        [torch.full(b[1:], float(i + 1)) for i in range(b[0])])
+    state[f"{name}.scnd_mntm"] = torch.stack(
+        [torch.full(shape[1:], float(i + 1) * 0.5) for i in range(shape[0])])
+
+model_data = {"step": 42, "code": "# toy\n", "weights": weights}
+optim_data = {"step": 42, "t_step": 42, "state": state}
+
+cfg = GPTConfig(sequence_len=256, vocab_size=V, n_layer=NL, n_head=NH, n_kv_head=NH,
+                n_embd=DM, window_pattern="SSSL")
+new_opt = lambda m: m.setup_optimizer(unembedding_lr=LRS["unembedding_lr"],
+                                      embedding_lr=LRS["embedding_lr"],
+                                      matrix_lr=LRS["matrix_lr"],
+                                      weight_decay=LRS["weight_decay"],
+                                      scalar_lr=LRS["scalar_lr"])
+
+# ---------------------------------------------------------------- 1. weights
+with torch.device("meta"):
+    ref = GPT(cfg)
+ref_sd = ref.state_dict()
+for tag, mant in [("with mantissas", state), ("model only", None)]:
+    sd = convert(model_data, mant)
+    assert set(sd) == set(ref_sd), (f"key mismatch: missing "
+                                    f"{sorted(set(ref_sd) - set(sd))[:4]}, extra "
+                                    f"{sorted(set(sd) - set(ref_sd))[:4]}")
+    with torch.device("meta"):
+        model = GPT(cfg)
+    model.to_empty(device="cpu")
+    model.init_weights()
+    want_dtype = {k: v.dtype for k, v in model.state_dict().items()}
+    model.load_state_dict(sd, strict=True, assign=True)   # strict: shapes + names
+    bad = {k: (sd[k].dtype, want_dtype[k]) for k in sd if sd[k].dtype != want_dtype[k]}
+    assert not bad, f"dtype mismatch vs a fresh model: {bad}"
+    for key, src, idx in [("lm_head.weight", "lm_head", None),
+                          ("transformer.h.2.attn.c_q.weight", "W_Q", 2),
+                          ("transformer.h.2.attn.c_proj.weight", "W_O", 2),
+                          ("transformer.h.3.mlp.c_fc.weight", "W_in", 3),
+                          (f"transformer.h.{VE[1]}.attn.ve_gate.weight", "ve_gate", 1)]:
+        full = fp32_src[src] if idx is None else fp32_src[src][idx]
+        want = full if mant else (upper(fp32_src[src]) if idx is None
+                                  else upper(fp32_src[src])[idx]).float()
+        assert torch.equal(sd[key], want), f"{tag}: {key} value mismatch"
+    for key, src, idx in [("transformer.wte.weight", "input_embeds", None),
+                          (f"value_embeds.{VE[0]}.weight", "value_embeds", 0)]:
+        want = weights[src] if idx is None else weights[src][idx]
+        assert torch.equal(sd[key], want), f"{tag}: {key} should be the bf16 live weight"
+    print(f"[OK] weights, {tag}: {len(sd)} tensors, strict load, values verified")
+
+ve_keys = sorted(k for k in ref_sd if k.endswith("ve_gate.weight"))
+assert ve_keys == [f"transformer.h.{i}.attn.ve_gate.weight" for i in VE], ve_keys
+print(f"[OK] ve slot->layer map: slots 0..{NVE - 1} -> layers {VE}")
+
+# ---------------------------------------------------------------- 2. optimizer
+ref_opt = new_opt(model)
+param_key = {id(p): k for k, p in model.named_parameters()}
+flat = [p for g in ref_opt.param_groups for p in g["params"]]
+index_of_key = {param_key[id(p)]: i for i, p in enumerate(flat)}
+print(f"reference optimizer: {len(ref_opt.param_groups)} groups, {len(flat)} params")
+
+ROLE = {"W_Q": "attn.c_q", "W_K": "attn.c_k", "W_V": "attn.c_v", "W_O": "attn.c_proj",
+        "W_in": "mlp.c_fc", "W_out": "mlp.c_proj"}
+def key_of(entry):
+    name, slot = entry
+    return {"lm_head": "lm_head.weight", "input_embeds": "transformer.wte.weight",
+            "smear_gate": "smear_gate.weight"}.get(name) \
+        or (f"value_embeds.{VE[slot]}.weight" if name == "value_embeds"
+            else f"transformer.h.{VE[slot]}.attn.ve_gate.weight" if name == "ve_gate"
+            else name if slot is None
+            else f"transformer.h.{slot}.{ROLE[name]}.weight")
+
+ADAMW = ([("lm_head", None), ("input_embeds", None)]
+         + [("value_embeds", j) for j in range(NVE)]
+         + [("resid_lambdas", None), ("x0_lambdas", None), ("smear_gate", None),
+            ("smear_lambda", None), ("backout_lambda", None)])
+matrix = _matrix_params(NL, VE)
+shape_of = lambda e: tuple(weights[e[0]].shape[1:])
+
+# world_size 3 is deliberate: it leaves the Muon groups RAGGED (the ve_gate
+# group hands rank 2 nothing, the attention group one padding slot), which is
+# exactly what d24's 12-param ve_gate group does across 8 ranks.
+for W in (1, 2, 3, 4):
+    shards = [convert_optimizer(model_data, optim_data, W, r, LRS) for r in range(W)]
+
+    for sd_r in shards:                       # structure: a real optimizer accepts it
+        opt = new_opt(model)
+        opt.load_state_dict(sd_r)             # raises on group/param-count mismatch
+        for g in opt.param_groups:
+            if g["kind"] != "muon":
+                continue
+            assert "momentum_buffer" in opt.state[g["params"][0]], "muon state off params[0]"
+            assert not any(opt.state.get(p) for p in g["params"][1:]), "muon state leaked"
+
+    for entry in ADAMW:                       # AdamW: shards must reassemble exactly
+        pidx = index_of_key[key_of(entry)]
+        full = state[f"{entry[0]}.exp_avg"]
+        if entry[1] is not None:
+            full = full.view(NVE, -1, full.shape[-1])[entry[1]]
+        pieces = [s["state"][pidx]["exp_avg"] for s in shards]
+        got = pieces[0] if pieces[0].shape == full.shape else torch.cat(pieces, 0)
+        assert torch.equal(got, full), f"W={W} {key_of(entry)}: adamw state misassembled"
+        assert all(s["state"][pidx]["step"] == 42 for s in shards), "step not carried"
+
+    for shape in sorted({shape_of(e) for e in matrix}):   # Muon: chunked by group
+        members = [e for e in matrix if shape_of(e) == shape]
+        first = index_of_key[key_of(members[0])]
+        chunk = -(-len(members) // W)
+        for n, e in enumerate(members):
+            r, k = n // chunk, n % chunk
+            assert torch.equal(shards[r]["state"][first]["momentum_buffer"][k],
+                               state[f"{e[0]}.frst_mntm"][e[1]]), f"W={W} {e}: momentum"
+            got = shards[r]["state"][first]["second_momentum_buffer"][k]
+            ours = state[f"{e[0]}.scnd_mntm"][e[1]]
+            if tuple(got.shape) == tuple(ours.shape):
+                assert torch.equal(got, ours), f"W={W} {e}: second moment"
+            else:                              # the square-c_proj mean fallback
+                assert e[0] == "W_O" and shape[-2] == shape[-1], f"unexpected fallback {e}"
+                assert torch.allclose(got, ours.mean().expand_as(got)), f"W={W} {e}: fallback"
+        for r in range(W):                     # padding slots stay zero
+            for k in range(chunk):
+                if r * chunk + k >= len(members):
+                    assert not shards[r]["state"][first]["momentum_buffer"][k].any(), "padding"
+    print(f"[OK] optimizer, world_size={W}: structure + routing + reassembly verified")
+
+# ---------------------------------------------------------------- 3. it actually steps
+# At world_size=1 the non-distributed optimizer's buffers have exactly the shapes
+# DistMuonAdamW's rank-0 shard carries, so this runs the real update kernels over
+# our tensors rather than only checking bookkeeping.
+opt = new_opt(model)
+opt.load_state_dict(convert_optimizer(model_data, optim_data, 1, 0, LRS))
+for p in model.parameters():
+    p.grad = torch.randn_like(p) * 1e-3
+before = {k: v.detach().clone() for k, v in model.named_parameters()}
+mom_before = {i: opt.state[g["params"][0]]["momentum_buffer"].clone()
+              for i, g in enumerate(opt.param_groups) if g["kind"] == "muon"}
+opt.step()
+stuck = [k for k, v in model.named_parameters() if torch.equal(v, before[k])]
+assert not stuck, f"params did not move: {stuck[:5]}"
+for i, m0 in mom_before.items():
+    assert not torch.equal(m0, opt.state[opt.param_groups[i]["params"][0]]["momentum_buffer"]), \
+        f"muon group {i}: momentum buffer did not advance"
+assert all(opt.state[g["params"][0]]["step"] == 43
+           for g in opt.param_groups if g["kind"] == "adamw"), "step did not advance from 42"
+assert all(torch.isfinite(v).all() for _, v in model.named_parameters()), "non-finite param"
+print(f"[OK] warm-started MuonAdamW.step(): all {len(before)} params updated, "
+      "buffers advanced, step 42 -> 43")
+
+# ---------------------------------------------------------------- 4. the guard bites
+# Outside the square case the two second-moment axes carry different information,
+# and there is no faithful conversion. Refuse rather than quietly mean-fill.
+bad_state = dict(state)
+bad_state["W_in.scnd_mntm"] = torch.zeros(NL, 1, DM)   # wrong axis, NON-square bank
+try:
+    convert_optimizer(model_data, {"step": 42, "t_step": 42, "state": bad_state}, 1, 0, LRS)
+    sys.exit("FAIL: a non-square second-moment axis mismatch was accepted")
+except AssertionError as e:
+    assert "non-square" in str(e), e
+print("[OK] non-square second-moment axis mismatch refused")
+print("PASS")