cayley-24L2048-131k-3L-mlp_in-v3
A 1.31B-parameter GPT with a 3-level CayleySAE inserted at mlp_in in every
transformer block, trained on 16B tokens of FineWeb-Edu. First successful clean
completion at this scale.
Headline
- Val loss: 2.8501 (iter 10173, final)
- Training tokens: 16B (FineWeb-Edu)
- Wall clock: ~16h on 8× RTX PRO 6000 Blackwell (96 GB)
This is the new canonical 24L/2048d 131k-3L model. It beats the spike-corrupted
cayley-24L2048-131k-3L-mlp_in-v2-cosine
best of 2.9522 by
0.102 nats, using 16B tokens vs 26B.
Backbone
- 24 transformer blocks, d_model 2048, 16 heads (head_dim 128)
- ~1.31B total parameters (~1.21B non-embedding)
- seq_len 1024
- RMSNorm, RoPE position embeddings
CayleySAE
Inserted at mlp_in in every block: RMSNorm → CayleySAE → MLP. Output is
dense d=2048; sparsity lives in the intermediate code.
- 3 levels with hierarchy
11,16,0;15,32,256;17,64,256
- L0:
n=11 (2048 coords), k=16
- L1:
n=15 (32k leaves), k=32, parent budget 256
- L2:
n=17 (131k leaves), k=64, parent budget 256
- 112 active features per token (16 + 32 + 64)
- Parameter-free algebraic dictionary; only per-feature biases are learned
cayley-per-parent-budget, cayley-score-standardize, and
cayley-forward-standardized (the "zombie fix" — see report 29 in repo)
all enabled
n0 is forced to log2(d_model) = 11; this requires the t=11 primitive
polynomial x^11 + x^2 + 1 in deeptopk/f2_algebra.py.
Training recipe
| Knob | Value |
|---|
| Optimizer | Muon (2D weights) + AdamW (embeddings, biases) |
| Peak Muon LR | 8e-3 |
| Min Muon LR | 1.5e-4 |
| Peak AdamW LR | 1.4e-3 (decoupled, 5.7× muon ratio) |
| Min AdamW LR | 3e-5 |
| LR schedule | linear_warmdown |
warmdown_frac | 0.5 |
| Warmup iters | 200 |
| Batch size (per rank) | 24 |
| Gradient accumulation | 64 (global) |
| Tokens per iter | 1,572,864 |
| Total iters | 10,173 |
| World size | 8× RTX PRO 6000 Blackwell |
| Dataset | FineWeb-Edu-100B |
Warmup 0 → 200 iters; flat phase 200 → 5087 iters; warmdown 5087 → 10173
(linear 8e-3 → 1.5e-4 for muon, 1.4e-3 → 3e-5 for adamw).
This run was resumed mid-stream from iter 3240 of an earlier 26B/wf=0.9
schedule, with a revised 16B/wf=0.5 plan. Sub-Chinchilla wf=0.5 was chosen per
the saturation-dependent warmdown rule (D=16B = 0.61× Chinchilla for 1.31B
params).
Recipe deltas vs prior 24L/2048d 131k-3L attempts
The two prior attempts to train this configuration failed in distinct ways:
- v1 (linear_warmdown wf=0.9, lockstep LR 1.2e-2): bailed at iter 5250 on
throughput economics; trajectory was healthy (val 2.926 at 5.1B tokens).
- v2-cosine (cosine schedule, lockstep LR 1.2e-2 → 1.2e-3): reached val
2.9522 best, then spiked at iter 27,750 (49% through cosine, LR ~6.7e-3)
with grad_norm 0.10 → 0.27. Never recovered; final 2.99.
This run (v3) made four changes:
cayley_forward_standardized=true — z-scores forwarded into reconstruction
to avoid latent rescaling pathology- RoPE position encoding (replacing learned)
- Muon peak 1.2e-2 → 8e-3 (precautionary given zombie fix on 24L)
- AdamW LR decoupled from Muon (5.7× muon ratio)
linear_warmdown was retained over cosine — at the same LR band that triggered
the v2 spike (~6.7e-3), this run's schedule passed through cleanly.
Training health
| Signal | Outcome |
|---|
| grad_norm in warmdown | stable 0.07–0.08 throughout |
| L0 dead features (of 2048) | ≤0.4% throughout (single-point peak at iter 7000); 0.0% at completion |
| L1 / L2 dead features | 0/30720 and 0/98304 throughout |
l0_mu_mean | converged from −0.0066 → −0.0010 (unbiased) |
l0_sigma_mean | monotonically declined from 0.5500 → 0.4759 |
Hierarchy quick-eval (final)
| Metric | Value |
|---|
weighted_mean_refinement_gain_nats | 1.9214 |
mean_normalized_refinement_gain | 0.4262 |
frac_well_resolved | 1.000 (100%) |
mean_eff_parent_count | 17.18 |
context_coherence/mean_parent_minus_sibling | 0.0551 |
context_coherence/frac_positive_parent_minus_sibling | 75.6% |
child_selection_ablation/delta_ce_mean | −0.0023 ± 0.0013 SE (1.7σ; helps CE) |
Files
ckpt.pt — PyTorch checkpoint (9.1 GB). Contains model,
optimizer_states, config, model_config, iter_num, best_val_loss.config.json — training config snapshot.
Loading
1import torch
2from sparse_nanogpt.model import GPT
3from sparse_nanogpt.config import DeepTopKGPTConfig
4
5ckpt = torch.load("ckpt.pt", map_location="cpu", weights_only=False)
6model_config = DeepTopKGPTConfig(**ckpt["model_config"])
7model = GPT(model_config)
8model.load_state_dict(ckpt["model"])
Citation
Part of the
Sparse NanoGPT project.
