優化器狀態應該存在哪裡？用於記憶體高效混合專家訓練的分層狀態分配

摘要

优化器状态是混合专家（MoE）模型训练中内存预算的最大单一开销项：对于一个67.8亿参数的MoE语言模型，AdamW需要占用50.6 GB的一阶和二阶动量来更新12.6 GB的bfloat16权重。我们研究SkewAdam优化器，其构建基于以下观察：MoE的三种参数群体——密集骨干网络、专家模块和路由模块——在规模与梯度统计特性上差异显著，不应共享相同状态。SkewAdam为骨干网络（占参数5%）保留float32动量及因式分解二阶动量，为专家模块（占参数95%）仅保留因式分解二阶动量，为路由模块（占参数<0.01%）保留精确二阶动量。由此产生的状态仅占用1.29 GB，相当于AdamW的2.6%，训练峰值内存从81.4 GB降至31.3 GB，完全适配40 GB加速器的预算。在基于相同初始化的受控比较中（训练8.2M token），SkewAdam达到验证困惑度108.4，领先于AdamW（126.8）、Muon（120.2）和Lion（393.7），并将路由负载均衡度稳定在均匀分布基线的1%以内。这种层级分配并非困惑度提升的核心原因：层级消融实验在使用二十倍状态量时仍能复现相同效果，而共享因式化估计量但丢弃动量的Adafactor优化器则落后40个困惑度点。层级设计以零精度代价换取内存节省，真正的精度提升来自动量保持——而这恰恰是统一优化器同样具备的特性。通过扫描基线方法的学习率可缩小差距但无法消除：最优调参后的AdamW达到118.5，Adafactor达到139.7。这些结果表明，优化器状态的位置分布，其重要性至少不亚于状态量的大小。

English

Optimizer state is the largest single line item in the memory budget of mixture-of-experts (MoE) training: on a 6.78B-parameter MoE language model, AdamW keeps 50.6 GB of first and second moments to update 12.6 GB of bfloat16 weights. We study SkewAdam, an optimizer built on the observation that the three parameter populations of an MoE - the dense backbone, the experts, and the router - differ enough in size and gradient statistics that they should not receive the same state. SkewAdam keeps float32 momentum plus a factored second moment for the backbone (5% of parameters), a factored second moment alone for the experts (95%), and an exact second moment for the router (<0.01%). The resulting state occupies 1.29 GB, 2.6% of AdamW's, and peak training memory falls from 81.4 GB to 31.3 GB, within the budget of a 40 GB accelerator. In a controlled comparison from identical initializations over 82M tokens, SkewAdam reaches validation perplexity 108.4, ahead of AdamW (126.8), Muon (120.2), and Lion (393.7), and settles router load balance to within 1% of its uniform floor. The allocation is not what earns that perplexity: a tier ablation matches it with twenty times the state, and Adafactor, which shares the factored estimator but drops momentum, plateaus 40 points behind. The tiers buy memory at no cost to accuracy; the accuracy comes from keeping momentum, which a uniform optimizer shares too. Sweeping the baselines' learning rates narrows but does not close the gap: the best tuned AdamW reaches 118.5, tuned Adafactor 139.7. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.