decoderstack-d24 / README.md
ChrisMcCormick's picture
Add model card, meta.json, logs, tokenizer, training source, and nanochat converter
11f07f9 verified
|
Raw History Blame
8.47 kB
metadata
license: mit
language:
  - en
tags:
  - decoderstack
  - nanochat
  - pretraining

DecoderStack d24

A 1.38B-parameter d24 model pre-trained from scratch with DecoderStack β€” 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 once converted β€” see 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.

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, mirrored here under 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.

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 branch (varlen, matching how these weights were trained):

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 β€” the converter is committed at e664830 (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 (ClimbMix, pre-tokenized to 32k binary shards). Validation is a rebuild of shard 06542 to match nanochat's pinned val split.