Correct optimizer-state notes: SFT needs only the weights; W_O axis is a no-op at d24

README.md

@@ -83,8 +83,9 @@ 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.
+the LR/momentum hold ends after update 1949 (`N - round(0.65*N)`), so 1950 is the capture
+to branch from if the horizon is ever extended. (Extending it needs a load path that does
+not exist yet in either codebase — see [Optimizer state](#optimizer-state).)
 
 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`
@@ -146,8 +147,19 @@ 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:
+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.
+
+**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
+loader back in the stream, and DecoderStack's pre-tokenized binary loader has no equivalent
+state to hand over. The data order could not be continued no matter what the optimizer held.
+
+Every buffer does have a counterpart, and the precisions line up on everything except the
+embedding tables:
 
 | DecoderStack | nanochat | precision |
 |---|---|---|
@@ -156,19 +168,25 @@ counterpart, and the precisions line up on everything except the embedding table
 | `.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`.
+DecoderStack is the more precise of the two on the embeddings, deliberately: only the
+*gradients* are bf16 for those tables (they are the model's largest tensors), while the
+moment math stays fp32 — the kernel upcasts on the way in. There is no bf16 AdamW variant;
+the two AdamW kernels differ in whether the param carries a mantissa, not in moment dtype.
+
+`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`.
+
+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`.
 
 ## Provenance
 
 Trained by [`chrisjmccormick/stacks`](https://github.com/chrisjmccormick/stacks) — the
 converter is committed at
-[`e664830`](https://github.com/chrisjmccormick/stacks/commit/e6648306442e0479433b287ede02d271997365c0)
+[`2e199ac`](https://github.com/chrisjmccormick/stacks/commit/2e199acd743c6dcdb4d768daae9b73500c9354a8)
 (`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,

base_checkpoints/d24_decoderstack/meta_001950.json

@@ -13,7 +13,7 @@
   "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.",
+  "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 and for SFT (chat_sft builds a fresh optimizer and only optionally warm-starts it). NOT sufficient for pre-training 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 -- the data order could not be continued regardless. 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,

base_checkpoints/d24_decoderstack/meta_005568.json

@@ -13,7 +13,7 @@
   "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.",
+  "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 and for SFT (chat_sft builds a fresh optimizer and only optionally warm-starts it). NOT sufficient for pre-training 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 -- the data order could not be continued regardless. 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,

convert_ckpt_to_nanochat.py

@@ -214,6 +214,22 @@ 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.
+#
+# 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
+# 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.
+#
 # Every buffer we keep has a nanochat counterpart, and the precisions line up on
 # everything except the two embedding tables:
 #
@@ -228,36 +244,34 @@ if __name__ == "__main__":
 #
 # 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.
+# 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 ours is the more
+# precise of the two, deliberately: it is only the GRADIENTS that are bf16 for
+# those two tables (they are the biggest tensors in the model, so fp32 grads
+# would double their scatter and comm traffic, and bf16 matches the autograd
+# baseline's numerics). The moment math stays fp32 -- adamw_step_fused upcasts on
+# the way in, `grad = grad.to(exp_avg.dtype)`. There is no bf16 AdamW variant in
+# the file: the two AdamW kernels differ in whether the param carries a mantissa
+# (adamw_step_fused vs adamw_step_fused_fp32), not in moment dtype.
 #
-# 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.
+# W_O's reduction axis differs between the two, and at d24 it costs nothing.
+# NorMuon's factored second moment is a per-neuron mean-square; nanochat infers
+# the neuron axis from the shape (`red_dim = -1 if shape[-2] >= shape[-1] else
+# -2`) while DecoderStack states it (m.W_O.residual_dim = -2), so the two
+# disagree on a square c_proj -- ours is (1, 1536) where nanochat's is (1536, 1).
+# But polar express returns a ~orthonormal update, and a square orthonormal
+# matrix has ~uniform neuron norms along either axis: there is no variance to
+# reduce, the rescale is a ~no-op, and the run is unaffected by the choice.
+# The explicit axis earns its keep only when n_heads * d_head != d_model. Above
+# 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).
 #
-# 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:
+# 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:
 #   - 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})`