a3ca42a678
add comment
2026-04-13 14:17:23 +02:00
9822cc7424
use nn.init and initialize smear gate's weight as well
2026-04-13 14:03:18 +02:00
94b73ad29a
fix: initialize smear and backout lambdas in init_weights
2026-04-03 20:39:55 +00:00
a825e63f81
Autoresearch round 2: smear, backout, and hyperparameter tuning
...
New architectural features:
- Smear: mix previous token embedding into current position via learned
gate, providing cheap bigram-like info (works in training + KV cache)
- Backout: subtract learned fraction of mid-layer residual before logit
projection to remove low-level features
Hyperparameter tuning:
- Muon momentum warmdown 0.97→0.90 during LR warmdown phase
- Non-uniform per-layer init: resid_lambdas 1.15→1.05, x0_lambdas 0.20→0.05
- c_fc init scale 0.4x, QK norm scale 1.2, sliding window seq_len/4
- Speedrun data:params ratio reduced to 8
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-03-14 17:03:06 +00:00
6ed7d1d82c
All of these improvements were developed by Claude running autonomously over ~2 days using autoresearch. I didn't touch anything - incredible. All tuning was done on d12 but generalized easily to larger models (e.g. d24 in particular). This means we will also get a new "Time to GPT-2" Leaderboard entry, which I will push separately.
...
Optimizer & schedule changes:
- Increase unembedding LR 0.004 -> 0.008, weight decay 0.2 -> 0.28
- Per-group Adam betas and weight decay (instead of shared global betas)
- Muon beta2 0.95 -> 0.9, momentum warmup target 0.95 -> 0.97 over 400 steps
- Warmup: ratio-based -> absolute steps (default 40)
- Warmdown ratio 0.5 -> 0.65, final LR fraction 0.0 -> 0.05
- Weight decay schedule: linear -> cosine decay
- Polar express norm factor 1.02 -> 1.01
Architecture & init changes:
- VE gate: channels 32 -> 12, scale range 2x -> 3x, init small positive
- Add post-QK-norm scaling (q,k *= 1.15) for sharper attention
- Embedding init std 1.0 -> 0.8, MLP c_fc init 0.5x smaller
- RoPE base theta 10K -> 100K
- Short attention window: seq_len/2 -> ~seq_len/3 (ceil to 128 tile)
- Logit softcap 20 -> 15
2026-03-09 20:45:17 +00:00
1076f97059
delete autocast, an unnecessary thorn in my side, manage dtypes directly
2026-03-04 23:55:30 +00:00
aba30cb037
tune logit softcap?
2026-03-03 00:38:53 +00:00
2e17723817
Fix generate() crash when top_k=0 ( #467 )
...
Prevent a crash in generate() by skipping top-k filtering when top_k is set to 0
2026-01-30 09:21:02 -08:00
41bb2eac32
Combine AdamW and Muon into single MuonAdamW optimizer, cleaner, ty @chrisjmccormick for idea/help
2026-01-29 00:52:08 +00:00
74554be3b5
revert engram, not seeing an improvement at larger scale
2026-01-28 20:07:39 +00:00
c8d93beed2
add engram-lite, add log, tune scaling laws analysis scripts
2026-01-27 22:31:17 +00:00
59e36cc727
first version of engram following modded nanogpt style
2026-01-25 18:59:51 +00:00
85b3e95e09
320 experiments just to tune the adam beta1 of x0 a little bit up from 0.8 to 0.96
2026-01-25 00:04:02 +00:00
e7ed2082b8
update the default GPTConfig kwargs otherwise they are confusing
2026-01-17 21:16:46 +00:00
e85db6b4a4
alternating design
2026-01-16 23:52:12 +00:00
9a88194c3f
simply one VE per layer, works best
2026-01-16 22:08:52 +00:00
0b58d70e99
full ve version works very well
2026-01-16 21:16:47 +00:00
e3f58b838e
ranked version
2026-01-16 20:59:42 +00:00
8203efa919
implement flash attention 3 fallback to pytorch sdpa by touching as few lines of code as possible in main files and keeping all implementation to a single file. add tests. add helpful warning messages for the user.
2026-01-16 17:37:51 +00:00
fbc1484e8c
add alternating window size patterns for the GPT layers, following GPT-3. Experimented a bit and found the pattern SSSL to work well - 3 short, 1 long alternating. This is now the new default and the plots look quite a bit better on flops vs. bpb
2026-01-11 21:49:54 +00:00
2ff7d51252
integrate Flash Attention 3. +9% tok_per_sec for d12 with ctx even as low as 2048 out of the box nice. also, ready to tune windows huge
2026-01-11 20:33:19 +00:00
aa530cdad5
Add learnable lambdas that gate the residual connection and a skip connection to the input embeddings, solid bump to val_bpb
2026-01-11 18:47:35 +00:00
2c4473dd1b
Big Muon optimizer changes inspired by latest of modded-nanogpt. Added Polar Express, Adafactor-style variance reduction, cautious weight decay, schedule weight decay linearly to ramp down to zero. Tuned optimum weight decay for multiple model sizes d8, d12, d16, d20 and found a scaling law with optimum wd \propto 1/channels^2, including it as default into code. --weight_decay of base_train is now default on and configured optimally according to all of these experiments. Solid bump to val_bpb observed as a result of these changes.
2026-01-11 16:56:59 +00:00
ccf4b7f9bf
nudge hyperparameters of the base script with the results of the sweeps and miniseries. vocab size down to 32K. D:N ratio from 20 to 8. add miniseries script
2026-01-07 22:11:59 +00:00
48abd7d85f
simplify, clarify and slightly tune model initialization. should be very slightly better possibly, but certainly a lot clearer
2026-01-01 21:15:09 +00:00
e1770a3061
remove spurious cast, gets compiled away anyway but it's confusing people
2025-12-27 23:07:48 +00:00
d314e96aa2
formatting
2025-12-09 12:48:46 +01:00
f1bf69d562
feat: pad vocab size to 64 for DDP optimizers and efficiency
2025-12-09 12:38:18 +01:00
849d95ae1f
remove unnecessary check to make the logic in CausalSelfAttention.forward() clearer
2025-12-08 18:30:37 -08:00
bffdb2ef91
group common code to make things neater in gpt logit computation
2025-12-09 02:01:05 +00:00
16788eed3c
fix(model): apply float32 cast before logits softcapping
...
This change ensures that the logits softcapping operation (tanh) is performed in float32 precision rather than bfloat16. Previously, the code cast to float32 after the tanh operation, which meant the non-linearity was computed with bfloat16 precision
2025-11-23 20:12:09 +05:30
5c93a56be5
remove unnecessary check
2025-11-19 16:31:41 -05:00
11e68bf442
Fix comment: rotary embeddings final dimension size
2025-11-17 11:32:56 -05:00
bc1fca39f3
mqa -> gqa to reduce confusion
2025-11-15 15:43:37 +00:00
a088b7a6ec
use enable_gqa of pytorch sdpa, allows us to delete some code, didnt realize it's available
2025-10-21 18:07:33 +00:00
306bc380ab
add support for CPU and for MPS. I had to change a few cosmetic things. I also discovered I think a bit of a bug, where I was casting wte to bfloat16 in the wrong place (the model init) instead of in init_weights
2025-10-16 10:04:43 -07:00
3a5e0bc50b
initial commit
2025-10-13 06:49:24 -07:00
Sofie Van Landeghem