Tiered optimizer states free more than half of GPU RAM for Mixture‑of‑Experts training. The twist is that a simple reallocation of Adam’s moments—rather than any new learning rule—delivers both the memory cut and a perplexity win.
Before SkewAdam, practitioners accepted that AdamW’s first‑ and second‑moment buffers dominate the memory budget of MoE models; a 6.78 B‑parameter model already needs roughly 50 GB just for optimizer state, forcing multi‑GPU pipelines or aggressive gradient checkpointing.
Peak training memory collapses from 81.4 GB to 31.3 GB when SkewAdam’s tiered buffers are used, with the optimizer’s state shrinking to 1.29 GB (2.6 % of AdamW) — enough to fit a single 40 GB accelerator [1]. The authors achieve this by keeping full‑precision momentum only for the dense backbone, applying factored variance to the expert bank, and storing an exact second moment solely for the router.
Validation perplexity improves to 108.4, beating AdamW’s 126.8, Muon’s 120.2, and Lion’s 393.7 on a controlled 82 M‑token run [1]. The paper stresses that this gain stems from preserving momentum across all tiers; an ablation with twenty‑times more state but no tiering matches the perplexity, while Adafactor—sharing the factored estimator but dropping momentum—lags by 40 points.
The study reports results on a single 6.78 B‑parameter MoE model and a modest token budget, leaving open whether the same memory‑accuracy tradeoff scales to denser architectures or longer pretraining runs [1]. It also assumes that developers can integrate a custom optimizer without incurring extra engineering overhead.
If tiered state allocation truly generalizes, the default training stack for large MoE systems should replace AdamW with SkewAdam, eliminating the need for multi‑GPU memory hacks and enabling larger batch sizes on commodity hardware.
