SkewAdam: A Tiered Optimizer Strategy for Memory-Efficient MoE Training

Introduction

Mixture-of-Experts models promise larger parameter counts without activating every parameter for every token. But their training memory budget is still difficult to manage. Beyond weights, gradients, and activations, optimizer state can become the dominant line item. The paper “Where Should Optimizer State Live?” focuses exactly on that problem and proposes SkewAdam, an optimizer designed around the internal structure of MoE models.

In the reported 6.78B-parameter MoE language model, bfloat16 weights take 12.6GB. AdamW, however, keeps 50.6GB of first- and second-moment state. That makes optimizer state larger than the weights themselves by a wide margin. SkewAdam starts from a simple observation: the dense backbone, expert layers, and router do not have the same size or gradient behavior, so they should not necessarily receive the same optimizer state.

Key points

• State is allocated by parameter role. For the dense backbone, which represents about 5% of parameters, SkewAdam keeps float32 momentum plus a factored second moment. For the experts, which represent about 95%, it keeps only a factored second moment. For the tiny router, under 0.01% of parameters, it keeps an exact second moment.
• The memory reduction is large. Optimizer state drops from AdamW’s 50.6GB to 1.29GB, or 2.6% of the original. Peak training memory falls from 81.4GB to 31.3GB, putting the experiment within a 40GB accelerator budget.
• The reported validation result improves rather than degrades. In a controlled comparison from identical initialization over 82M tokens, SkewAdam reaches validation perplexity 108.4, compared with 126.8 for AdamW, 120.2 for Muon, and 393.7 for Lion.
• Tiering saves memory; momentum matters for accuracy. The paper notes that a tier ablation can match the perplexity with much more state, while Adafactor, which uses a related factored estimator but drops momentum, plateaus far behind.

Why it matters

The broader implication is that optimizer design for MoE models should be topology-aware. Treating every parameter tensor the same may be convenient, but it ignores how uneven MoE models are: experts dominate parameter count, the backbone is smaller but central, and the router is tiny yet important for load balance.

SkewAdam reframes the memory problem as an allocation problem. Instead of asking only how to compress optimizer state globally, it asks which parts of the model deserve which kind of state. That distinction is important because the paper’s results suggest that aggressive memory savings do not have to come at the cost of validation quality, at least in the reported setting.

There are still open questions. The evidence comes from a specific model scale, training budget, and evaluation setup, so the result should be validated across larger runs and other MoE architectures. Still, the message is useful: for MoE training, where optimizer state lives can matter as much as how much state is stored.

Source: Hugging Face Daily Papers