Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

AI summary

Overview

Research area: Large-scale neural network training — hyperparameter optimization and scaling laws for Mixture-of-Experts (MoE) language models.

Technical level: Advanced. The paper assumes familiarity with Mixture-of-Experts routing, Maximal Update Parameterization (μP), the Muon optimizer, Multi-head Latent Attention (MLA), and scaling-law methodology.

One-sentence scope: The paper proposes a two-step framework that predicts the optimal learning rate for trillion-token-scale MoE pretraining by transferring it across model widths via μP and extrapolating it across token budgets via a linear scaling law, then validates the prediction by pretraining a 155B-total / 17B-active-parameter model on 10 trillion tokens.

What This Paper Is About

Training large MoE models requires choosing a learning rate that is highly sensitive to both model size and token budget, so every change in scale normally demands a new and extremely expensive hyperparameter sweep. The authors target the problem that existing transfer frameworks such as μP and μ-Transfer were designed for dense models under width scaling, leaving it unclear whether they generalize to MoE, where sparsity (the ratio of active to total experts) is a second scaling axis. Their goal is to estimate the optimal learning rate for a very large, very long-horizon MoE run using only a handful of small proxy runs.

Key Contributions

1. μP adaptation for MoE architectures: The authors formulate a μP adaptation for MoE models that use Multi-head Latent Attention (MLA) and the Muon optimizer, covering zero-shot hyperparameter transfer when scaling both model width and the total number of experts (sparsity expansion), and extending the empirical scope of μP-based MoE scaling.
2. Token-scale extrapolation in a two-step framework: They introduce a two-step predictive framework that extrapolates optimal learning rates to unseen, long token horizons (e.g., 10 trillion tokens) from cost-effective small-scale proxy runs, substantially reducing hyperparameter tuning overhead.
3. Large-scale validation: They apply the framework to pretrain a 155B-total / 17B-active-parameter MoE foundation model from scratch over 10 trillion tokens, reporting stable training and competitive evaluations.
4. Methodological choices for proxy efficiency: They terminate proxy runs during the stable phase of the Warmup-Stable-Decay scheduler without decay and apply Exponential Moving Average (EMA) to weights, so multiple checkpoints at desired token intervals can be extracted from a single run rather than requiring one run per token budget.

Main Findings

• μP enables learning rate transfer across MoE width. Under μP, the optimal learning rate identified in the base proxy MoE (0.6B total, 0.3B active) transferred consistently to models scaled to 2x (2.2B total, 0.7B active), 4x (8B total, 1.5B active), and 8x (30.7B total, 3.6B active) the base width. Under Standard Parameterization (SP), the optimal learning rate shifted with width and failed to transfer.
• Transfer also holds for dense MLA models. In the dense baseline, μP preserved the optimal learning rate from a 0.24B base proxy across 2x (0.7B), 4x (2.27B), and 8x (8.02B) width scales, while SP diverged.
• Proxy and target curves agree at each token scale. At 40B, 60B, 80B, and 100B tokens, the quadratic fits of validation loss versus log learning rate for the 5.6B-total / 1.8B-active proxy and the 2x width-scaled 20.7B-total / 3.8B-active held-out model showed highly consistent curvature and vertex locations. Validation losses at a learning rate of 2×10⁻³ fell directly on the quadratic curves fitted using the remaining four learning rates.
• Optimal learning rate drifts downward as the token budget grows. The paper reports a slight downward trend in the optimal learning rate as the token budget increases.
• Log-log linear regression extrapolates to 10T tokens with R² = 0.95. Fitting optimal learning rates estimated across token scales from 255B to 502B (using the 10.8B-total / 3.3B-active proxy trained for approximately 500B tokens) produced a linear fit with R² = 0.95, predicting an optimal learning rate of 3.85×10⁻⁴ for 10T-token pretraining.
• Compute savings versus 2D sweeps. Model-scale search at 1.5x and 2x wider than the proxy would incur an additional 240.3 ZFLOPs beyond the proxy run cost of 64.8 ZFLOPs. The full-scale target pretraining requires approximately 98x the total compute of the proxy runs used for learning rate prediction.
• Stable full-scale training. The 155B-total / 17B-active model trained on 10 trillion tokens showed a stable pretraining loss with no loss spikes, which the authors present as evidence that the extrapolated learning rate is valid.
• Pareto-frontier comparison. Using training compute estimated as 6ND with N set to active parameters, the model achieved higher MMLU-Pro accuracy than dots.llm1 and GLM-4.5-Air at comparable or lower estimated training compute. The numeric benchmark scores themselves are presented in Figure 7 and are not stated as numbers in the paper text.
• Batch size is deliberately excluded from transfer. The authors fix batch size to maximize hardware throughput and treat the learning rate as the primary transfer target, citing contradictory findings in the literature about how optimal batch size scales.

Methodology in Plain English

The authors split the problem of choosing a learning rate into two independent steps.

Step 1 — Remove the model-size axis with μP. Instead of sweeping learning rates separately for every model size, they organize model parameters by shape. Parameters with exactly one expandable dimension (embeddings, biases, expert FC2 weights) are "vector-like" and get μP initialization only. Parameters with two expandable dimensions (FFN, attention, router, and expert FC1 weights) are "matrix-like" and get both μP initialization and learning rate scaling, with the scaling factor set to fan_in_base / fan_in. Only matrix-like hidden weights are learning-rate-scaled, following Wortsman et al. (2024). Depth is held fixed because depth scaling is known to be unstable for μP transfer, and the head dimension is held fixed while the number of heads scales with hidden size. When scaling up, they fix the number of active experts and the expert intermediate dimension while increasing the total expert count and hidden size, so model scale, active scale, and sparsity move together rather than independently. They validate this by sweeping learning rates on MoE models at 1x, 2x, 4x, and 8x width and checking whether the same learning rate wins each time.

Step 2 — Extrapolate along the token axis. Running a learning rate sweep for every possible token budget would be expensive, and decaying the learning rate early during proxy runs risks biasing loss estimates. So proxy runs are stopped during the stable phase of the WSD schedule and EMA is applied to weights, which lets them pull out many checkpoints at different token intervals from one run. EMA is updated as θ_EMA⁽ᵗ⁾ = α·θ_EMA⁽ᵗ⁻¹⁾ + (1−α)·θ⁽ᵗ⁾ with α = 0.6. Checkpoints are EMA-updated roughly every 2B tokens, and every-10B-token checkpoints are used for analysis. For each token budget, they fit a second-order polynomial of validation loss against log learning rate, take the vertex as the optimal learning rate, and then run a linear regression of log(optimal learning rate) against log(token budget). They restrict this regression to points after a batch-size increase has stabilized the dynamics.

Step 3 — Apply and check. They predict the optimal learning rate for a 10-trillion-token run of a 155B-total / 17B-active MoE model and pretrain it from scratch. Data mixture starts at 45% English, 12.5% Math/STEM, 27.5% Code, and 15.0% Multilingual, then shifts at the 6T-token mark to 22.5% English, 27.5% Math/STEM, 25.0% Code, and 25.0% Multilingual. Batch size scheduling is applied after 200B tokens. All experiments run on NVIDIA H200 GPUs using an internal fork of Megatron-LM with the Muon optimizer.

Why This Matters

The paper's central claim is that the optimal learning rate for a 10-trillion-token, 155B-parameter MoE pretraining run can be predicted from small proxy experiments rather than found by exhaustive sweeps — a procedure the authors state is computationally infeasible at full scale. If the approach generalizes, it lowers the cost of hyperparameter decisions that sit on the critical path of frontier-scale training.

Impact on research: The work extends μP-based transfer beyond the width-only, dense-model settings it was designed for, into MoE architectures that combine MLA attention, Muon optimization, and sparsity expansion. It also connects hyperparameter transfer to token-dimension scaling laws, which prior MoE scaling research (Clark et al., 2022; Ludziejewski et al., 2024) largely treated as a search over architectural configurations rather than a transfer problem. The paper notes that prior MoE μP work (Małaśnicki et al., 2025) was limited to AdamW and to fixing both total and active expert counts while scaling only hidden dimensions, which does not generalize to MoE designs with large numbers of fine-grained experts.

Real-world applications:

• Reducing the compute budget required to configure large-scale MoE pretraining runs.
• Guiding training-recipe decisions for long-horizon runs where trial-and-error sweeps are impossible.
• Informing efficient scaling of model capacity through expert count, following the practical sparsity route many open MoE models now take.
• Providing a reusable transfer protocol for teams training MoE models with MLA attention and the Muon optimizer.

Industry relevance: The concrete validation target is a production-scale foundation model (155B total, 17B active parameters) trained on 10 trillion tokens at Kakao Corp., making the results directly relevant to organizations that pretrain MoE models and must justify large-scale compute allocations before a run begins.

Future Directions

• Broadening architecture and optimizer coverage. The framework is demonstrated only for MoE architectures with MLA and the Muon optimizer; the authors state that extending it to other MoE structures and optimizers is important future work.
• Per-expert learning rate adaptation. Because top-k routing assigns varying numbers of tokens to individual experts, the effective batch size, gradient noise scale, and potentially the optimal update scale may differ per expert. The authors flag this as a potential source of gains but leave it to future work given the experimental and engineering cost.
• Disentangling sparsity from width. The large-scale recipe expands sparsity jointly with width, so the effect of the sparsity dimension alone cannot be separated from that of width. A controlled large-scale study of μP transfer along the sparsity axis alone would be needed.
• Direct verification of the predicted learning rate at full scale. The authors acknowledge that exhaustively sweeping at full scale to confirm the optimality of 3.85×10⁻⁴ is computationally infeasible, so validation rests on the stable loss trajectory and benchmark scores rather than a full-scale optimality check. The truncated Appendix E begins a small-scale validation of the extrapolation, but its details are not included in the available content.

Target Audience

This paper is most useful to machine learning engineers and researchers who pretrain or fine-tune large MoE language models and need to make hyperparameter decisions under tight compute budgets. It is also relevant to researchers working on parameterization theory and hyperparameter transfer (μP, μ-Transfer), on scaling laws for training token budgets, and on MoE architecture design. Readers will get the most value if they already understand μP, MoE routing, and the Muon optimizer; the methodological details assume that background. Those looking for a beginner-level introduction to MoE training or for absolute benchmark score tables will not find what they need here — the evaluation results are presented as figures rather than as numbers in the text.

Authors’ abstract

Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization ($μ$P) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity ($R^2=0.95$). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.