Summary (Overview)

• This paper proposes a compute-efficient two-step hyperparameter transfer framework for large-scale Mixture-of-Experts (MoE) pretraining, specifically targeting the optimal learning rate (LR) without exhaustive sweeps. .
• Step 1: Formulates a Maximal Update Parameterization (µP) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal LRs transfer consistently across width-scaled models (from 0.6B to 30.7B total parameters).
• Step 2: Establishes a token-scaling law—a linear regression in log-log space between the estimated optimal LR and token budget (achieving $R^2 = 0.95$)—that enables extrapolation to trillion-token horizons (e.g.,, 10T tokens) from small proxy runs. .
• The methodology is validated by pretraining a 155B total, 17B active parameter MoE foundation model over 10T tokens from scratch, achieving stable training dynamics (no loss spikes) and competitive benchmark scores (lying on the Pareto frontier vs. comparable open-weight models).
• The framework decouples model-scale sweeps (via µP transferability) from token-scale sweeps (via the scaling law), replacing prohibitive 2D sweeps (Figure 1) with a lightweight 1D search along the token dimension.

Introduction and Theoretical Foundation

Background: MoE architectures (Shazeer et al.,, 2017) expand model capacity without proportional compute increases, enabling stronger performance than dense counterparts at comparable inference costs. However, they introduce additional hyperparameters (e.g.,, expert routing, load balancing), making tuning more complex和expensive. The learning rate is particularly critical (Bengio, 2012; Smith, 2018; Kaplan et al.,,,2020; Yang et al.,,,2022; Li et al.,,,2025b) yet highly sensitive to both model size and token budget—requiring fresh searches at every scale change.

.

Motivation: Conventional hyperparameter scaling approaches require exhaustive 2D sweeps across both model scale (M) and token scale (D) to jointly predict the target scale (Figure 1a), incurring prohibitive compute. For MoE models exceeding 100B parameters, width scaling alone becomes intractable due to inference costs; instead, increasing the total number of experts while keeping active experts fixed (sparsity expansion) offers a practical path—but whether existing hyperparameter transfer frameworks (designed for dense width scaling) generalize to sparsity remains unclear.

.

Theoretical Foundation:

• Maximal Update Parameterization (µP) (Yang & Hu, 2021; Yang et al.,,,2022): Enables zero-shot transfer of optimal hyperparameters across model widths by scaling parameter initialization and learning rates according to tensor shape (Table 1,Table 2). The key principle is that matrix-like parameters (containing two infinitely expandable dimensions) receive both µP initialization and LR scaling, while vector-like parameters (one expandable dimension) receive initialization only.

• Spectral-condition view of µP (Yang et al.,,,2023): Scaling sparsity (increasing total experts while fixing active experts) does not change the fan-in/fan-out of individual experts beyond width scaling, thereby remaining compatible with the same µP scaling rule—coupling sparsity and width as a joint scaling axis.

.

• Token-scaling law: The optimal LR decays slightly as the token budget increases, following a power-law relationship that can be captured by linear regression in log-log space (Eq.(3)).

Methodology

µP Adaptation for MoE (Section 2.1):

• Parameter classification (Table 1): For MoE layers, router parameters and expert FC1 weights are classified as matrix-like (two expandable dimensions); expert FC2 weights are vector-like (effective input dimensionality bounded by fixed active experts and MoE intermediate dimension).
• Initialization & LR scaling (Table 2):
  • I/O Embeddings & All Biases (Vector-like): Init. Var = 0.04, LR Scaling Factor = 1
  • Hidden Weights (Vector-like: Init. Var = $fan\_in_{base}/fan\_in$, LR Scaling Factor = 1
  • Hidden Weights(Matrix-like: Init. Var = $fan\_in_{base}/fan\_in$, LR Scaling Factor = $fan\_in_{base}/fan\_in$
• Scaling strategy: Fix active experts per token and MoE intermediate dimension; increase total experts and hidden dimension proportionally. For MLA, low-rank projection dimensions for query/key-value are kept fixed during width scaling, so LR scaling on up-projection matrices reduces to 1 (no effect).
• Depth fixed: Depth scaling is known to be unstable for µP transfer; only width scaling is considered.

Estimating Optimal LRs over Short Budgets (Section 2.2.1):

• Train small proxy models over short token horizons across a range of LRs, using the Warmup-Stable-Decay (WSD) scheduler (Hu et al.,,,2024) matching the target warmup steps. .

• Terminate proxy runs during the stable phase without decay and apply Exponential Moving Average(EMA) to model weights. EMA parameters are updated as:

$$\theta_ {\mathrm{EMA}} ^ {(t)} = \alpha \cdot \theta_ {\mathrm{EMA}} ^ {(t - 1)} + (1 - \alpha) \cdot \theta^ {(t)}\tag{1}$$

with smoothing factor $\alpha = 0.6$, enabling extraction of multiple checkpoints at desired token intervals from a single run. . This mimics learning rate decay effects (Morales-Brotons et al.,,,2024; Baidu-ERNIE-Team,2025) and is particularly relevant given the large global batch size of 32M tokens. .

• For each token scale (B), fit a second-order polynomial between validation loss (L)and log-transformed LR(log η):

$$\mathcal {L} (\eta) = a (\log \eta) ^ {2} + b (\log \eta) + c\tag{2}$$

The vertex(log $\eta ^ { * } = - b / 2 a ,$ ) yields the estimated optimal LR($\eta ^ { * } = \exp ( - b / 2 a ) )$.

Extrapolating to Long Horizons (Section 2.2.2):

• Perform linear regression in log-log space between estimated optimal LR($\eta^ {*}$)and token budget(B):

$$\log (\eta^ {*}) = \beta \cdot \log (B) + \gamma\tag{3}$$

• When batch size scheduling is applied, restrict regression to data points collected after training dynamics stabilize following the batch size increase.

Empirical Validation / Results

µP Transferability across Width (Section 3.2):

• Base proxy model: hidden dim $d_ {\mathrm{model}} = 256$, total experts $n_ {\mathrm{experts}} = 16$, attention heads = 4, using MLA architecture. Scale width by $2\times$, $4\times$, $8\times$ (up to 30.7B total, 3.6B active parameters), trained on 1.3B tokens. .
• Result (Figure 3): Under µP, the optimal LR identified in the base proxy consistently transfers across all width scales. Under Standard Parameterization(SP),, the optimal LR shifts as width scales, failing to transfer.

.

Token-Scale Dynamics (Section 3.3.1):

• Proxy model: 5.6B total, 1.8B active; held-out target:20.7B total, 3.8B active(2× width), trained up to 100B tokens.
• Result (Figure 4): As token budget increases, optimal LR exhibits a slight downward trend. At each token scale, fitted parabolas for both proxy and target models exhibit highly consistent curvature和vertex locations, confirming LR transfer across widths at each token scale. The × markers (validated at $2 \times 1 0 ^ { - 3 }$) fall directly on the quadratic curves fitted using the remaining four LRs.

.

Extrapolation to 10T Tokens (Section 3.3.2):

• Base proxy:10.8B total, 3.3B active(1/4× width of target), trained for approximately 500B tokens. Estimate optimal LR at 10B-token intervals using EMA weights computed every 2B tokens.
• Result (Figure 5): Linear regression (restricted to data points after 255B tokens, post batch-size increase) achieves $R^2 = 0.95$, predicting an optimal LR of $3.85 \times 1 0 ^ { - 4 }$ for 10T-token training.

Large-Scale Validation (Section 3.3.3):

• Target model: 155B total, 17B active parameters, trained over 10T tokens from scratch. Stage 1 data mixture: initially 45% English, 12.5% Math/STEM, 27.5% Code, 15.0% Multilingual; adjusted at 6T tokens to enhance underrepresented domains. Batch size scheduling applied after 200B tokens.
• Result (Figure 6): Training loss remains highly stable throughout, with no loss spikes, confirming the validity of the extrapolated LR.
• Result (Figure 7): Benchmark scores across English(MMLU, MMLU-Pro, BBH), Multilingual(Global-MMLU: Ko, Ja, Vi, Zh), Math(MATH, GSM8K), and Code(MBPP, HumanEval) confirm successful training.
• Result (Figure 8): Comparison with comparable open-weight MoE base models(dots.llm1, GLM4.5-Air, Hunyuan-A13B, DeepSeek-V4-Flash) using 6ND compute estimation (N = active parameters): our model lies on the Pareto frontier, achieving higher MMLU-Pro accuracy than dots.llm1 and GLM-4.5-Air at comparable or lower estimated training compute.

Compute Savings: 2D sweeps require scaling across model size ($1.5\times$, $2\times$ wider than Proxy), incurring an additional 240.3 ZFLOPs beyond the Proxy run cost of 64.8 ZFLOPs(Figure 2a).. The total compute of the Target is approximately 98× larger than the sum of the Proxy runs(Figure 2b)..

Theoretical and Practical Implications

• Theoretical contribution: Provides the first systematic study of µP-based hyperparameter transfer for large-scale MoE architectures with MLA and the Muon optimizer, extending prior work that was limited to dense models or restrictive MoE settings(e.g.,, fixing total/active expert counts). It also introduces a token-scaling law for optimal LR, decoupling model-scale from token-scale search—replacing costly 2D sweeps with lightweight 1D extrapolation.

• Practical contribution: Demonstrates that proxy training on small models is sufficient to determine the optimal LR for extensive training of large-scale MoEs—reducing unnecessary computation by orders of magnitude(e.g.,, avoiding 240+ ZFLOPs of additional sweeps).. The successful pretraining of a 155B-parameter model over 10T tokens validates the methodology's practical effectiveness at extreme scale.

• Design implications: The framework decouples LR optimization from batch size—acknowledging that batch size is a system-level variable frequently adjusted for hardware throughput, and the literature remains divided on its scaling behavior. This robustness to batch size choice makes the approach broadly applicable across hardware environments.

• Caveats: The authors acknowledge that while the extrapolated LR ($3.85 \times 1 0 ^ { - 4 }$) is validated indirectly via stable training and competitive benchmarks, exhaustive full-scale sweeps to definitively verify optimality remain computationally infeasible.

Conclusion

The paper proposes a compute-efficient two-step framework for hyperparameter transfer in large-scale MoE pretraining, combining:

1. µP-based width transferability to eliminate model-scale sweeps, and
2. A linear scaling law across token budgets to enable direct extrapolation to trillion-token horizons.

This approach successfully predicted the optimal LR ($3.85 \times ^{-4}$) for a 155B-total/17B-active MoE model trained over 10T tokens, achieving stable loss trajectories and competitive benchmark performance.

Future directions:

• Extending the framework to diverse MoE structures and other optimizers(beyond MLA/Muon).
• Investigating per-expert learning rate adaptation, since top-k routing can result in varying token counts per expert, potentially affecting gradient noise and optimal update scales—though requiring careful consideration of routing dynamics, data mixture effects, and engineering costs.
• Conducting a controlled large-scale study of µP transfer along the sparsity axis alone (since the current recipe expands sparsity jointly with width, making their effects difficult to disentangle)..