Model card: step 5568 is now a ready-made nanochat checkpoint (+8-rank optimizer shards)

README.md

@@ -85,9 +85,12 @@ checkpoints/                        native DecoderStack captures (the source of
   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)
+base_checkpoints/d24_decoderstack/  READY-MADE nanochat checkpoint -- no conversion needed
+  model_005568.pt       4.2 GB      fp32 masters + bf16 embeddings, 175 tensors
+  optim_005568_rank0.pt 1.0 GB      ZeRO-2 optimizer shards for an 8-GPU SFT run
+  ...    _rank7.pt      1.0 GB      (8 x 1.0 GB; see Optimizer state)
+  meta_005568.json                  model_config + user_config + run metadata
+  meta_001950.json                  metadata only -- step 1950 is not converted
 code/run_full_d24_w8.py             the exact training script this run executed
 logs/
   d24_decoderstack_20260801_023219.log   full training log
@@ -109,77 +112,96 @@ 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.
+**Step 5568 is already converted.** `base_checkpoints/d24_decoderstack/` is a nanochat
+checkpoint dir — download it and the tokenizer into a nanochat base dir and it loads. No
+conversion step, and the 11 GB optimizer capture is not part of it.
 
 ```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")
+import os
+from huggingface_hub import snapshot_download
+snapshot_download("ChrisMcCormick/decoderstack-d24",
+                  local_dir=os.path.expanduser("~/.cache/nanochat-ds"),
+                  allow_patterns=["base_checkpoints/d24_decoderstack/*", "tokenizer/*"])
 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/
+export NANOCHAT_BASE_DIR=~/.cache/nanochat-ds
 ```
 
-(`~/.cache/nanochat` is nanochat's default base dir; override with `NANOCHAT_BASE_DIR`.)
+**Use a base dir of its own, not a shared `~/.cache/nanochat`.** nanochat resolves the
+tokenizer from one global `tokenizer/` per base dir, and this vocabulary is not the one any
+other d24 release uses (see [⚠️ Tokenizer](#️-tokenizer)). A base dir holding both this
+checkpoint and, say, `nanochat-varlen-d24-2026-03-22` can only be right about one of them,
+and `build_model` only asserts the vocab *size* — which matches — so it will not warn you.
 
-`--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):
+Then, on
+[`varlen-gsm8k`](https://github.com/chrisjmccormick/nanochat/tree/varlen-gsm8k) (the branch
+whose SFT/RL pipeline this checkpoint targets; `fa-varlen` also works for eval):
 
 ```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")
+from nanochat.checkpoint_manager import load_model
+model, tokenizer, meta = load_model("base", torch.device("cuda"), phase="eval",
+                                    model_tag="d24_decoderstack", step=5568)
 ```
 
-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.
+Pass `model_tag` explicitly — with more than one `d24*` dir under `base_checkpoints/`,
+`find_largest_model` picks between them on depth alone and both are 24.
 
-**DecoderStack itself cannot load these back yet.** The capture path is write-only by
-design; a load/resume path is future work.
+Verified on an H100 (2026-08-01): 175 tensors, 1,384,122,122 params, shapes and dtypes
+identical to a freshly built nanochat `GPT` at this config, `strict=True` load, and
+**2.4839 nats/token (0.7255 bpb)** on held-out English — in line with the run's own 0.719
+val bpb — with coherent greedy continuations (`The capital city of France is` → ` Paris.`).
+A scrambled mapping scores 10.4.
 
-### Optimizer state
+### Re-converting yourself
 
-**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.
-
-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:
+`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. You need it to convert **step 1950**, or to emit optimizer shards for a GPU count
+other than 8:
 
 ```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 \
+    --out   $NANOCHAT_BASE_DIR/base_checkpoints/d24_decoderstack \
     --world-size 8
 ```
 
+`--optim` is optional for the weights. 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, or the optimizer shards, which need it.
+
+**DecoderStack itself cannot load these back yet.** The capture path is write-only by
+design; a load/resume path is future work.
+
+### Optimizer state
+
+**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.
+
+The `optim_005568_rank{0..7}.pt` files here remove that "slightly worse". They are the
+ZeRO-2 shards `DistMuonAdamW` expects, **cut for an 8-GPU run** — the shard count must equal
+the GPU count of the SFT run, because rank *r* holds only its own slice. For any other
+world size, re-run the converter with `--world-size N` (see above); `--world-size 1` emits
+a single rank0 file holding the full unsharded state.
+
+Verified on an H100 (2026-08-01), all 8 shards: 10 groups / 175 params / 23 state entries
+each, shapes equal to what `DistMuonAdamW` allocates per rank at this config, accepted by
+`load_state_dict`, and **reassembling bit-exactly** back to the original capture — including
+the ragged 12-param `ve_gate` group, where chunk 2 × 8 ranks leaves ranks 6–7 holding only
+zero padding. End to end, `chat_sft`'s own load block warm-starts from them (23 entries,
+step 5568) and steps.
+
+`meta_005568.json` carries a `user_config` block, which is where `chat_sft` reads the
+pretraining LRs from — without it SFT silently falls back to `unembedding_lr=0.004` instead
+of this run's **0.008**.
+
 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
@@ -215,9 +237,10 @@ square, where the two axes would carry genuinely different information. The expl
 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).
 
-**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.
+The optimizer path is verified twice over: against a real nanochat `GPT` and `MuonAdamW` at
+a toy config (structure, per-rank sharding at world sizes 1–4, reassembly, a warm-started
+`optimizer.step()`), and against the full d24 capture at `--world-size 8` on an H100
+(2026-08-01) — see above.
 
 ## Provenance
 
@@ -233,3 +256,8 @@ trainer with three launcher overrides (`micro_batch_tokens` 32768→65536,
 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.
+
+The nanochat-format checkpoint and its optimizer shards were produced and verified on
+2026-08-01 against nanochat
+[`varlen-gsm8k`](https://github.com/chrisjmccormick/nanochat/tree/varlen-gsm8k) `d9dffe9`,
+torch 2.9.1, 1×H100.