cayley-24L2048-131k-3L-mlp_in-v1
Current canonical 24L/2048d CayleySAE 3-level model at mlp_in.
1.3B-param GPT (24 layers, 16 heads, d=2048) with a 3-level CayleySAE
sparsity bottleneck inserted at mlp_in in every block. Trained on
FineWeb-Edu-100B for 10B tokens with a cosine LR schedule
(peak muon_lr 6e-3, floor 6e-4) on 4× B200.
Checkpoint is ckpt_best.pt from iter 18,500 of a 19,074-iter schedule —
best val reached during the run. Training ran to completion.
Architecture
| Parameter | Value |
|---|
| n_layer | 24 |
| n_head | 16 |
| n_embd | 2048 |
| block_size | 1024 |
| vocab_size | 50304 |
| bias | False |
| norm | RMSNorm (affine) |
| MLP | GELU, 4× expansion |
| pos_encoding | learned |
| tokenizer | GPT-2 (tiktoken) |
| dtype | bfloat16 |
| sparsity_mode | cayley |
| cayley_locations | mlp_in |
| cayley_levels | 11,16,0; 15,32,256; 17,64,256 |
| cayley_per_parent_budget | True |
| cayley_score_standardize | True |
Level layout ([L, k, delta]):
- L0 = (2^11 = 2,048 features, k=16)
- L1 = (2^15 = 32,768 features, k=32, delta=256)
- L2 = (2^17 = 131,072 features, k=64, delta=256)
Total: 165,888 features per layer, 112 active per token.
Training
| Parameter | Value |
|---|
| optimizer | Muon (hidden 2D) + AdamW (embeddings) |
| muon_lr | 6e-3 peak → 6e-4 floor |
| adamw_lr | 6e-3 peak → 6e-4 floor |
| lr_schedule | cosine |
| warmup_iters | 200 |
| muon_momentum | 0.95 |
| muon_ns_steps | 5 |
| batch_size | 64 |
| seq_len | 1024 |
| grad_accum_steps (global) | 8 |
| tokens/iter | 524,288 |
| max_iters | 19,074 |
| uploaded checkpoint iter | 18,500 (best val) |
| tokens seen at upload | ~9.70B |
| dataset | FineWeb-Edu-100B |
| best_val_loss | 2.9299 |
Campaign context
This is run v4 of the 3-level 24L2048 CayleySAE campaign. Prior runs:
| run | tokens | peak LR | best val | notes |
|---|
| v2-cosine (resume) | 26B | 1.2e-2 | ~3.00 | spiked at iter 24.4k, regressed 0.16 nats |
| v3-cosine | 20B | 1.2e-2 | 2.9058 | mild grad_norm bump at iter 25k, val resilient |
| v4-halflr (this) | 10B | 6e-3 | 2.9299 | clean run, no spike |
v4 is picked as the canonical upload for two reasons:
- Completed cleanly — no grad_norm episode of concern.
- Half the compute of v3 and a flat, stable training trajectory with val
0.024 nats behind v3. The loss gap is real but the stability makes this
the more reliable artifact for downstream interpretability work.
A successor experiment is queued (v5-cosine-highlr, 10B @ peak 1.2e-2) as
a direct LR-only control to pin down whether v3's ~0.024-nat win over v4 is
the LR or the longer 20B schedule.
Files
ckpt.pt — torch checkpoint dict with keys model, config,
model_config, iter_num=18500, best_val_loss=2.9299,
wandb_step_offset, parity_trigger_iter=-1. Torch 2.4+.config.json — full training config used by the run (not the model_config
only — includes optimizer / schedule / data knobs).
Loading
1import torch
2from deeptopk.src.sparse_nanogpt.model import GPT, DeepTopKGPTConfig
3
4ckpt = torch.load("ckpt.pt", map_location="cpu", weights_only=False)
5model_cfg = DeepTopKGPTConfig.from_dict(ckpt["model_config"])
6model = GPT(model_cfg)
7state = {k.removeprefix("module.").removeprefix("_orig_mod."): v
8 for k, v in ckpt["model"].items()}
9model.load_state_dict(state)
10model.eval()
Related
markhenry/cayley-24L2048-131k-3L-mlp_in-20b-v3-cosine — v3 run (longer,
higher LR, lower val, grad_norm episode in the middle).markhenry/vanilla-24L2048-parity-cold — 1.3B vanilla GPT at matched val
loss (2.7926) for alignment-tax comparisons.markhenry/cayley-32k-2L-mlp_in — smaller 12L/1024d 2-level variant.aemack-org/cayley-mlp-in-1p3b — RoPE variant on the same chassis with
larger 2^20 leaf dictionary.
Model card authored by Claude (Opus 4.7) during training campaign on 2026-04-20.
