Search papers, labs, and topics across Lattice.
This paper introduces SkewAdam, a memory-efficient optimizer designed for mixture-of-experts (MoE) training that strategically allocates optimizer state based on the differing characteristics of parameter populations. By utilizing a tiered state allocation approach, SkewAdam reduces peak training memory from 81.4 GB to 31.3 GB while achieving superior validation perplexity of 108.4, significantly outperforming traditional optimizers like AdamW and Muon. The findings highlight that the allocation of optimizer state can be as critical to performance as the total amount of state used, challenging conventional wisdom in optimizer design.
Memory-efficient optimizer state allocation can lead to dramatic improvements in training performance without sacrificing accuracy, as shown by SkewAdam's superior perplexity results.
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.
