Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
Merged summary
TL;DR - SkewAdam allocates different optimizer states to MoE backbones, experts, and routers, cutting optimizer-state memory to 1.29 GB while outperforming tested baselines. It enables a 6.78B-parameter MoE model to train within a 40 GB accelerator budget.
- Reduces optimizer state from AdamW’s 50.6 GB to 1.29 GB and peak training memory from 81.4 GB to 31.3 GB.
- Uses momentum plus factored second moments for the backbone, factored second moments for experts, and exact second moments for the router.
- Reaches 108.4 validation perplexity versus 126.8 for AdamW, 120.2 for Muon, and 393.7 for Lion over 82M tokens.
- Ablations indicate tiering provides memory savings without accuracy loss, while retained momentum drives the performance advantage.
Sources (1)
Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
TL;DR - SkewAdam allocates different optimizer states to MoE backbones, experts, and routers, cutting optimizer-state memory to 1.29 GB while outperforming tested baselines. It enables a 6.78B-parameter MoE model to train within a 40 GB accelerator budget.
- Reduces optimizer state from AdamW’s 50.6 GB to 1.29 GB and peak training memory from 81.4 GB to 31.3 GB.
- Uses momentum plus factored second moments for the backbone, factored second moments for experts, and exact second moments for the router.
- Reaches 108.4 validation perplexity versus 126.8 for AdamW, 120.2 for Muon, and 393.7 for Lion over 82M tokens.
- Ablations indicate tiering provides memory savings without accuracy loss, while retained momentum drives the performance advantage.