Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training

Hugging Face Daily Papers Papers
long-context mixture-of-experts memory-optimization training gpu-memory pipelining offloading

Summary

This paper introduces techniques to manage memory peaks in training large Mixture-of-Experts models with long context lengths, including Pipelined LLEP, Ring-DTP, SCO, and OffloadStreamAdamW, which enable fixed GPU working sets and improve throughput up to 10.4x.

Training a Mixture-of-Experts (MoE) model at long context or large batch size fails as soon as any one component's peak allocation exceeds device memory, so the target is every peak at once, not the average footprint. Four are left unbounded by the parallelism plans in common use, and each grows differently: expert dispatch with the routing matrix, the vocabulary projection with tokens times vocabulary, gradient checkpoint boundaries with depth times sequence length, and optimizer state with parameter count. Which one runs out first changes with the model, the context length, and the device count, so lowering the largest only exposes the next. We bound all four with schedules whose GPU working set is fixed at launch: PipelinedLLEP extends least-loaded expert parallelism with a cap on the tokens each source contributes to a dispatch chunk, Ring-DTP circulates activations or weight shards around a ring at the vocabulary projection and folds each block of logits into an online log-sum-exp, Selective checkpoint offload (SCO) keeps the one long-lived tensor of each checkpoint boundary in CPU memory, and OffloadStreamAdamW turns the serial CPU Adam update of optimizer offload into a bucket pipeline. All four change only the order and granularity of computation and data movement, so the loss and gradients stay exact. In matched component tests, they cut the MoE dispatch peak by up to 59.3% without losing throughput, the vocabulary projection peak by 86.6%, and the offloaded optimizer step by 2.05times faster. Composed on MoE models from 120B to 667B parameters, they train at 1M context length, 8--32times the reach of a tuned FSDP2 baseline, and up to 10.4times its throughput.
Original Article
View Cached Full Text

Cached at: 09/17/26, 10:53 AM

Paper page - Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training

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

Abstract

TrainingaMixture-of-Experts(MoE)modelatlongcontextorlargebatchsizefailsassoonasanyonecomponent’speakallocationexceedsdevicememory,sothetargetiseverypeakatonce,nottheaveragefootprint.Fourareleftunboundedbytheparallelismplansincommonuse,andeachgrowsdifferently:expertdispatchwiththeroutingmatrix,thevocabularyprojectionwithtokenstimesvocabulary,gradientcheckpointboundarieswithdepthtimessequencelength,andoptimizerstatewithparametercount.Whichonerunsoutfirstchangeswiththemodel,thecontextlength,andthedevicecount,soloweringthelargestonlyexposesthenext.WeboundallfourwithscheduleswhoseGPUworkingsetisfixedatlaunch:PipelinedLLEPextendsleast-loadedexpertparallelismwithacaponthetokenseachsourcecontributestoadispatchchunk,Ring-DTPcirculatesactivationsorweightshardsaroundaringatthevocabularyprojectionandfoldseachblockoflogitsintoanonlinelog-sum-exp,Selectivecheckpointoffload(SCO)keepstheonelong-livedtensorofeachcheckpointboundaryinCPUmemory,andOffloadStreamAdamWturnstheserialCPUAdamupdateofoptimizeroffloadintoabucketpipeline.Allfourchangeonlytheorderandgranularityofcomputationanddatamovement,sothelossandgradientsstayexact.Inmatchedcomponenttests,theycuttheMoEdispatchpeakbyupto59.3%withoutlosingthroughput,thevocabularyprojectionpeakby86.6%,andtheoffloadedoptimizerstepby2.05timesfaster.ComposedonMoEmodelsfrom120Bto667Bparameters,theytrainat1Mcontextlength,8--32timesthereachofatunedFSDP2baseline,andupto10.4timesitsthroughput.

View arXiv pageView PDFAdd to collection

Get this paper in your agent:

hf papers read 2609\.14306

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/2609.14306 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

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

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2609.14306 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

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.

Understanding Is Done Early: A Depth Division of Labor in Large Language Models and Its Use for Unbounded-Context Memory

arXiv cs.CL

This paper introduces CoMem, a method that exploits the depth-wise division of labor in LLMs to cache intermediate residual tensors and recompute only upper layers for retrieval, enabling bounded read compute and memory independent of stored-context length. Evaluated on Qwen3-8B, CoMem achieves strong long-context performance with significant memory savings and prefill speedups.