优化器状态应该放在哪里？用于内存高效混合专家训练的分层状态分配

优化器状态是专家混合 (MoE) 训练的内存预算中最大的单行项目。在 6.78B 参数的 MoE 语言模型上，AdamW 保留 50.6 GB 的第一和第二时刻来更新 12.6 GB 的 bfloat16 权重。我们研究了 SkewAdam，这是一种基于观察的优化器，即 MoE 的三个参数群体在大小和梯度统计方面差异很大，因此它们不应该接收相同的状态。这些人群是密集的骨干网、专家和路由器。 SkewAdam 保持 float32 动量加上主干网的因子二阶矩（参数的 5%）、专家的因子二阶矩（95%）和路由器的精确二阶矩（<0.01%）。最终的状态占用 1.29 GB 或 AdamW 的 2.6%，峰值训练内存从 81.4 GB 降至 31.3 GB，在 40 GB 加速器的预算范围内。在对超过 82M 词元的相同初始化进行的受控比较中，SkewAdam 的验证困惑度达到 108.4，领先于 AdamW (126.8)、Muon (120.2) 和 Lion (393.7)，并将路由器负载平衡稳定在其统一下限的 1% 以内。分配并不是造成这种困惑的原因。层消融在携带二十倍状态的同时达到相同的值，因此层购买的是内存而不是准确性。同一平台运行的内容是分开的。消除动量需要 31 个困惑点（调整后的 Adafactor，139.7），并用完整的二阶矩替换分解的二阶矩及其更新裁剪，成本为 10（调整后的 AdamW，118.5），因此调整后的基线都没有达到未调整的分层策略。这些结果表明，优化器状态所在的位置至少与它的数量一样重要。