Skip to content

README.md

Latest commit

 

History

History
379 lines (273 loc) · 7.97 KB

README.md

File metadata and controls

379 lines (273 loc) · 7.97 KB
Raw

minigpt-muon-benchmark

Single-GPU from-scratch GPT pretraining experiments comparing AdamW against Hybrid Muon.

This repository is both:

  1. a learning project for understanding LLM pretraining from scratch, and
  2. a small controlled benchmark of AdamW vs Hybrid Muon on MiniGPT models.

The benchmark is intentionally modest: one GPU, public datasets, small GPT-style models, seed-matched runs, and validation-loss comparisons. It is not a frontier-model reproduction.

What this repo does

This repo trains small decoder-only GPT models from random initialization using next-token prediction.

It compares:

AdamW
vs
Hybrid Muon

Hybrid Muon means:

Muon for hidden 2D matrix weights:
  - attention q/k/v/o matrices
  - MLP matrices

AdamW fallback for everything else:
  - token embeddings
  - position embeddings
  - LM head
  - LayerNorm weights
  - biases
  - vectors/scalars
  - non-2D parameters

The main question:

At the same model, seed, dataset cache, batch setting, and token budget, does Hybrid Muon reach lower validation loss than AdamW?

Current results

57M MiniGPT

Model:

MiniGPT-Dense-60M-v1
parameters: ~57.47M
context length: 512
tokenizer: openai-community/gpt2
precision: bf16

Training setup:

micro_batch_size: 16
grad_accum_steps: 8
effective_batch_tokens: 65,536
steps per run: 4,000
tokens per run: 262,144,000
seeds: 1337, 2024, 2025

Aggregate results:

Dataset Steps Seeds AdamW mean ± std Hybrid Muon mean ± std Mean delta Relative improvement
TinyStories 4000 3 1.4658 ± 0.0028 1.4116 ± 0.0016 0.0542 3.69%
FineWeb-Edu 500M cache 4000 3 3.9747 ± 0.0150 3.8645 ± 0.0144 0.1103 2.77%
FineWeb-Edu 500M cache 8000 1 3.7377 3.6614 0.0763 2.04%

Positive delta means AdamW validation loss minus Muon validation loss. Positive values mean Muon was lower.

125M MiniGPT

Model:

MiniGPT-Dense-125M-v1
parameters: 123,963,648
context length: 512
tokenizer: openai-community/gpt2
precision: bf16

Dataset:

source: HuggingFaceFW/fineweb-edu
config: sample-10BT
train cache: 500,000,000 tokens
validation cache: 5,000,000 tokens

Aggregate result over 3 seed-matched runs:

Metric AdamW Hybrid Muon
Best validation loss mean ± std 3.8056 ± 0.0146 3.7049 ± 0.0181
Mean validation-loss delta n/a 0.1007 ± 0.0038
Mean relative improvement n/a 2.65%

Main observation so far:

In these small single-GPU experiments, Hybrid Muon consistently reached lower validation loss than AdamW at the same token budget. AdamW remained faster in raw tokens/sec.

What this does not prove

This repo does not prove that Muon is universally better than AdamW.

Limitations:

  • models are small: ~57M and ~124M parameters
  • runs are single-GPU
  • context length is only 512 tokens
  • FineWeb-Edu uses a 500M-token extracted cache, not the full dataset
  • the implementation is educational/research-oriented, not production training infrastructure

A defensible claim is:

On these seed-matched MiniGPT pretraining benchmarks, Hybrid Muon achieved lower validation loss than AdamW at the same model size, dataset cache, and token budget, while AdamW was faster in raw throughput.

Repository structure

.
├── gpu_benchmark/
│   ├── train_gpt.py                  # main pretraining script
│   ├── batch_sweep.py                # batch-size/memory sweep helper
│   ├── compare_optimizer_runs.py     # comparison/plotting helper
│   ├── README.md                     # GPU benchmark usage notes
│   ├── notebooks/                    # notebook runners for EC2 experiments
│   └── downloaded_runs/              # downloaded summaries, CSVs, plots
│
├── visual_labs/
│   ├── PRETRAINING_HANDHOLDING_GUIDE.md
│   ├── REAL_WORLD_ANALOGIES.md
│   ├── PROJECT_ROADMAP.md
│   ├── scripts/                      # VS Code / Jupyter-style # %% labs
│   └── notebooks/                    # notebook versions of the visual labs
│
├── reports/
│   ├── 57m_adamw_vs_muon_results.md
│   ├── 125m_adamw_vs_muon_seed1337_results.md
│   └── 125m_adamw_vs_muon_3seed_results.md
│
├── papers/
│   └── 2502.16982v1-muon-scalable-llm-training.*
│
└── spec.md                           # training process spec

Core concepts

Pretraining

Pretraining trains a randomly initialized model to predict the next token.

Example:

input:  The capital of France is
target: Paris

Token budget

We usually measure LLM training by tokens processed, not only by epochs.

For the main runs:

context_length = 512
micro_batch_size = 16
grad_accum_steps = 8

So:

tokens_per_step = 512 × 16 × 8 = 65,536

For 4000 steps:

tokens_seen = 65,536 × 4000 = 262,144,000

Seed-matched benchmark

A seed controls randomness: model initialization, data sampling, dropout, and generation randomness.

For optimizer comparisons, seed matching matters:

AdamW seed 1337
Muon  seed 1337

This makes the comparison cleaner because both optimizers start from the same random setup.

Datasets

TinyStories

source: roneneldan/TinyStories
train tokens: 473,992,236
validation tokens: 4,765,918

Used for early pipeline validation and small-model optimizer comparison.

FineWeb-Edu

source: HuggingFaceFW/fineweb-edu
config: sample-10BT
train cache: 500,000,000 tokens
validation cache: 5,000,000 tokens

This is an extracted local token cache from the larger sample-10BT source. It is not the full FineWeb-Edu dataset.

Quick start

Local toy smoke test

python3 gpu_benchmark/train_gpt.py \
  --dataset toy \
  --data-dir /tmp/minigpt_toy_data \
  --run-dir /tmp/minigpt_toy_adamw \
  --optimizer adamw \
  --model-config tiny_debug \
  --precision bf16 \
  --micro-batch-size 2 \
  --grad-accum-steps 2 \
  --max-steps 2 \
  --max-minutes 0 \
  --eval-every-steps 1 \
  --save-every-steps 1 \
  --validation-tokens 4096

Muon toy smoke test

python3 gpu_benchmark/train_gpt.py \
  --dataset toy \
  --data-dir /tmp/minigpt_toy_data \
  --run-dir /tmp/minigpt_toy_muon \
  --optimizer muon \
  --model-config tiny_debug \
  --precision bf16 \
  --micro-batch-size 2 \
  --grad-accum-steps 2 \
  --max-steps 2 \
  --max-minutes 0 \
  --eval-every-steps 1 \
  --save-every-steps 1 \
  --validation-tokens 4096

Safe long-running training

The training script supports:

  • checkpoint saving
  • resume from checkpoint
  • STOP-file safe exit
  • metrics JSONL logging
  • periodic validation
  • sample generation
  • plot generation

To stop a run safely:

touch <run-dir>/STOP

The script should finish the current safe point, save the latest checkpoint, write summaries, and exit cleanly.

Learning materials

If you are learning pretraining from scratch, start here:

visual_labs/PRETRAINING_HANDHOLDING_GUIDE.md

Then read:

visual_labs/REAL_WORLD_ANALOGIES.md
visual_labs/PROJECT_ROADMAP.md
visual_labs/notebooks/
visual_labs/scripts/

The visual labs cover:

  • embeddings
  • attention projections
  • MLPs
  • LM heads
  • KV cache intuition
  • MoE routing
  • AdamW vs Muon toy experiments
  • tiny GPT training

Reports

Main result reports:

reports/57m_adamw_vs_muon_results.md
reports/125m_adamw_vs_muon_3seed_results.md

Generated artifacts include CSVs and plots under:

gpu_benchmark/downloaded_runs/

Recommended repo name

Recommended GitHub repository name:

minigpt-muon-benchmark

Reason:

  • accurate: it is a MiniGPT benchmark
  • specific: it is about Muon
  • not overclaiming: it does not imply frontier-scale reproduction
  • searchable: includes the relevant technical keywords

Alternative names:

single-gpu-muon-pretraining
adamw-vs-muon-minigpt
muon-pretraining-lab