mlsanityData to 14 Oct 2026

Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training

arXiv 2607.19058 ↗First posted 21 Jul 2026cs.LG cs.AIcited by 0

On the map: Mixture-of-experts → Efficient MoE training and inference systems

Abstract (arXiv)

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.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts. Builds on 23 of its 30 references that are in the map. First posted on arXiv 21 Jul 2026.

Builds on

23 of 30 references.

Closest ancestors · 1

Same direction, newest first

  1. 2025Muon is Scalable for LLM Training24 Feb

Other references · 22

  1. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,886
  2. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,100
  3. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  4. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
  5. 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,883
  6. 2024Mixtral of Expertscited by 1,823
  7. 2020GLU Variants Improve Transformercited by 1,804
  8. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
  9. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,346
  10. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
  11. 2017Mixed Precision Trainingcited by 1,156
  12. 2018Adafactor: Adaptive Learning Rates with Sublinear Memory Costcited by 905
10 more
  1. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 867
  2. 2016Training Deep Nets with Sublinear Memory Costcited by 848
  3. 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 463
  4. 2023Symbolic Discovery of Optimization Algorithmscited by 434
  5. 20218-bit Optimizers via Block-wise Quantizationcited by 403
  6. 2022ST-MoE: Designing Stable and Transferable Sparse Expert Modelscited by 399
  7. 2015Deep Learning with Limited Numerical Precisioncited by 376
  8. 2019A Study of BFLOAT16 for Deep Learning Trainingcited by 232
  9. 2024Adam-mini: Use Fewer Learning Rates To Gain Morecited by 137
  10. 2020Revisiting BFloat16 Trainingcited by 17

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training
  2. 2026MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training
  3. 2024Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Training
  4. 2026FlashOptim: Optimizers for Memory-Efficient Training
  5. 2025Mixture-of-Experts Can Surpass Dense LLMs Under Strictly Equal Resource
  6. 2025Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient
  7. 2026Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization
  8. 2024MoE-CAP: Benchmarking Cost, Accuracy and Performance of Sparse Mixture-of-Experts Systems