Validation loss 3.28 with AdamW in less than 10k iterations (local bs=64)

[svaiter]

Fix partially issue #5 (bullet point Reaching 3.28).

Validation loss on 10B FineWeb tokens was around 3.6–3.8 after 8k iterations on 8×H100 (2 nodes on jzay) but classical target is ≤3.29 due to the issue 481 of llm.c project by A. Kaparthy

Main cause is that the model in this repo was a mix between classic GPT-2 (learned positional embeddings, LayerNorm, no attention scaling) and new stuff like sinusoidal init. Taking modded-nanogpt commit 844e5fd a bit more rigorously leads to 3.28 loss after ~9k iterations with AdamW, like the run 2 of modded-nanogpt.

At the moment, only the behavior of theses changes on AdamW was tested.

The commits in this PR were prepared using Claude Opus 4.7.

Main changes

• Removed the generate code in benchmark_utils/model_gpt2.py(never used for this benchmark).

Architecture changes

Modified file: benchmark_utils/model_gpt2.py

• RoPE rotary embeddings applied to Q/K in attention, replacing the learned
  positional embedding table (wpe removed).
• RMSNorm (parameter-free) before attention and MLP and at the output,
  replacing LayerNorm / ln_f.
• Attention residual scale 1/sqrt(2*n_layer) applied to the attention
  output before the residual add.
• Default PyTorch init instead of the custom scaled normal init. The
  init_func hook is kept so the sinusoidal-init experiment still works.
• vocab_size 50257 (real GPT-2 vocab) instead of 50304.

Hyperparameters

Modified files: solvers/adam.py, benchmark_utils/lr_scheduler.py

• learning_rate=1.8e-3, betas=(0.9, 0.98), weight_decay=0.1
• num_steps=9536, warmup_iters=256, warmdown_iters=2048
• New trapezoidal LR schedule get_lr_trapezoidal (linear warmup → plateau →
  linear warmdown). The existing cosine-like get_lr is retained for the
  Muon/SOAP/Scion solvers.

[svaiter]

Default parameters of Muon needs to be fixed. I propose that we still work based on this architecture, and then looks at further modded-nanogpt commits to see what are the good hyperparameters for Muon.