Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts
TL;DR - This paper proposes a two-step method for transferring optimal learning rates from small proxy models to large Mixture-of-Experts models and trillion-token training runs. It could substantially reduce the cost of hyperparameter sweeps for large-scale pretraining.
- Adapts Maximal Update Parameterization (ÎĽP) to MoE architectures using Multi-head Latent Attention and the Muon optimizer.
- Demonstrates consistent learning-rate transfer across models of different widths.
- Fits a token-budget scaling law that predicts optimal learning rates up to 10 trillion tokens with (R^2=0.95).
- Validates the method by stably pretraining a 155B-parameter MoE model with 17B active parameters.