Parameter Golf
Notes and logs from attempting the parameter golf challenge — squeezing BPB out of a tiny transformer under 16MB.
notes and logs from my attempts at actually attempting this https://github.com/sijirama/parameter-golf/tree/main
notes
- Depth recurrence reusing the same layer weights multiple times. Instead of 9 different layers each with unique weights, we have 3 layers but we run them 3 times each, same weights, multiple passes.
logs
- 20/mar 1:50am
- as of when this section was written the current baseline (which is the default code in the train_gpt file ) has a score of 1.2244, this ran for the normal 10 minutes
- the non record score OF THE SAME ARCHITECTURE had a score of 1.2074, this ran for 4 HOURS
- this means that the bottleneck here is not the compute, it's the ARCHITECTURE
- their
9layer 512dim 1024vocab TiedEmbeddings 4 KV headsdoesn't work well, there's room for improvemnt - ### baseline stats (before any changes)
| metric | value |
|---|---|
| architecture | 9 layers, 512 dim, 1024 vocab |
| unique blocks | 9 (no weight tying) |
| total params | 17,059,912 |
| emb params | 524,288 |
| block params | 16,533,576 |
| skip params | 2,048 |
| int8 size | 17.06 MB |
| estimated compressed (conservative) | 6.82 MB |
| estimated compressed (optimistic) | 4.87 MB |
| actual compressed (from real run) | 15.82 MB |
| actual compression ratio | 3.91x |
| val BPB (10 min run) | 1.2244 |
| val BPB (4 hour run) | 1.2074 |
| KV heads | 4 |
| tied embeddings | yes |
- > plan 1
- Increase layers, increase width - with depth recurrence we can go deeper and wider
- depth recurrence - implement Albert-style weight tying across more layers - https://arxiv.org/pdf/1909.11942
- 4 unique layers, 1 skip weight for each decoder 2 passes, 8 effective layers
| metric | recurrent v2 (our run) |
|---|---|
| architecture | 4 unique blocks, 768 dim, 8 effective layers |
| unique blocks | 4 (recurrent U-Net) |
| total params | 17,316,896 |
| emb params | 786,432 |
| block params | 16,527,392 |
| skip params | 3,072 |
| int8 size | 17.32 MB |
| compressed size | 12.45 MB |
| compression ratio | 3.89x |
| step time | 65ms |
| total steps | 9,227 |
| final val BPB | 1.5317 |
| final val loss | 2.5862 |
| sequence length | 1024 |
| num heads | 8, kv heads 4 |
| status | beats baseline per step, loses on total steps |
terrible really, didn't even catch up to baseline, we can do better
- > plan 2
- 2048 instead of 1024 sequence length
- scaling window eval
- 0.04 weigth decay on Muon - will help with quantization
- int6 quantization currently only int8
| metric | recurrent v4 (512 dim) |
|---|---|
| architecture | 4 unique blocks, 512 dim, 8 effective layers |
| unique blocks | 4 (recurrent U-Net) |
| total params | 7,874,592 |
| emb params | 524,288 |
| block params | 7,348,256 |
| skip params | 2,048 |
| int8 size | 7.87 MB |
| compressed estimate | 3.15 MB |
| step time | 39ms |
| projected total steps | ~15,384 |
| val bpb step 1000 | 1.4294 |
| val bpb step 2000 | interrupted |
| sequence length | 1024 |
| eval method | sliding window |
| muon weight decay | 0.04 |
| status | pod died at step 2000 due to insufficient balance |
Summary text:
v4 was our best run to date. Dropping model_dim from 768 to 512 was the key insight, it brought step time down to 39ms, giving us a projected 15,000+ steps in 10 minutes compared to 9,200 in our best previous run. The smaller model also learned faster, hitting 1.4294 BPB at step 1000, our best early trajectory across all runs. For context, v2 hit 1.8952 at step 1000. We also added sliding window evaluation and Muon weight decay in this run. The architecture is very interesting, recurrent U-Net transformer with shared weights across encoder and decoder passes. The run was cut off at step 2000 when Runpod balance ran out. This config is confirmed as the right direction, just needs a clean uninterrupted run to get a final score.
v5
| metric | recurrent v5 (final) |
|---|---|
| architecture | 4 unique blocks, 512 dim, 8 effective layers, recurrent U-Net |
| total params | 7,874,592 |
| emb params | 524,288 |
| block params | 7,348,256 |
| skip params | 2,048 |
| int8 size | 7.87 MB |
| compressed size | 5.94 MB |
| compression ratio | 3.84x |
| step time | 38.96ms |
| total steps | 15,403 |
| sequence length | 1024 |
| eval method | sliding window |
| muon weight decay | 0.04 |
| final val BPB | 1.2750 |
| final val loss | 2.1527 |
| baseline BPB | 1.2244 |
| gap to baseline | +0.0506 |
| compressed submission | 5.99 MB (well under 16MB) |
| status | complete, clean run |
step:15400/20000 train_loss:2.1422 train_time:599916ms step_avg:38.96ms
step:15403/20000 val_loss:2.1378 val_bpb:1.2662 train_time:600051ms step_avg:38.96ms
stopping_early: wallclock_cap train_time:600051ms step:15403/20000
peak memory allocated: 8975 MiB reserved: 9124 MiB
Serialized model: 30465637 bytes
Code size: 51057 bytes
Total submission size: 30516694 bytes
Serialized model int8+zlib: 5937084 bytes (payload:7932032 raw_torch:7952474 payload_ratio:3.84x)
Total submission size int8+zlib: 5988141 bytes
final_int8_zlib_roundtrip val_loss:2.1527 val_bpb:1.2750 eval_time:36110ms
final_int8_zlib_roundtrip_exact val_loss:2.15271180 val_bpb:1.27501508v6
| metric | recurrent v6 |
|---|---|
| architecture | 6 unique blocks, 576 dim, 12 effective layers, recurrent U-Net |
| total params | 14,541,744 |
| emb params | 589,824 |
| block params | 13,948,464 |
| skip params | 3,456 |
| int8 size | 14.54 MB |
| compressed size | 10.79 MB |
| compression ratio | 3.89x |
| step time | 95.88ms |
| total steps | 6,245 |
| sequence length | 1024 |
| eval method | sliding window |
| muon weight decay | 0.04 |
| final training BPB | 1.2216 |
| final roundtrip BPB | 1.2302 |
| baseline BPB | 1.2244 |
| gap to baseline | +0.0058 |
| compressed submission | 10.79 MB (under 16MB) |
| status | complete, clean run |
Summary:
v6 was the most promising recurrent run yet and nearly beat the baseline. Scaling from 4 to 6 unique blocks and 512 to 576 dim gave significantly better learning per step — hitting 1.2216 BPB at step 6,245 compared to v5 which hit 1.2662 at step 15,403. The bigger model learned far more efficiently per step. The problem is the step time jumped to 95ms giving only 6,245 total steps vs v5's 15,403. The model was still dropping fast at the wallclock cap — the warmdown alone dropped BPB from ~1.31 to 1.22 in the final 1,245 steps. We're only 0.0058 from beating the baseline with half the steps of a standard transformer. The architecture is working — the bottleneck is purely step speed. The next experiment should find the sweet spot between model capacity and step speed.
step:5000/20000 val_loss:2.2144 val_bpb:1.3115 train_time:481824ms step_avg:96.36ms
step:5200/20000 train_loss:2.3823 train_time:500813ms step_avg:96.31ms
step:5400/20000 train_loss:2.3145 train_time:519869ms step_avg:96.27ms
step:5600/20000 train_loss:2.1966 train_time:538450ms step_avg:96.15ms
step:5800/20000 train_loss:2.3571 train_time:557304ms step_avg:96.09ms
step:6000/20000 train_loss:2.0767 train_time:577040ms step_avg:96.17ms
step:6000/20000 val_loss:2.0910 val_bpb:1.2384 train_time:577093ms step_avg:96.18ms
step:6200/20000 train_loss:2.1416 train_time:595483ms step_avg:96.05ms
step:6245/20000 val_loss:2.0626 val_bpb:1.2216 train_time:598937ms step_avg:95.91ms
stopping_early: wallclock_cap train_time:598937ms step:6245/20000
peak memory allocated: 14976 MiB reserved: 15398 MiB
Serialized model: 57010425 bytes
Code size: 51057 bytes
Total submission size: 57061482 bytes
Serialized model int8+zlib: 10737613 bytes (payload:14637248 raw_torch:14667514 payload_ratio:3.89x)
Total submission size int8+zlib: 10788670 bytes
final_int8_zlib_roundtrip val_loss:2.0771 val_bpb:1.2302 eval_time:45838ms
final_int8_zlib_roundtrip_exact val_loss:2.07711513 val_bpb:1.23024045v7
| metric | recurrent v7 |
|---|---|
| architecture | 6 unique blocks, 576 dim, 12 effective layers, recurrent U-Net |
| total params | 14,541,744 |
| emb params | 589,824 |
| block params | 13,948,464 |
| skip params | 3,456 |
| int8 size | 14.54 MB |
| compressed size | 10.44 MB |
| compression ratio | 3.89x |
| step time | 75.75ms |
| total steps | 7,922 |
| sequence length | 1024 |
| eval stride | 256 |
| val loss every | 3000 steps |
| matrix lr | 0.02 |
| scalar lr | 0.02 |
| tied embed lr | 0.03 |
| muon weight decay | 0.04 |
| final training BPB | 1.2128 |
| final roundtrip BPB | 1.2171 |
| baseline BPB | 1.2244 |
| gap to baseline | -0.0073 (BEATS BASELINE) |
| compressed submission | 10.44 MB (under 16MB) |
| status | complete, clean run, first submission candidate |
summary:
v7 is the first run to beat the openai baseline. the key changes from v6 were lower learning rates across the board (matrix and scalar lr down to 0.02, embed lr down to 0.03) and a better eval stride of 256 instead of 512. the lower lr made a big difference on the quantization side, the penalty dropped from 0.007 in earlier runs to just 0.0043 here, meaning the model survives int8 compression much better.
the model was still learning fast at the wallclock cap, hitting 1.2128 at step 7922 before the 10 minutes ran out. the warmdown did its usual thing and the final roundtrip came in at 1.2171. we have 5.56mb of compressed budget still unused which means there is room to scale up or try new things.
the main bottleneck right now is step time. 75ms steps only gets us 7922 steps in 10 minutes. eval is also expensive, the final eval alone took 99 seconds. tomorrow the plan is to look at sliding window attention to reduce the quadratic attention cost, which should make steps faster and also make eval cheaper, giving us more training steps and potentially a much better final score.
step:7922/20000 val_loss:2.0476 val_bpb:1.2128 train_time:600089ms step_avg:75.75ms
stopping_early: wallclock_cap train_time:600089ms step:7922/20000
peak memory allocated: 14976 MiB reserved: 15398 MiB
Serialized model: 57010425 bytes
Code size: 51053 bytes
Total submission size: 57061478 bytes
Serialized model int8+zlib: 10384540 bytes (payload:14637248 raw_torch:14667514 payload_ratio:3.89x)
Total submission size int8+zlib: 10435593 bytes
final_int8_zlib_roundtrip val_loss:2.0549 val_bpb:1.2171 eval_time:99707ms
final_int8_zlib_roundtrip_exact val_loss:2.05494747 val_bpb:1.21711110