循序漸進地擴展：大規模混合專家模型的計算高效超參數遷移

摘要

混合專家（MoE）架構能顯著擴展模型容量，同時避免計算成本等比例增加。然而，在模型規模與詞元預算皆處於極端尺度下，透過掃蕩方式優化其超參數——尤其是學習率——在計算上仍是難以負荷的。本文提出一個節省計算資源的兩階段超參數遷移框架，透過跨模型寬度遷移來估計大規模 MoE 模型訓練的最佳學習率，進而外推至兆級詞元訓練視野。首先，我們為採用多頭潛在注意力（MLA）與 Muon 優化器的 MoE 架構，制定最大化更新參數化（μP）的適配方案，證明最佳學習率能在不同寬度縮放的模型之間一致地遷移。其次，我們沿著詞元維度建立可預測的標度律，將此可遷移性進一步延伸。透過對有限預算下的小型代理模型所推導出的最佳值應用線性回歸，我們成功地將理想學習率外推至大規模訓練視野（例如 10 兆詞元），且具備高度擬合度（R²=0.95）。因此，這表明在小型模型上進行代理訓練，便足以確定大規模 MoE 廣泛訓練的最佳學習率。我們將所提出的方法應用於從零開始預訓練基礎模型（總參數 155B，啟用參數 17B），穩定的訓練與評估結果驗證了：僅需極少的消融成本，即可準確預測全規模目標模型的最佳配置。

English

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.