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

Hugging Face Daily Papers Papers
mixture-of-experts optimizer memory-efficient tiered-allocation skeadam training

Summary

SkewAdam is a novel optimizer for mixture-of-experts models that tier allocates optimizer state across backbone, experts, and router, reducing memory footprint to 2.6% of AdamW while achieving better validation perplexity in controlled comparisons.

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.
Original Article
View Cached Full Text

Cached at: 07/22/26, 10:42 AM

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

Source: https://huggingface.co/papers/2607.19058

Abstract

Optimizerstateisthelargestsinglelineiteminthememorybudgetofmixture-of-experts(MoE)training:ona6.78B-parameterMoElanguagemodel,AdamWkeeps50.6GBoffirstandsecondmomentstoupdate12.6GBofbfloat16weights.WestudySkewAdam,anoptimizerbuiltontheobservationthatthethreeparameterpopulationsofanMoE-thedensebackbone,theexperts,andtherouter-differenoughinsizeandgradientstatisticsthattheyshouldnotreceivethesamestate.SkewAdamkeepsfloat32momentumplusafactoredsecondmomentforthebackbone(5%ofparameters),afactoredsecondmomentalonefortheexperts(95%),andanexactsecondmomentfortherouter(<0.01%).Theresultingstateoccupies1.29GB,2.6%ofAdamW’s,andpeaktrainingmemoryfallsfrom81.4GBto31.3GB,withinthebudgetofa40GBaccelerator.Inacontrolledcomparisonfromidenticalinitializationsover82Mtokens,SkewAdamreachesvalidationperplexity108.4,aheadofAdamW(126.8),Muon(120.2),andLion(393.7),andsettlesrouterloadbalancetowithin1%ofitsuniformfloor.Theallocationisnotwhatearnsthatperplexity:atierablationmatchesitwithtwentytimesthestate,andAdafactor,whichsharesthefactoredestimatorbutdropsmomentum,plateaus40pointsbehind.Thetiersbuymemoryatnocosttoaccuracy;theaccuracycomesfromkeepingmomentum,whichauniformoptimizersharestoo.Sweepingthebaselines’learningratesnarrowsbutdoesnotclosethegap:thebesttunedAdamWreaches118.5,tunedAdafactor139.7.Whereoptimizerstatelives,theseresultssuggest,mattersatleastasmuchashowmuchofitthereis.

View arXiv pageView PDFGitHub1Add to collection

Get this paper in your agent:

hf papers read 2607\.19058

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2607.19058 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2607.19058 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2607.19058 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Gefen: Optimized Stochastic Optimizer

arXiv cs.LG

Gefen is a memory-efficient optimizer that reduces AdamW's memory footprint by ~8x by automatically sharing second-moment estimates and quantizing the first moment using a learned codebook, while maintaining AdamW-level performance.

MawForge: Memory-Bounded Expert Materialization for Local Mixture-of-Experts Inference

arXiv cs.LG

MawForge presents a memory-bounded approach for serving large Mixture-of-Experts (MoE) language models on constrained unified-memory machines by storing the full model on disk and materializing expert tensors into a bounded cache on demand. Experiments on a MacBook Pro M5 Pro demonstrate effective serving of 34GB and 25GB models within a 24GB memory envelope, with analysis of cache size trade-offs and speculative decoding results.