Add model card, meta.json, logs, tokenizer, training source, and nanochat converter

README.md

@@ -0,0 +1,180 @@
+---
+license: mit
+language:
+- en
+tags:
+- decoderstack
+- nanochat
+- pretraining
+---
+
+# DecoderStack d24
+
+A 1.38B-parameter d24 model pre-trained from scratch with
+[**DecoderStack**](https://github.com/chrisjmccormick/stacks) — a single-file training
+stack with a handwritten forward/backward and a written-out optimizer: no autograd, no
+`torch.optim`, no `nn.Module`. bf16, varlen FlashAttention, 8xH100, 1h51m.
+
+The architecture is the nanochat d24 architecture, so these weights load into
+[nanochat](https://github.com/karpathy/nanochat) once converted — see
+[Loading into nanochat](#loading-into-nanochat).
+
+## Results
+
+Full horizon: 5,568 steps x 2^20 tokens = 5.84B tokens (data:param ratio 8).
+
+```
+step 05567/05568 (99.98%) | loss: 2.388180 | lrm: 0.05 | dt: 1198.08ms | tok/sec: 875,210 | bf16_mfu: 53.76 | total time: 110.74m
+step:5568/5568 val_bpb:0.719042
+CORE metric: 0.2517 | total CORE eval time: 80.46s
+peak memory allocated: 71569 MiB reserved: 72736 MiB
+total training time: 110.74m
+minimum validation bpb: 0.719042
+```
+
+| | value |
+|---|---|
+| parameters | 1,384,122,122 |
+| min val bpb | 0.719042 |
+| CORE | 0.2517 |
+| train time | 110.74 min on 8xH100 80GB |
+| throughput | 1,212 ms/step, 864,856 tok/sec (steps 20–99) |
+| precision | bf16 (no FP8) |
+
+val bpb curve: 0:3.175835 / 250:0.997472 / 500:0.904081 / 1000:0.846058 / 2000:0.809199 /
+3000:0.773139 / 4000:0.747155 / 5000:0.726908 / 5568:0.719042
+
+**Against upstream nanochat**, run on the same box on the same rebuilt validation shard,
+DecoderStack tracks 0.2–1.4% behind through step 1000 (where the upstream baseline was
+stopped) while running ~6% faster per step (1,212 vs 1,291 ms/step). Full head-to-head in
+[`logs/baseline_report.md`](logs/baseline_report.md).
+
+`ChrisMcCormick/nanochat-varlen-d24-2026-03-22` is the closest sibling release
+(0.724772 val bpb, 0.2614 CORE, 97.4 min) but it is **not** a controlled comparison: it
+trained with FP8 and a different 32k tokenizer.
+
+## ⚠️ Tokenizer
+
+These weights are tied to the 32k vocabulary in
+[`ChrisMcCormick/climbmix_32k_8_170`](https://huggingface.co/datasets/ChrisMcCormick/climbmix_32k_8_170),
+mirrored here under [`tokenizer/`](tokenizer). It is a **different vocabulary** from the
+one in `ChrisMcCormick/nanochat-varlen-d24-2026-03-22` — 31,474 of 32,759 merge ids
+differ. Both are 32,768 tokens and both load without complaint, so nothing will warn you.
+Pairing these weights with the wrong tokenizer produces garbage, not slightly-worse text.
+
+## Contents
+
+```
+checkpoints/                        native DecoderStack captures (the source of record)
+  model_step001950.pt   2.8 GB      {step, code, weights} -- bf16 live weights + fp32 scalars
+  optim_step001950.pt  11.1 GB      {step, t_step, state} -- mantissas + Muon/AdamW moments
+  model_step005568.pt   2.8 GB
+  optim_step005568.pt  11.1 GB
+base_checkpoints/d24_decoderstack/  nanochat-format target dir
+  meta_001950.json                  model_config + run metadata
+  meta_005568.json                  (the converter writes model_NNNNNN.pt in here)
+code/run_full_d24_w8.py             the exact training script this run executed
+logs/
+  d24_decoderstack_20260801_023219.log   full training log
+  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
+```
+
+Two capture points: **5568** is the end of the run, **1950** is the last uncooled state —
+the LR/momentum hold ends after update 1949 (`N - round(0.65*N)`), so 1950 is the one to
+resume from if you want to train the horizon longer.
+
+DecoderStack writes `code` (the full training script, `open(sys.argv[0]).read()`) into
+every model capture, so each `.pt` carries its own exact source. `code/run_full_d24_w8.py`
+is that string extracted from `model_step005568.pt`.
+
+## Loading into nanochat
+
+`convert_ckpt_to_nanochat.py` unbanks the per-layer matrix banks, renames tensors to
+nanochat's module paths, maps value-embedding bank slots back to their layers, and
+rebuilds the fp32 masters from the bf16 live weights plus the uint16 mantissa held in the
+optimizer file.
+
+```bash
+pip install huggingface_hub torch
+python - <<'PY'
+from huggingface_hub import hf_hub_download
+R = "ChrisMcCormick/decoderstack-d24"
+for f in ["checkpoints/model_step005568.pt",          # 2.8 GB
+          "checkpoints/optim_step005568.pt",          # 11.1 GB -- optional, see below
+          "base_checkpoints/d24_decoderstack/meta_005568.json",
+          "convert_ckpt_to_nanochat.py",
+          "tokenizer/tokenizer.pkl", "tokenizer/token_bytes.pt"]:
+    hf_hub_download(R, f, local_dir="decoderstack-d24")
+PY
+
+cd decoderstack-d24
+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
+mkdir -p ~/.cache/nanochat/tokenizer && cp tokenizer/* ~/.cache/nanochat/tokenizer/
+```
+
+(`~/.cache/nanochat` is nanochat's default base dir; override with `NANOCHAT_BASE_DIR`.)
+
+`--optim` is optional. It only supplies the mantissas — the lower 16 bits of the fp32
+masters. Without it the fp32 parameters carry bf16 precision, which costs about 0.0005
+nats/token (2.4833 vs 2.4828 on a held-out paragraph); skip the 11 GB download unless you
+want the exact master.
+
+Then, on the [`fa-varlen`](https://github.com/chrisjmccormick/nanochat/tree/fa-varlen)
+branch (varlen, matching how these weights were trained):
+
+```python
+import os, torch
+from nanochat.checkpoint_manager import build_model
+ckpt_dir = os.path.expanduser("~/.cache/nanochat/base_checkpoints/d24_decoderstack")
+model, tokenizer, meta = build_model(ckpt_dir, 5568, torch.device("cuda"), "eval")
+```
+
+The conversion is verified: 175 tensors, shapes and dtypes identical to a freshly built
+nanochat `GPT` at this config, loads with `strict=True`, and scores 2.4828 nats/token
+(0.725 bpb) on held-out English — in line with the run's own 0.719 val bpb — with
+coherent greedy continuations.
+
+**DecoderStack itself cannot load these back yet.** The capture path is write-only by
+design; a load/resume path is future work.
+
+### Optimizer state
+
+Not converted, and nanochat could not resume from it if it were. Every buffer does have a
+counterpart, and the precisions line up on everything except the embedding tables:
+
+| DecoderStack | nanochat | precision |
+|---|---|---|
+| `.frst_mntm` | Muon `momentum_buffer` | fp32 both |
+| `.scnd_mntm` | Muon `second_momentum_buffer` | fp32 both |
+| `.exp_avg` / `.exp_avg_sq` | AdamW `exp_avg` / `exp_avg_sq` | fp32 both, **except** `wte` and `value_embeds`, which nanochat allocates `zeros_like(p)` on a bf16 param — so bf16 there, fp32 here |
+| `.mantissa` (uint16) | — | nanochat's fp32 param *is* the master; consumed by the converter to rebuild it |
+
+Two things block a faithful resume regardless: `W_O`'s NorMuon second moment reduces along
+a different axis in each codebase (nanochat infers it from the shape and lands on `-1` for
+a square `c_proj`; DecoderStack sets `residual_dim = -2` deliberately, because W_O's heads
+are stored transposed relative to QKV), and nanochat's resume needs a
+`dataloader_state_dict` that DecoderStack's pre-tokenized binary loader has no equivalent
+of. The remaining mapping — param ordering and per-rank sharding — is written out at the
+bottom of `convert_ckpt_to_nanochat.py`.
+
+## Provenance
+
+Trained by [`chrisjmccormick/stacks`](https://github.com/chrisjmccormick/stacks) — the
+converter is committed at
+[`e664830`](https://github.com/chrisjmccormick/stacks/commit/e6648306442e0479433b287ede02d271997365c0)
+(`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,
+`max_num_docs` 96→192, `NUM_TRAIN_SHARDS` 20→80); the rationale for each is in
+`logs/run_full_d24_w8.sh`.
+
+Training data: [`ChrisMcCormick/climbmix_32k_8_170`](https://huggingface.co/datasets/ChrisMcCormick/climbmix_32k_8_170)
+(ClimbMix, pre-tokenized to 32k binary shards). Validation is a rebuild of shard 06542 to
+match nanochat's pinned val split.

base_checkpoints/d24_decoderstack/meta_001950.json

@@ -0,0 +1,87 @@
+{
+  "step": 1950,
+  "val_bpb": null,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 32768,
+    "n_layer": 24,
+    "n_head": 12,
+    "n_kv_head": 12,
+    "n_embd": 1536,
+    "window_pattern": "SSSL"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 1048576,
+  "note": "DecoderStack-medium (d24) capture converted for nanochat. This meta was built after the fact from the run log (full_d24_w8.log); DecoderStack's capture writes only {step, code, weights} and does not emit a meta.json of its own. model_config is the nanochat GPTConfig equivalent of StackConfig -- the two architectures are the same model (both count 1,384,122,122 parameters). Sufficient for load/eval; NOT sufficient for resume: there is no dataloader_state_dict, loop_state, or optimizer state here, and DecoderStack's pre-tokenized binary loader has no state nanochat's tokenizing loader could consume. device_batch_size/max_seq_len are the nanochat-shaped restatement of micro_batch_tokens=65,536 (= 32 x 2048) -- DecoderStack trains varlen, so its micro-batch is a token count, not a sequence count. TOKENIZER: these weights use the 32k vocab shipped in the ChrisMcCormick/climbmix_32k_8_170 dataset repo (mirrored under tokenizer/ in this repo). It is NOT the same vocabulary as ChrisMcCormick/nanochat-varlen-d24-2026-03-22 -- 31,474 of 32,759 ids differ. Pairing these weights with that tokenizer produces garbage.",
+  "stack_config": {
+    "n_layers": 24,
+    "d_model": 1536,
+    "d_vocab": 32768,
+    "d_smr_gate": 24,
+    "n_q_heads": 12,
+    "n_kv_heads": 12,
+    "n_o_heads": 12,
+    "d_qk": 128,
+    "d_vo": 128,
+    "seq_len": 2048,
+    "short_win_size": 768,
+    "full_ctxt_layers": [
+      3,
+      7,
+      11,
+      15,
+      19,
+      23
+    ],
+    "d_ve_gate": 12,
+    "ve_layers": [
+      1,
+      3,
+      5,
+      7,
+      9,
+      11,
+      13,
+      15,
+      17,
+      19,
+      21,
+      23
+    ],
+    "d_mlp": 6144,
+    "micro_batch_tokens": 65536,
+    "total_batch_size": 1048576,
+    "num_iterations": 5568,
+    "max_num_docs": 192,
+    "num_train_shards": 80,
+    "val_tokens": 10485760,
+    "val_loss_every": 250,
+    "save_steps": [
+      1950,
+      5568
+    ],
+    "world_size": 8,
+    "grad_accum_steps": 2,
+    "precision": "bf16",
+    "run_id": "2026-08-01_023219-d24"
+  },
+  "schedules": {
+    "batch_lr_scale": 1.4142135623730951,
+    "muon_weight_decay_peak": 0.059738,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "matrix_lr": 0.02,
+    "muon_momentum": "0.85->0.97 over 400 steps, cooled to 0.90",
+    "muon_beta2": 0.9,
+    "muon_ns_steps": 5,
+    "unembedding_lr": 0.008,
+    "embedding_lr": 0.3,
+    "value_embedding_lr_mult": 0.5,
+    "scalar_lr": 0.5,
+    "smear_lr": 0.2,
+    "adamw_eps": 1e-10
+  },
+  "val_bpb_note": "not measured at this step (val_loss_every=250); bracketed by step 1750 = 0.815902 and step 2000 = 0.809199. This is the last step before the LR/momentum cooldown begins."
+}

base_checkpoints/d24_decoderstack/meta_005568.json

@@ -0,0 +1,93 @@
+{
+  "step": 5568,
+  "val_bpb": 0.719042,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 32768,
+    "n_layer": 24,
+    "n_head": 12,
+    "n_kv_head": 12,
+    "n_embd": 1536,
+    "window_pattern": "SSSL"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 1048576,
+  "note": "DecoderStack-medium (d24) capture converted for nanochat. This meta was built after the fact from the run log (full_d24_w8.log); DecoderStack's capture writes only {step, code, weights} and does not emit a meta.json of its own. model_config is the nanochat GPTConfig equivalent of StackConfig -- the two architectures are the same model (both count 1,384,122,122 parameters). Sufficient for load/eval; NOT sufficient for resume: there is no dataloader_state_dict, loop_state, or optimizer state here, and DecoderStack's pre-tokenized binary loader has no state nanochat's tokenizing loader could consume. device_batch_size/max_seq_len are the nanochat-shaped restatement of micro_batch_tokens=65,536 (= 32 x 2048) -- DecoderStack trains varlen, so its micro-batch is a token count, not a sequence count. TOKENIZER: these weights use the 32k vocab shipped in the ChrisMcCormick/climbmix_32k_8_170 dataset repo (mirrored under tokenizer/ in this repo). It is NOT the same vocabulary as ChrisMcCormick/nanochat-varlen-d24-2026-03-22 -- 31,474 of 32,759 ids differ. Pairing these weights with that tokenizer produces garbage.",
+  "stack_config": {
+    "n_layers": 24,
+    "d_model": 1536,
+    "d_vocab": 32768,
+    "d_smr_gate": 24,
+    "n_q_heads": 12,
+    "n_kv_heads": 12,
+    "n_o_heads": 12,
+    "d_qk": 128,
+    "d_vo": 128,
+    "seq_len": 2048,
+    "short_win_size": 768,
+    "full_ctxt_layers": [
+      3,
+      7,
+      11,
+      15,
+      19,
+      23
+    ],
+    "d_ve_gate": 12,
+    "ve_layers": [
+      1,
+      3,
+      5,
+      7,
+      9,
+      11,
+      13,
+      15,
+      17,
+      19,
+      21,
+      23
+    ],
+    "d_mlp": 6144,
+    "micro_batch_tokens": 65536,
+    "total_batch_size": 1048576,
+    "num_iterations": 5568,
+    "max_num_docs": 192,
+    "num_train_shards": 80,
+    "val_tokens": 10485760,
+    "val_loss_every": 250,
+    "save_steps": [
+      1950,
+      5568
+    ],
+    "world_size": 8,
+    "grad_accum_steps": 2,
+    "precision": "bf16",
+    "run_id": "2026-08-01_023219-d24"
+  },
+  "schedules": {
+    "batch_lr_scale": 1.4142135623730951,
+    "muon_weight_decay_peak": 0.059738,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "matrix_lr": 0.02,
+    "muon_momentum": "0.85->0.97 over 400 steps, cooled to 0.90",
+    "muon_beta2": 0.9,
+    "muon_ns_steps": 5,
+    "unembedding_lr": 0.008,
+    "embedding_lr": 0.3,
+    "value_embedding_lr_mult": 0.5,
+    "scalar_lr": 0.5,
+    "smear_lr": 0.2,
+    "adamw_eps": 1e-10
+  },
+  "results": {
+    "core_metric": 0.2517,
+    "min_val_bpb": 0.719042,
+    "peak_memory_mib": 71569,
+    "total_training_time_min": 110.74,
+    "gpu": "8x NVIDIA H100 80GB HBM3"
+  }
+}

convert_ckpt_to_nanochat.py

@@ -0,0 +1,271 @@
+# Convert a DecoderStack-medium (d24) capture into a nanochat checkpoint.
+#
+# DecoderStack writes its own two-file capture (see § Checkpoint capture in
+# decoderstack_medium_pt-sft.py):
+#
+#   model_stepNNNNNN.pt   {step, code, weights: {name: tensor}}
+#                         bf16 live weights + fp32 scalars, banked over layers.
+#   optim_stepNNNNNN.pt   {step, t_step, state: {"name.attr": tensor}}
+#                         mantissa / frst_mntm / scnd_mntm / exp_avg / exp_avg_sq,
+#                         all-gathered to full size (world-agnostic).
+#
+# nanochat wants one flat state_dict per capture, in its own module-path key
+# names, next to a meta_NNNNNN.json:
+#
+#   <out_dir>/model_NNNNNN.pt   torch.save(model.state_dict())
+#   <out_dir>/meta_NNNNNN.json  {"step", "val_bpb", "model_config", ...}
+#
+# This script does that translation. The architectures are the same model --
+# DecoderStack-medium is a flattened port of nanochat d24 -- so every tensor has
+# a home and nothing is reshaped or transposed: DecoderStack's banks index the
+# layer on dim 0 and each slice already uses F.linear's (out, in) convention.
+# The only real work is naming, unbanking, and dtype.
+#
+# DTYPE, and why the optimizer file matters
+# -----------------------------------------
+# nanochat holds fp32 master weights for everything except the two embedding
+# tables, which init_weights() casts to COMPUTE_DTYPE (bf16). DecoderStack holds
+# bf16 LIVE weights everywhere plus a uint16 `mantissa` in the optimizer file --
+# the fp32 master's bit pattern is (live_bf16_bits << 16) | mantissa. So:
+#
+#   nanochat fp32 params  <- fp32 master  = live + mantissa   (needs --optim)
+#   nanochat bf16 params  <- live bf16 as-is                  (mantissa dropped,
+#                            which is correct: nanochat's bf16 embedding IS its
+#                            master, it has no lower bits to carry)
+#
+# Without --optim the fp32 params are filled with the bf16 live values upcast to
+# fp32 -- the right dtype, but only bf16 precision. That is what you get if you
+# only downloaded the model file, and it is fine for eval: on a held-out English
+# paragraph the d24 step-5568 capture scores 2.4828 nats/token with the masters
+# 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.
+#
+# 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
+#
+# Then, in nanochat (branch fa-varlen):
+#   from nanochat.checkpoint_manager import build_model
+#   model, tokenizer, meta = build_model(checkpoint_dir, 5568, device, "eval")
+#
+# 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
+# different 32k tokenizer produces garbage, not slightly-worse text.
+import argparse
+import json
+import os
+import shutil
+
+import torch
+
+
+def fp32_master(live: torch.Tensor, mantissa: torch.Tensor | None) -> torch.Tensor:
+    """Rebuild the fp32 master from bf16 live bits + the stashed lower 16 bits.
+
+    Mirrors fp32_master() in the training script. int32 rather than uint32
+    because CUDA has no uint32 shifts as of torch 2.9; the truncating .to(int16)
+    and the <<16 discard of sign-extension bits make the two equivalent. With no
+    mantissa this degrades to a plain upcast (the mantissa bits read as zero),
+    which is exactly the bf16-precision fallback documented above.
+    """
+    if mantissa is None:
+        return live.float()
+    assert mantissa.shape == live.shape, f"mantissa {tuple(mantissa.shape)} != live {tuple(live.shape)}"
+    bits = (live.view(torch.int16).to(torch.int32) << 16) | \
+           (mantissa.view(torch.int16).to(torch.int32) & 0xFFFF)
+    return bits.view(torch.float32)
+
+
+def ve_layers(n_layer: int) -> list[int]:
+    """Layers carrying a value embedding, in bank-slot order.
+
+    nanochat's has_ve(): alternating layers, last layer always included. Identical
+    to StackConfig.ve_layers, and the ascending order matches the VE bank's slot
+    order (cfg.ve_index), so slot j belongs to layer ve_layers(n_layer)[j].
+    """
+    return [i for i in range(n_layer) if i % 2 == (n_layer - 1) % 2]
+
+
+def convert(model_data: dict, optim_state: dict | None) -> dict:
+    """DecoderStack weights dict -> nanochat state_dict."""
+    w = model_data["weights"]
+    mant = {} if optim_state is None else optim_state
+
+    def master(name, i=None):
+        """fp32 param: live + mantissa. The mantissa of a sharded weight was
+        all-gathered to full size at capture, so it already lines up 1:1. Pass
+        `i` to rebuild one slice of a bank -- worth it on the MLP banks, where
+        materializing the whole thing in fp32 would cost ~0.9 GB per bank."""
+        live, m = w[name], mant.get(f"{name}.mantissa")
+        if i is not None:
+            live, m = live[i], (None if m is None else m[i])
+        return fp32_master(live, m)
+
+    n_layer = w["W_Q"].shape[0]
+    sd = {}
+
+    # --- Embeddings: bf16 in nanochat, so the live weights go in untouched. ---
+    sd["transformer.wte.weight"] = w["input_embeds"]
+    for slot, layer in enumerate(ve_layers(n_layer)):
+        sd[f"value_embeds.{layer}.weight"] = w["value_embeds"][slot]
+
+    # --- lm_head: AdamW in both, but nanochat keeps it fp32 (it is a Linear,
+    #     not an Embedding, so init_weights() never casts it). ---
+    sd["lm_head.weight"] = master("lm_head")
+
+    # --- Per-layer matrices: unbank dim 0. Every slice is already (out, in). ---
+    banks = {
+        "W_Q":   "transformer.h.{i}.attn.c_q.weight",
+        "W_K":   "transformer.h.{i}.attn.c_k.weight",
+        "W_V":   "transformer.h.{i}.attn.c_v.weight",
+        "W_O":   "transformer.h.{i}.attn.c_proj.weight",
+        "W_in":  "transformer.h.{i}.mlp.c_fc.weight",
+        "W_out": "transformer.h.{i}.mlp.c_proj.weight",
+    }
+    for name, template in banks.items():
+        for i in range(n_layer):
+            sd[template.format(i=i)] = master(name, i).contiguous()
+
+    # --- VE gates: banked by VE SLOT, not by layer, so they unbank through the
+    #     same slot->layer map as the value embeddings. ---
+    for slot, layer in enumerate(ve_layers(n_layer)):
+        sd[f"transformer.h.{layer}.attn.ve_gate.weight"] = master("ve_gate", slot).contiguous()
+
+    # --- Scalars: fp32-live in both, no mantissa, same names. ---
+    sd["resid_lambdas"] = w["resid_lambdas"]
+    sd["x0_lambdas"] = w["x0_lambdas"]
+    sd["smear_gate.weight"] = w["smear_gate"]
+    sd["smear_lambda"] = w["smear_lambda"]
+    sd["backout_lambda"] = w["backout_lambda"]
+    return sd
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--model", required=True, help="DecoderStack model_stepNNNNNN.pt")
+    p.add_argument("--optim", default=None,
+                   help="matching optim_stepNNNNNN.pt; supplies the mantissas that make "
+                        "the fp32 params exact. Omit for a bf16-precision conversion.")
+    p.add_argument("--out", required=True, help="output checkpoint dir (nanochat model_tag dir)")
+    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")
+    args = p.parse_args()
+
+    # 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
+    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}"
+        optim_state = optim_data["state"]
+        n_mant = sum(1 for k in optim_state if k.endswith(".mantissa"))
+        print(f"loaded {args.optim}: {len(optim_state)} state tensors, {n_mant} mantissas")
+    else:
+        print("no --optim: fp32 params will carry bf16 precision (upcast, not exact masters)")
+
+    sd = convert(model_data, optim_state)
+    total = sum(t.numel() for t in sd.values())
+    by_dtype = {}
+    for t in sd.values():
+        by_dtype[t.dtype] = by_dtype.get(t.dtype, 0) + t.numel()
+    print(f"converted: {len(sd)} tensors, {total:,} params "
+          + ", ".join(f"{n:,} {str(d).replace('torch.', '')}" for d, n in by_dtype.items()))
+
+    os.makedirs(args.out, exist_ok=True)
+    model_path = os.path.join(args.out, f"model_{step:06d}.pt")
+    torch.save(sd, model_path)
+    print(f"wrote {model_path} ({os.path.getsize(model_path):,} bytes)")
+
+    if args.meta:
+        meta_path = os.path.join(args.out, f"meta_{step:06d}.json")
+        shutil.copyfile(args.meta, meta_path)
+        with open(meta_path, encoding="utf-8") as f:
+            meta_step = json.load(f)["step"]
+        assert meta_step == step, f"meta step {meta_step} != model step {step}"
+        print(f"wrote {meta_path}")
+    else:
+        print(f"NOTE: nanochat also needs meta_{step:06d}.json in {args.out} "
+              "(model_config lives there, not in the .pt)")
+
+    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:
+            f.write(model_data["code"])
+        print(f"wrote {code_path} (the exact training script for this capture)")
+
+
+if __name__ == "__main__":
+    main()
+
+
+# -----------------------------------------------------------------------------
+# OPTIMIZER-STATE NOTES
+# -----------------------------------------------------------------------------
+# Every buffer we keep has a nanochat counterpart, and the precisions line up on
+# everything except the two embedding tables:
+#
+#   DecoderStack              nanochat (MuonAdamW / DistMuonAdamW)     precision
+#   ------------------------  --------------------------------------  ---------
+#   .frst_mntm      fp32      Muon  state["momentum_buffer"]           fp32 both
+#   .scnd_mntm      fp32      Muon  state["second_momentum_buffer"]    fp32 both
+#   .exp_avg        fp32      AdamW state["exp_avg"]                   see below
+#   .exp_avg_sq     fp32      AdamW state["exp_avg_sq"]                see below
+#   .mantissa     uint16      (no counterpart -- nanochat's fp32 param IS
+#                              the master; consumed above to rebuild it)
+#
+# nanochat allocates its Muon buffers as `dtype=p.dtype` and its AdamW buffers as
+# `torch.zeros_like(p)`. Its Muon params and lm_head are fp32, so those match us.
+# But wte and value_embeds are bf16 params, so THEIR AdamW moments are bf16 --
+# where ours are fp32. That is the one precision difference, and it is ours that
+# is the more precise of the two: we pair the embeddings with a mantissa so a
+# single AdamW kernel serves every param, instead of carrying a second bf16-live
+# variant (see the dtype scheme in the training script).
+#
+# Two things would still block a faithful resume INTO nanochat, so this script
+# does not pretend to offer one:
+#
+# 1. W_O's second moment is a different quantity. NorMuon's factored second
+#    moment is a per-neuron mean-square, and the two codebases disagree about
+#    which axis holds the neurons for the attention output projection. nanochat
+#    infers it from the shape -- `red_dim = -1 if shape[-2] >= shape[-1] else -2`
+#    -- which for a square (1536, 1536) c_proj picks -1. DecoderStack sets it
+#    explicitly (m.W_O.residual_dim = -2) because W_O's heads are stored
+#    transposed relative to QKV. So our W_O.scnd_mntm is (1, 1536) where
+#    nanochat's is (1536, 1); they are not transposes of each other, they are
+#    reductions along different axes. Everything else agrees (W_in -1, W_out -2,
+#    QKV -1, ve_gate -1), because there the shape heuristic happens to land on
+#    the same axis we chose deliberately.
+#
+# 2. nanochat's resume needs meta_data["dataloader_state_dict"] to place its
+#    tokenizing loader back in the stream. DecoderStack reads pre-tokenized
+#    binary shards through a completely different loader and has no such state to
+#    hand over, so the data order could not be continued regardless of optimizer
+#    state.
+#
+# If you do want the optimizer state anyway (say, to warm-start rather than
+# resume), the remaining mapping 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})`
+#     order -- (12,12) ve_gates, (1536,1536) c_q/c_k/c_v/c_proj interleaved in
+#     block order, (1536,6144) mlp.c_proj, (6144,1536) mlp.c_fc.
+#   - AdamW state for params with >= 1024 elements is sliced over dim 0 by rank;
+#     smaller ones are replicated. Ours is captured all-gathered to full size,
+#     so it just needs re-slicing (value_embeds first reshaped from its flattened
+#     (num_ves * vocab, kv_dim) row axis back to (num_ves, vocab, kv_dim)).
+#   - Muon state is stacked per group and chunked: rank r owns params
+#     [r*ceil(K/W) : (r+1)*ceil(K/W)] of the group, zero-padded when K % W != 0.

logs/baseline_report.md

@@ -0,0 +1,50 @@
+# d24 bf16, total batch 2^20: upstream nanochat vs single-file -v3
+
+All runs on the same 8xH100 box unless the label says GH200. `bf16_mfu` is **not** comparable across the two pipelines (different FLOPs/token estimates: upstream 4.775e9, ours 4.860e9) — compare **ms/step** and **tok/sec**, which measure the same 2^20 tokens per step.
+
+| run                                   | parsed     |   ms/step (steps 20-99) | tok/sec   |   bf16_mfu* |
+|---------------------------------------|------------|-------------------------|-----------|-------------|
+| ours w8 64K micro (full run)          | 5568 steps |                    1212 | 864,856   |       53.13 |
+| ours w8 32K micro +FP8                | 200 steps  |                    1040 | 1,008,030 |       61.92 |
+| ours w8 32K micro                     | 100 steps  |                    1249 | 839,838   |       51.59 |
+| upstream w8 bf16                      | 1017 steps |                    1291 | 812,435   |       49.03 |
+| upstream w8 +fp8                      | 200 steps  |                    1081 | 970,455   |        4.86 |
+| fork varlen branch w8 bf16            | 200 steps  |                    1278 | 820,783   |       49.54 |
+| fork varlen+fp8 w8 (2026-03-22, full) | 5568 steps |                    1066 | 983,648   |       59.37 |
+| ours w1 (1xH100, partial)             | 23 steps   |                    9620 | 109,002   |       53.57 |
+| ours w1 (1x GH200)                    | 100 steps  |                    9635 | 108,834   |       53.48 |
+
+| derived   | value   |
+|-----------|---------|
+
+## Correctness: does world=8 reproduce the world=1 curve?
+
+The loader assigns every rank's documents from one sequential stream and total batch is fixed, so the global batch per optimizer step is identical at every world size — these are directly comparable.
+
+| run                            |   step 0 |   step 25 |   step 50 |   step 75 |   step 100 |
+|--------------------------------|----------|-----------|-----------|-----------|------------|
+| -v3 world=8 (8xH100)           |  3.17614 |   1.75957 |   1.57767 |   1.45099 |    1.37943 |
+| -v3 world=1 (GH200, committed) |  3.17614 |   1.75956 |   1.57762 |   1.4503  |    1.3794  |
+
+| run                                 | min val bpb         | peak alloc MiB   | peak reserved MiB   | train time   |
+|-------------------------------------|---------------------|------------------|---------------------|--------------|
+| -v3 world=8                         | 1.379430            | 40049            | 65750               | 1.85m        |
+| -v3 world=1 GH200                   | 1.379400            | 49283            | 71046               | 14.29m       |
+| world=1 reference band (eval100a-e) | 1.379318 - 1.379563 | 49,088 - 49,334  | --                  | --           |
+
+
+## Head-to-head val bpb (both on shard_06542)
+
+|   step |   -v3 (full run) |   upstream baseline | delta   |
+|--------|------------------|---------------------|---------|
+|      0 |         3.17584  |            3.16882  | +0.22%  |
+|    250 |         0.997472 |            0.984017 | +1.37%  |
+|    500 |         0.904081 |            0.898168 | +0.66%  |
+|    750 |         0.866224 |            0.861772 | +0.52%  |
+|   1000 |         0.846058 |            0.841709 | +0.52%  |
+
+The upstream baseline was stopped at step 1000, so it bounds the comparison; beyond that the -v3 curve stands alone.
+
+Full-run val bpb curve: 0:3.175835 / 250:0.997472 / 500:0.904081 / 750:0.866224 / 1000:0.846058 / 1250:0.832731 / 1500:0.822766 / 1750:0.815902 / 2000:0.809199 / 2250:0.798533 / 2500:0.789216 / 2750:0.781275 / 3000:0.773139 / 3250:0.766625 / 3500:0.759539 / 3750:0.753329 / 4000:0.747155 / 4250:0.741616 / 4500:0.736278 / 4750:0.731419 / 5000:0.726908 / 5250:0.722909 / 5500:0.719696 / 5568:0.719042
+
+_Produced by `make_baseline_report.py`; sources listed in its docstring._

logs/run_full_d24_w8.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+# THE FULL-HORIZON d24 RUN: 5,568 steps, world=8, bf16, total batch 2^20.
+# ~1.87 h at the measured 1,208 ms/step. This is the run the checkpoints come
+# from, so it is the deliverable.
+#
+# Config, and why each value:
+#   micro_batch_tokens 32768 -> 65536  grad accum 4 -> 2. Measured 1,208 vs
+#       1,249 ms/step (~3.2% faster) at 71,568 MiB peak allocated against the
+#       card's 81,559 -- 8.8 GB spare. Confirmed end-to-end at this size
+#       including checkpoint capture and CORE eval (db64k_confirm.log).
+#   max_num_docs 96 -> 192             Worst-case packing over all 79 train
+#       shards this run reads is 162 docs at 65,536 tokens, so 163 is the
+#       minimum (the assert is `len(cum_lengths) < max_num_docs`). 192 gives
+#       18% margin at zero measured cost: an A/B at 163/192/256 came out
+#       1207.6/1209.7/1207.3 ms -- ghost cu_seqlens entries do not register at
+#       d24 (scan_max_docs_sweep.log).
+#   NUM_TRAIN_SHARDS 20 -> 80          Horizon is 5,838,471,168 tokens; shards
+#       1..79 hold 7.9B raw ~= 6.7B usable after seq_len truncation, +876M
+#       margin. All 79 are already on disk (prefetch_train_shards.log).
+#
+# Left at the file's defaults deliberately: num_iterations 5568, warmup 40,
+# val_loss_every 250, save_checkpoint True, save_steps (1950,) + the always-saved
+# final -- the two capture points Chris chose (1950 is the last uncooled state;
+# the cooldown covers steps 1950-5567).
+#
+# NOTE: val bpb is NOT comparable to the eval100a-e baselines -- the validation
+# shard was rebuilt from shard_06542 to match nanochat's pinned val split
+# (regen_val_shard_06542.py). It IS comparable to the upstream baseline curve in
+# upstream_speedrun_d24_bf16_world8.log.
+#
+# Checkpoints land in $RUN_DIR/logs/<run_id>/ (under ~/.cache). Push them with
+#   python ~/agent-ops/stacks/2026-07-31_0822am_fable-rewrite-handoff/push_checkpoints.py \
+#       <run_dir>/logs/<run_id> <hf-model-repo>
+# BEFORE releasing the box.
+set -eu
+source ~/nanochat/.venv/bin/activate
+source ~/env.sh
+export OMP_NUM_THREADS=1
+export DATA_PATH="$HOME/.cache/stacks"
+export PYTHONUNBUFFERED=1
+
+SRC=~/agent-ops/stacks/2026-07-31_0822am_fable-rewrite-handoff
+SESS=~/agent-ops/stacks/2026-07-31_0627pm_8xh100-d24-baseline-world8
+RUN_DIR="$DATA_PATH/run_full_d24_w8"
+mkdir -p "$RUN_DIR"
+cp "$SRC/decoderstack_medium_pt-sft-fable-v3.py" "$RUN_DIR/run_full_d24_w8.py"
+cd "$RUN_DIR"
+
+sed -i 's/    micro_batch_tokens: int = 32768/    micro_batch_tokens: int = 65536/' run_full_d24_w8.py
+sed -i 's/^    max_num_docs = 96$/    max_num_docs = 192/' run_full_d24_w8.py
+# NB: no `$` anchor here -- this line carries a trailing comment, and a
+# $-anchored pattern silently matched nothing (the earlier eval run-scripts
+# have the same latent no-op; harmless there because NUM_TRAIN_SHARDS only
+# gates DOWNLOADS, never the loader, which globs every train_*.bin on disk).
+sed -i 's/^NUM_TRAIN_SHARDS = 20/NUM_TRAIN_SHARDS = 80/' run_full_d24_w8.py
+
+echo "=== applied run-copy edits ==="
+grep -n "micro_batch_tokens: int\|max_num_docs = \|NUM_TRAIN_SHARDS = \|num_iterations: int\|save_steps:\|save_checkpoint: bool" run_full_d24_w8.py
+
+# Fail loudly on a sed that matched nothing, rather than starting a ~2 h run on
+# a config nobody verified.
+check() { grep -qE "$1" run_full_d24_w8.py || { echo "EDIT FAILED: $1"; exit 1; }; }
+check '^    micro_batch_tokens: int = 65536'
+check '^    max_num_docs = 192$'
+check '^NUM_TRAIN_SHARDS = 80'
+check '^    num_iterations: int = 5568$'
+check '^    save_steps:      tuple = \(1950,\)$'
+check '^    save_checkpoint: bool = True$'
+echo "=== all run-copy edits verified ==="
+echo "=== run dir: $RUN_DIR ==="
+
+torchrun --standalone --nproc_per_node=8 run_full_d24_w8.py 2>&1 | tee "$SESS/full_d24_w8.log"

tokenizer/token_bytes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfe1000a3605b109246baaeb9266fd6f151c1f7eddcd0541a083aa585977784f
+size 132649

tokenizer/tokenizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb740019e0eb230dc8ed46a2197479c2bc9476f121a8532fdabdbb9f797f54f0
+size 412035