オプティマイザの状態はどこに配置すべきか？メモリ効率的なMixture-of-Experts学習のための階層的状態割り当て

要旨

オプティマイザ状態は、混合エキスパート（MoE）学習のメモリ予算において最大の単一項目である。6.78BパラメータのMoE言語モデルでは、AdamWは12.6GBのbfloat16重みを更新するために、50.6GBの一次モーメントと二次モーメントを保持する。本研究では、MoEの3つのパラメータ集団（高密度バックボーン、エキスパート、ルータ）がサイズと勾配統計量において十分に異なり、同一の状態を受け取るべきではないという観察に基づいたオプティマイザであるSkewAdamを研究する。SkewAdamは、バックボーン（パラメータの5%）にはfloat32の運動量と因子分解された二次モーメントを、エキスパート（95%）には因子分解された二次モーメントのみを、ルータ（0.01%未満）には正確な二次モーメントを保持する。結果として得られる状態は1.29GBを占め、AdamWの2.6%に相当し、ピークトレーニングメモリは81.4GBから31.3GBに減少し、40GBアクセラレータの予算内に収まる。同一初期化から8200万トークンにわたる制御比較において、SkewAdamは検証パープレキシティ108.4を達成し、AdamW（126.8）、Muon（120.2）、Lion（393.7）を上回り、ルータの負荷バランスを均一下限の1%以内に収めた。この割り当て自体がパープレキシティをもたらすわけではない。階層アブレーションでは、20倍の状態量で同様の結果を示し、因子分解推定器を共有するが運動量を削除した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.