You can deploy this model and start earning money today!

markhenry/
cayley-24L2048-131k-3L-mlp_in-20b-v3-cosine

Name: cayley-24L2048-131k-3L-mlp_in-20b-v3-cosine
Availability: InStock

Views

No views yet

cayley-24L2048-131k-3L-mlp_in-20b-v3-cosine

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 with a cosine LR schedule (peak 1.2e-2, floor 1.2e-3) on 4× B200.

This checkpoint is from iter 35,250 of a 38,147-iter schedule — best val (2.9058) during the run. It sits deep inside a grad_norm episode that started around iter 17k and peaked at ~0.2 around iter 25k. Val loss was resilient through the episode (adjacent-eval numbers 2.906–2.910 in the floor cluster), but the 26B-token predecessor (v2-cosine, same config, longer schedule) spiked harder and regressed 0.16 nats without recovering — so treat weights from this checkpoint as taken from a schedule that was walking a line.

A successor run at half the peak LR (6e-3) and half the tokens (10B) is running to test whether the pathology is purely LR-driven or token-count- dependent.

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 = 2048 features, k=16)
L1 = (2^15 = 32,768 features, k=32, delta=256)
L2 = (2^17 = 131,072 features, k=64, delta=256)

Training

Parameter	Value
optimizer	Muon (hidden 2D) + AdamW (embeddings)
muon_lr	1.2e-2 peak → 1.2e-3 floor
adamw_lr	1.2e-2 peak → 1.2e-3 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
scheduled max_iters	38,147
uploaded checkpoint iter	35,250 (best_val)
training killed at iter	37,220 (early kill; floor plateau)
tokens seen at upload	~18.5B
dataset	FineWeb-Edu-100B
best_val_loss	2.9058

Grad-norm episode

iters 0–17k: grad_norm stable ~0.10–0.15.
iters ~17k–24k: slow upward creep (premonitory, no val regression).
iter ~25k: spike to ~0.2, no matching val reaction.
iters 25k–37k: elevated 0.13–0.16 band, val descended to the 2.906 floor cluster then plateaued.

For comparison, v2-cosine (same config, 26B-token schedule, resumed) spiked at ~iter 24.4k (~49% of schedule) and regressed val 0.16 nats without recovering. v3 was a fresh-init run at 20B and came through with only a mild grad_norm bump, though both runs show the same mid-schedule signature.

See mh/reports/ in the source repo for the full campaign context.

Files

ckpt.pt — torch checkpoint dict with keys model, optimizer_states, config, model_config, iter_num=35250, best_val_loss=2.9058, 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

python

import torch
from deeptopk.src.sparse_nanogpt.model import GPT, DeepTopKGPTConfig

ckpt = torch.load("ckpt.pt", map_location="cpu", weights_only=False)
model_cfg = DeepTopKGPTConfig.from_dict(ckpt["model_config"])
model = GPT(model_cfg)
state = {k.removeprefix("module.").removeprefix("_orig_mod."): v
         for k, v in ckpt["model"].items()}
model.load_state_dict(state)
model.eval()

Model card authored by Claude (Opus 4.7) during training campaign on 2026-04-19.

markhenry/cayley-24L2048-131k-3L-mlp_in-20b-v3-cosine