Please run data.py in the trainers folder first in the same directory where the training code lives.
Sample usage:
python train_muon.py --experiment_name muon
For logging:
python logger.py --experiment_name muon
Sample for profiling:
python train_muon.py --experiment_name muon --enable_profiling
This repo is an attempt to port all changes of nanogpt modded by Keller Jordan to MoE(Mixture of Experts) architecture
GitHub: Nanogpt-modded
Doc listing all changes: Nanogpt Modded Change Log
Architecture details:
- Mixtral-style Mixture of Experts (MoE)
- Top-2 experts chosen
- Noisy Top-K router to avoid expert bias
PRs are welcome.
Checklist of all completed changes:
- Tuned learning rate & rotary embeddings
- Introduced the Muon optimizer
- Muon improvements
- Pad embeddings, ReLU², zero-init projections, QK-norm
- Distributed the overhead of Muon
- Untied embedding and head
- Value and embedding skip connections, momentum warmup, logit softcap
- Bfloat16 activations
- U-net pattern skip connections & double lr
- 1024-ctx dense causal attention → 64K-ctx FlexAttention ( post this change due to some bug in flex attention for B,S,H values, all of the following scripts have a torch compile bug<<working to solve :)>>)
- Attention window warmup
- Value Embeddings
- U-net pattern value embeddings, assorted code optimizations
- Split value embeddings, block sliding window, separate block mask
- Sparsify value embeddings, improve rotary embeddings, drop an attn layer
- Lower logit softcap from 30 to 15
- FP8 head, offset logits, lr decay to 0.1 instead of 0.0
- Merged QKV weights, long-short attention, attention scale, lower Adam epsilon, batched Muon
- Reduced batch size
- Faster gradient all reduce
- Overlap computation and gradient communication