---
title: "Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training"
source: "arXiv AI/CL/LG（论文）"
source_tier: "t2"
url: "https://arxiv.org/abs/2607.19058v1"
permalink: "https://feed.astralor.com/items/itm_b4bae957f7521b6e57c79ab6"
published: "2026-07-21T12:47:15.000Z"
score: 83
channel: "ai"
tags: ["AdamW","SkewAdam","Tiered State Allocation","Momentum Optimization","Factored Second Moment","Mixture-of-Experts (MoE)"]
---

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

## AI commentary

> This paper is highly valuable for teams training MoE models: it demonstrates that with careful optimizer state allocation, one can dramatically reduce memory without sacrificing accuracy—and even improve performance. It is especially relevant for labs or enterprises constrained by 40 GB or 80 GB GPU memory. The tiered strategy and ablation on momentum retention are worth studying for optimizing training frameworks like DeepSpeed or Megatron.

## At a glance

SkewAdam assigns different optimizer state formats to the backbone, experts, and router of MoE, compressing AdamW's 50.6 GB state to 1.29 GB (2.6%), reducing peak training memory from 81.4 GB to 31.3 GB while achieving better validation perplexity (108.4 vs AdamW 126.8, Muon 120.2, Lion 393.7). Results show that where state lives matters as much as how much there is; retaining momentum is critical for accuracy, while tiered allocation saves memory without harming performance.

## Key points

- Optimizer state is the largest memory consumer in MoE training, occupying over 80% of AdamW's memory budget.
- SkewAdam uses float32 momentum plus factored second moment for the backbone (5% params), only factored second moment for experts (95%), and exact second moment for the router (<0.01%).
- Compared to AdamW, state memory reduces by 97.4%, peak memory by 61.5%, with improved perplexity.
- Momentum is critical for accuracy; merely reducing state (e.g., Adafactor) causes significant performance degradation.

## Full text

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.

---

Verify against the original before citing; curated by Astral Feed: https://feed.astralor.com/items/itm_b4bae957f7521b6e57c79ab6