SkewAdam 将混合专家（Mixture-of-Experts, MoE）训练显存降低了 60% 以上，并带来了显著的准确率提升，让开发者能够在单个 40 GB GPU 上运行 67.8 亿参数的 MoE 模型。

为什么 MoE 训练会遭遇显存瓶颈

混合专家（MoE）架构在保持稠密骨干网络（dense backbone）的同时，将每个输入路由到“专家”子网络的较小子集中。其优势在于模型可以扩展到数十亿参数，而计算量不会成比例增加。在实践中，优化器——特别是大多数团队使用的 AdamW 变体——消耗了大部分 GPU 显存。对于一个 67.8 亿参数的模型，仅优化器的动量（momentum）和方差（variance）张量就需要约 50 GB。再加上激活值（activations）和模型权重，峰值显存将超过 81.4 GB，迫使开发者不得不使用多 GPU 配置或采用更慢的训练计划。

SkewAdam 的不同之处

SkewAdam 并没有发明新的学习规则。它重新分配了 Adam 矩（moments）的存储位置：

  • 稠密骨干网络保留全精度动量，以维持稠密层所需的平滑更新。
  • 专家库存储方差的分解近似值（factored approximation），从而减少了每个专家需要保存的数据量。
  • 路由器（决定激活哪些专家）保留精确的二阶矩（second-moment）估计。

通过将这三个组件视为不同的“层级”（tiers），优化器在不那么重要的环节降低了冗余精度，而在最重要的环节保留了精度。

可衡量的影响

使用 SkewAdam 运行相同的 67.8 亿参数 MoE 模型可获得：

  • 峰值 GPU 显存：31.3 GB（从 81.4 GB 降至此）
  • 优化器状态占用：1.29 GB（从 50 GB 降至此）

缩减后的状态可以轻松放入单个 40 GB 加速器中。

不仅仅是节省，还有准确率提升

显存的削减通常会损害模型质量，但 SkewAdam 却提升了困惑度（perplexity）——这是语言模型的一项标准指标——从 126.8 (AdamW) 降至 108.4。在相同任务上，它的表现也优于 Muon (120.2) 和 Lion (393.7)。作者表示，这种提升源于在所有三个层级中都保留了动量；而像 Adafactor 这样完全舍弃动量的方法则表现落后。

开放性问题

SkewAdam 的结果是在一个 67.8 亿参数的模型上展示的。目前尚不清楚对于规模大一个数量级的模型，或者对于语言建模以外的任务，是否仍能保持相同的显存状态比例。

值得关注的方向

  • 更大规模 MoE 配置（数百亿参数）的基准测试。
  • 开源框架的采用情况以及是否被纳入流行的训练脚本中。
  • 社区对于潜在权衡（trade-offs）的反馈，例如在不同学习率调度下的收敛速度或稳定性。

来源：详细介绍优化器设计和实证结果的开发者帖子。