针对专家模型训练中优化器状态占用过大的问题，提出按参数类型分层分配状态内存的方案。
Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
- 根据参数类型（主干、专家、路由器）差异，分层分配优化器状态
- 内存从81.4GB降至31.3GB，仅需40GB显存即可训练
- 分层策略提升性能，优于未分层的先进优化器
在6.78亿参数的MoE语言模型上，AdamW需50.6GB优化器状态来更新12.6GB的bfloat16权重。本文提出SkewAdam，基于三个参数组（主干、专家、路由器）在规模与梯度统计上的差异，分别分配不同状态：主干使用float32动量+分解的二阶矩（占参数5%），专家仅用分解的二阶矩（95%），路由器用精确二阶矩（<0.01%）。最终状态仅1.29GB，占原方案2.6%，峰值显存从81.4GB降至31.3GB，可部署于40GB加速器。在相同初始化下训练8200万词元，SkewAdam验证困惑度达108.4，优于AdamW（126.8）、Muon（120.2）和Lion（393.7），且路由器负载均衡误差低于1%。消融实验表明，性能提升源于分层而非状态量；去除动量导致困惑度上升31点，替换为全量二阶矩则上升10点，说明分层设计本身是关键。
原文摘要 · Abstract (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 differ enough in size and gradient statistics that they should not receive the same state. Those populations are the dense backbone, the experts and the router. 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 or 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 reaches the same value while carrying twenty times the state, so the tiers buy memory rather than accuracy. Same-platform runs separate what does earn it. Removing momentum costs 31 perplexity points (tuned Adafactor, 139.7) and replacing the factored second moment and its update clipping with a full second moment costs 10 (tuned AdamW, 118.5), so neither tuned baseline reaches the untuned tiered policy. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。