SkewAdam：分层优化器将 MoE 显存占用降低 97%

MoE（混合专家）模型因参数规模庞大，训练时的显存占用一直是工程化部署的痛点。最近在 Reddit 的 r/MachineLearning 社区，作者 nuemaan 发布了一篇题为 SkewAdam 的预印本与开源代码，提出一种「分层状态分配」优化器思路，把 MoE 训练时的优化器状态显存开销削减了约 97%，并使一个 6.78B 参数的 MoE 模型可在单张 40 GB GPU 上完成训练。

核心思路：按参数角色分配精度

传统优化器（如 AdamW）对模型所有参数一视同仁，需要为每个参数维护动量与二阶动量等状态，这在 MoE 上代价尤其沉重。SkewAdam 把 MoE 的参数按角色拆成三层，分别施加不同精度的状态记录：

• Backbone（骨干，约 5% 参数）：保留动量 + 因子化二阶动量；
• Experts（专家层，约 95% 参数）：仅保留因子化二阶动量；
• Router（路由器，<0.01% 参数）：保留完整精确二阶动量。

这样做的好处是，占绝大多数参数的 Experts 层不再维护精度最高的状态量，从而把优化器状态显存压到极低水平；同时对占比很小但对训练稳定性关键的 Backbone 与 Router 保留充足精度，避免路由不稳定或骨干参数漂移。

硬件实测：显存与峰值占用双双显著下降

作者给出了硬件层面的对比数据。以一个 12.6 GB 规模的 6.78B 参数 MoE 为例：

• 优化器状态显存：由 AdamW 的 50.6 GB 降至 1.29 GB，缩减约 97.4%；
• 训练峰值显存：由 81.4 GB 降至 31.3 GB；
• 收敛表现与路由稳定性与基线一致（按作者自述）。

如此一来，整套训练可在一张 40 GB 显存的消费级或单卡服务器 GPU 上完成，无需多卡并行或激进量化，对中小团队复现与微调 MoE 尤为友好。

工程价值与待验证之处

对 MoE 训练而言，SkewAdam 直接命中了「优化器状态占大头」这一长期痛点。如果结果可复现，这一思路有望被集成进主流训练框架（如 Megatron、DeepSpeed、torchtitan 等），显著降低 MoE 研究的硬件门槛。

不过需要注意的是，目前该工作以预印本形式发布，论文尚未经同行评审，结论以作者自报数据为准；因子化二阶动量在不同模型规模、任务类型上的鲁棒性，也仍有待在外部团队中复现。代码已开源，对 MoE 训练感兴趣的工程师可在 GitHub 上自行验证。