SkewAdam：把优化器状态放到该放的位置，让 MoE 训练更省显存

导语

MoE（Mixture-of-Experts）模型常被视为扩大参数规模、控制计算量的有效路径，但训练时的显存账并不轻松。除了权重、梯度和激活值，优化器状态往往是隐藏的大头。论文《Where Should Optimizer State Live?》提出的 SkewAdam，试图回答一个更细的问题：优化器状态到底应该放在哪些参数上、以什么形式保存？

在作者的 6.78B 参数 MoE 语言模型实验中，bfloat16 权重只占 12.6GB，而 AdamW 为一阶和二阶矩维护的状态达到 50.6GB，成为单项最大显存开销。SkewAdam 的出发点是：MoE 里的 dense backbone、experts 和 router 在参数占比与梯度特性上差异明显，不应被同一种状态策略“一刀切”。

核心要点

• 分层分配优化器状态：SkewAdam 为约 5% 参数量的 dense backbone 保留 float32 动量和分解式二阶矩；为占 95% 参数的 experts 只保留分解式二阶矩；为占比不到 0.01% 的 router 保留精确二阶矩。
• 显存下降非常明显：优化器状态从 AdamW 的 50.6GB 降至 1.29GB，仅为其 2.6%；训练峰值显存从 81.4GB 降到 31.3GB，落入 40GB 加速卡的预算范围。
• 效果没有因省内存而牺牲：在相同初始化、82M tokens 的对照中，SkewAdam 验证困惑度为 108.4，优于 AdamW 的 126.8、Muon 的 120.2 和 Lion 的 393.7。
• 精度来自动量，而不是“省状态”本身：消融显示，分层策略主要带来内存收益；与 Adafactor 相比，保留 momentum 对性能很关键。

意义与影响

这项工作有价值的地方，不只是把数字压小，而是提出了一种更符合 MoE 结构的优化器设计思路。过去很多训练配置默认对所有参数使用统一优化器状态，但 MoE 的专家参数占比极高、路由器又极小且敏感，统一策略可能既浪费显存，也没有充分照顾关键模块。

如果后续在更大规模、更长训练和更多架构上得到验证，类似 SkewAdam 的“按参数角色分配状态”可能成为 MoE 训练系统的重要组件。它尤其适合显存预算受限但希望训练多十亿参数模型的团队：与其单纯依赖更大的 GPU，不如重新审视优化器状态这笔账。

当然，素材中的实验仍是特定模型和训练预算下的结果，是否能直接外推到更大数据、更复杂训练管线，还需要更多验证。但它清楚地提示了一点：在 MoE 训练中，优化器状态不仅要问“要多少”，还要问“该放在哪里”。

来源：Hugging Face Daily Papers