一句话总结

Kakao Corp. 与 Upstage AI 提出一种面向大规模 MoE 模型的计算高效超参数迁移框架，该框架将 µP 适配到多头潜在注意力（Multi-head Latent Attention）和 Muon 优化器，以在不同宽度之间迁移最优学习率，并利用 token 缩放定律 $( R ^ { 2 } = 0 . 9 5 )$(R2=0.95) 从小型代理模型外推到万亿 token 训练，成功预训练了一个 155B 参数模型。

核心贡献

• 针对具有多头潜在注意力（Multi-head Latent Attention）和 Muon 优化器的专家混合架构，提出了最大更新参数化（µP）适配方法，证明最优学习率能够在不同宽度缩放的模型之间一致迁移。
• 一个可预测的缩放定律将这种可迁移性沿 token 维度扩展，从而能够通过高保真线性回归（R²=0.95）从小型代理模型外推万亿 token 训练的最优学习率。
• 该两步框架通过从零开始预训练一个 155B 参数 MoE 基础模型（17B 激活参数）得到验证，稳定损失和有竞争力的基准结果表明，最优超参数能够以极小的消融成本被准确预测。

引言

专家混合（MoE）架构使大型语言模型能够扩展容量和专业化，同时保持较低的每 token 活跃计算量，但调整其学习率变得昂贵，因为最优值取决于模型规模、token 预算和专家路由。现有的超参数迁移方法（如最大更新参数化）主要针对稠密宽度缩放，因此它们在 MoE 稀疏性和大量专家数量下的可靠性仍不确定。作者提出一种两步、计算高效的框架，在模型宽度和专家数量之间迁移超参数，同时将最优学习率从小型代理运行外推到较长 token 范围。他们通过预训练一个总参数 155B、活跃参数 17B 的 MoE 模型、训练规模为 10 万亿 token 来验证该方法，避免了昂贵的多维搜索。

方法

作者提出一种两步超参数迁移框架，旨在高效确定大规模专家混合（MoE）模型的最优学习率。第一步将最大更新参数化（µP）适配到 MoE 架构，使最优学习率能够在不同模型宽度之间零样本迁移。第二步建立一条缩放定律，预测最优学习率随 token 预算变化的规律，从而能够从短代理运行外推到完整训练范围。

面向 MoE 的 µP 适配首先根据参数在宽度缩放下的行为对模型参数进行分类。由于没有任何参数的形状与词汇量等固定常量相关联，因此排除了标量类参数。向量类参数仅采用 µP 初始化，而矩阵类参数同时采用 µP 初始化和学习率缩放。基于先前发现，即仅对线性层缩放学习率即可保持 µP 性质，作者仅对矩阵类隐藏权重应用学习率缩放。由于已知深度缩放存在不稳定性，因此避免深度缩放；仅考虑宽度缩放，并保持层数固定。对于注意力，头维度保持恒定，头数随隐藏维度按比例缩放。

在缩放 MoE 模型时，保持每个 token 的活跃专家数和 MoE 中间维度固定，同时增加总专家数和隐藏维度。该策略将稀疏性提升为与宽度并列的关键缩放轴，使总规模、活跃规模和稀疏性耦合。这种耦合仍与 µP 兼容，因为通过增加专家数量改变活跃比例不会改变单个专家的扇入或扇出，其影响不超过宽度缩放带来的影响。此外，在较低稀疏性下进行代理搜索，并迁移到较高稀疏性目标，既保持搜索本身的计算效率，又仍然针对所需的稀疏区间。

第二步解决的问题是，即使使用 µP，在数万亿 token 上直接搜索学习率仍然代价过高。解决方案是在受限 token 预算内确定最优学习率，并外推到完整目标范围。该框架有意将学习率优化与批量大小解耦，将批量大小视为系统级变量，固定为最大化 GPU 吞吐量的取值。这使得学习率缩放定律对所选批量大小具有稳健性。

对于短 token 估计，作者使用 Warmup-Stable-Decay（WSD）调度器作为长范围设置，在有限 token 范围内训练小型代理模型。为避免过早引入衰减的计算成本和偏差，代理运行在稳定阶段终止而不进行衰减，并对模型权重应用指数移动平均（EMA）。EMA 平滑参数轨迹并模拟学习率衰减的效果，这在本文使用的大批量设置（全局批量大小为 32M token）中尤其有效。EMA 参数 $\theta_{\mathrm{EMA}}^{(t)}$θEMA(t)​ 更新为

$$\theta_{\mathrm{EMA}}^{(t)} = \alpha \cdot \theta_{\mathrm{EMA}}^{(t-1)} + (1 - \alpha) \cdot \theta^{(t)}$$θEMA(t)​=α⋅θEMA(t−1)​+(1−α)⋅θ(t)

平滑系数为 $\alpha = 0.6$α=0.6，确保最近 20B token 保留超过 1% 的影响。检查点通过 EMA 以约 2B token 间隔保存，每 10B token 处的检查点用于分析。对于每个 token 预算 $B$B，验证损失 $\mathcal{L}$L 被建模为对数变换学习率的二阶多项式：

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

顶点 $\log \eta^* = -b / 2a$logη∗=−b/2a 给出该预算下的估计最优学习率 $\eta^* = \exp(-b / 2a)$η∗=exp(−b/2a)。

为了外推到长 token 范围，作者在双对数空间中拟合估计最优学习率 $\eta^*$η∗ 与 token 预算 $B$B 之间的线性回归：

$$\log(\eta^*) = \beta \cdot \log(B) + \gamma.$$log(η∗)=β⋅log(B)+γ.

该回归模型能够在无需额外计算的情况下预测目标预算（例如 10T token）下的最优学习率。当使用批量大小调度时，只有批量大小增加后训练动态稳定下来之后收集的数据点才会被纳入回归，从而确保一致的优化状态。

实验

评估在 NVIDIA H200 GPU 上使用 µP、Muon 优化器和带批量大小调度的 WSD 调度器。实验首先验证 µP 使最优学习率能够在不同 MoE 模型宽度之间迁移，并且该迁移在不同 token 预算下持续成立，最优学习率随着处理更多 token 而略有下降。通过从小型代理运行外推线性拟合，预测并应用了在 10T token 上训练的 155B 参数 MoE 模型的最优学习率，结果训练损失稳定，并且该模型位于帕累托前沿，在相近或更低计算量下优于可比的开源权重模型。

稠密模型和 MoE 模型中的矩阵类参数根据其扇入与基础代理模型扇入之比进行学习率缩放，而向量类参数仅采用 µP 初始化。对于多头潜在注意力（MLA），低秩投影维度在宽度缩放期间保持固定，因此相应上投影矩阵的学习率缩放因子变为 1，不发生作用。在缩放 MoE 模型时，增加总专家数同时保持每个 token 的活跃专家数和中间维度固定，使稀疏性与宽度耦合，保持与基于 µP 的超参数迁移兼容。学习率缩放仅应用于矩阵类隐藏权重，缩放因子与扇入对基础代理模型扇入的比值成比例。在 MLA 中，固定的查询和键值低秩维度使上投影矩阵的学习率缩放降为 1，从而抵消其影响。MoE 缩放增加总专家数和隐藏维度，同时保持每个 token 的活跃专家数和 MoE 中间维度恒定，在相同 µP 规则下耦合稀疏性和宽度。向量类参数（例如嵌入、偏置、专家 FC2 权重）在宽度扩展期间接受 µP 初始化但不进行学习率缩放。

在面向专家混合模型的最大更新参数化中，参数根据其形状随模型宽度的缩放方式分类。矩阵类参数（如路由器权重和专家 FC1 权重）使用初始化方差和学习率按扇入倒数缩放，而向量类参数（包括嵌入、偏置和专家 FC2 权重）使用固定初始化且不进行学习率缩放。该分类能够在宽度和稀疏性联合缩放时实现稳定的超参数迁移。矩阵类参数（路由器权重、专家 FC1）同时接受 µP 初始化和与扇入成反比的学习率缩放。向量类参数（嵌入、偏置、专家 FC2）使用固定初始化方差和学习率缩放因子 1，即不随模型宽度调整。

评估在最大更新参数化下研究稠密模型和专家混合模型（包括多头潜在注意力）的宽度和稀疏性缩放。矩阵类参数（如路由器权重和专家 FC1 权重）根据扇入进行初始化和学习率缩放，而向量类参数（如嵌入、偏置和专家 FC2 权重）使用固定初始化且不进行学习率缩放。在 MLA 中，固定的低秩投影维度使上投影学习率因子等于 1，而 MoE 缩放通过增加总专家数同时保持活跃专家数和中间维度固定，使稀疏性与宽度耦合。