做大规模MoE模型训练的团队终于有了可靠的超参数迁移方案——不用为每个专家配置重新调参,直接复用密集模型的超参数即可,建议做预训练优化的点开看看。
研究者提出Complete-muE框架,解决了从密集FFN到混合专家(MoE)模型架构的超参数迁移问题。现有方法如μP和SDE无法处理MoE中专家数量变化带来的架构和每个专家token数同时改变的问题。Complete-muE通过双桥系统实现迁移:桥I利用激活宽度μP和归一化路由器尺度映射密集FFN到密集MoE;桥II通过激活专家缩放映射密集MoE到稀疏MoE,并处理一阶SDE学习率/权重衰减校正。实验表明,在语言模型和扩散模型预训练中,该框架能使超参数在多种MoE配置下保持稳定,实现“一次调参,迁移所有”的效果,显著加速MoE模型收敛。
Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models
We propose Complete-muE, a framework which targets hyperparameter transfer across dense FFN and any Mixture-of-Experts (MoE) setups in transformer blocks. Existing tools such as $μ$P (requires fixed architectue) or SDE (requires fixed per-step token count) cannot directly solve the hyperparameter transfer problem in MoE setups because Dense to MoE transfer or MoE total experts scaling changes both architecture and tokens per expert. Complete-muE solves this challenge with a two-bridge system: Bridge~I maps between dense FFN and Dense MoE by active-width $μ$P with a normalized router scale. Bridge~II maps between Dense MoE and sparse MoE by activated-expert scaling, where the first-order SDE LR/WD correction cancels while a bounded residual $σ_0$ shift remains. The resulting transfer rule, which we term as Complete muE, covers changes in activated experts, total capacity, granularity, and shared/group-balanced hybrids for MoE models as well as network width/depth, batch size, and duration changes for general Transformer models. Extensive language model and diffusion model pretraining experiments confirm that complete-muE yields relatively stable hyperparameter optima across model architectures and parameter counts -- with only minor drift consistent with the non-strict SDE behavior of Bridge~II. In practice this drift is small enough that hyperparameters tuned on a single dense reference transfer near-optimally to all MoE configurations -- \emph{tune dense once, transfer to all} is the practical recipe at the core of Complete-muE. This enables MoE models to achieve accelerated convergence speedup over dense models when scaling model capacity without costly hyperparameter search.