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

AI summary

Overview

Research area: Distributed systems and systems-for-machine-learning (cs.DC) — specifically memory management for large-scale distributed training of Mixture-of-Experts (MoE) models.

Technical level: Advanced. The paper assumes familiarity with expert parallelism, ZeRO/FSDP sharding, gradient checkpointing, tensor parallelism, and AdamW optimizer state.

Scope: The paper introduces four "bounded-streaming" operators that cap four different memory peaks in long-context MoE training, each of which is left unbounded by common parallelism plans, and evaluates them individually and composed on MoE models from 120B to 667B parameters.

What This Paper Is About

Training an MoE model at long context or large batch size fails as soon as any single component's peak memory allocation exceeds device memory — so the target is every peak at once, not the average footprint. The authors identify four live sets of tensors that common parallelism plans leave unbounded and that grow at different rates with the workload: expert dispatch (grows with the routing matrix), the vocabulary projection (grows with tokens × vocabulary), retained gradient checkpoint boundaries (grows with depth × sequence length), and optimizer state (grows with parameter count). The goal is to give each of the four a bound set at launch time, so that a workload can be checked for feasibility before it is run, while keeping the loss and gradients exact.

Key Contributions

1. PipelinedLLEP, an extension of least-loaded expert parallelism (LLEP) that delivers routes in chunks and caps how many tokens any source rank may place in a dispatch chunk, so receiver buffers follow a launch-time cap rather than the realized routing matrix. The number of chunks is derived from the token budget, and chunk membership uses a strided order so that a chunk's routes spread across destinations rather than concentrating on one.

2. Ring-DTP (ring data-tensor-parallel vocabulary projection), which lets each rank keep its own distinct batch while circulating either activations or weight shards around a ring until every batch has met every shard, folding each block of logits into an online log-sum-exp so the full tokens-by-vocabulary tensor is never formed.

3. SCO (selective checkpoint offload), which moves a host-budget subset of gradient checkpoint boundaries into pinned host memory and prefetches them back one layer before recomputation, and OffloadStreamAdamW, which turns the serial CPU AdamW update of optimizer offload into a bucket pipeline that uses the otherwise idle GPU as the update engine.

4. A composed per-rank budget inside the Mixture-of-Parallelisms (MoP) rank layout, in which every term is fixed by model or launch configuration once the per-rank token ceiling is set, together with matched component benchmarks and end-to-end results at 120B, 241B, and 667B parameters.

Main Findings

• PipelinedLLEP cuts the MoE dispatch peak by up to 59.3%. In the isolated expert-dispatch benchmark at 65,536 tokens per rank, H = 7168, I = 2048, 128 experts, top-8, and c = 6554 (K = 10), LLEP peaks at 52.7–53.0 GiB while PipelinedLLEP peaks at 21.4–22.9 GiB, a saving of 56.9–59.3% at 1.01–1.10× LLEP's speed. Standard expert parallelism runs out of memory at this shape as soon as the router is severely skewed.

• Chunk membership matters. With consecutive (rather than strided) chunks at a 65K-token batch, one destination receives routes at 6.0× the average, while strided chunks stay within 2.4× the average — worth up to 1.35× in speed and up to 1.37 GiB of peak memory. Striding overall runs 1.03–1.35× faster and saves up to 1.37 GiB of peak memory.

• Ring-DTP removes 86.6% of the vocabulary projection peak for under 5% more time. At P = 8, H = 7,168, V = 200,000 in FP32, with N = 16,384 (move-activations branch, since N falls below V/P = 25,000), the standard schedule peaks at 42.462 GiB versus 7.322 GiB for Ring-DTP (82.8% saved) at 1052.2 ms versus 1001.2 ms (+5.1%). With N = 32,768 (move-weights branch), the standard peaks at 79.521 GiB versus 10.622 GiB (86.6% saved) at 2155.8 ms versus 2065.2 ms (+4.4%). Doubling N nearly doubles the baseline peak (42.5 to 79.5 GiB) while Ring-DTP moves only from 7.3 to 10.6 GiB, so the saving widens as context grows.

• SCO trades HBM for host RAM monotonically. On gpt-oss-20b with 8xH200 GPUs and sequence and expert parallelism of degree eight, a configured global batch of 556432 tokens, the device peak falls from 139.790 GiB (0/47 boundaries offloaded) to 133.546 GiB at an 8 GiB host budget (21/47), 125.677 GiB at 16 GiB (42/47), and 123.728 GiB with full offload (47/47). Node RAM rises from 402.517 to 593.410 GiB over the same sweep. Throughput moves 1.9% across settings with no monotone trend, while saving 17.65% of HBM consumption; measured against the largest batch that runs without an out-of-memory error, the gain is 35.71%. The largest 10-step batch rises from 557,056 to 655,360 tokens.

• OffloadStreamAdamW runs the offloaded optimizer step 2.05× faster. On gpt-oss-20b with 8xH200 GPUs, against the AVX-vectorized CPU Adam kernel of ZeRO-Offload with the same fp32 master weights and moments on the host, the step falls from 3.95 s to 1.93 s with 4.234 GiB/GPU of staging. Two slots suffice; increasing staging slots leaves step time unchanged, consistent with the pipeline being transfer-bound at the 12Θ/W host-link floor.

• End-to-end, the composed stack reaches 8–32× the context of a tuned FSDP2 baseline. On MoE models of 120B, 241B, and 667B parameters at 16, 32, and 64 H200 GPUs, FSDP2-best exhausts device memory past 128K, 32K, and 64K tokens respectively, while the composed stack trains at one million tokens at all three scales. It is also faster at the longest length FSDP2-best reaches, by 7.6× at 128K on 120B and 10.4× at 64K on 667B. At 241B, FSDP2-best stops at 32K, four times below the shortest configuration of the composed stack.

• Floating-point rate rises with context. Per-GPU rate roughly doubles between 128K and 1M tokens at all three scales, from 91–110 to 213–233 TFLOP/s, while FSDP2-best stays below 40 TFLOP/s anywhere in its sweep. At those lengths the composed stack delivers 2.2× the baseline's rate at 667B scale. Token throughput falls with context at the same time, because attention work per token grows with context.

• Larger batches as well as longer contexts. The largest global batch that runs is 1.5M, 1.8M, and 3M distinct tokens per forward and backward pass — 12×, 7×, and 3× the baseline's — and the longest context that runs is 1M tokens at every scale.

• Exactness is preserved. All four operators change only the order and granularity of computation and data movement. No low-rank adapters, no quantized state, and no approximate routing or attention are used, so memory and throughput figures are directly comparable with standard full-parameter BF16 training, and training quality is reported as unchanged.

Methodology in Plain English

The authors start from a diagnosis rather than a single fix: they observe that a training step dies when any one component's peak allocation exceeds device memory, and that which component runs out first depends on the model, the context length, and the device count. Lowering only the largest peak merely exposes the next one. They therefore enumerate four "live sets" — groups of tensors that must be resident simultaneously — and hold each of them to the same three requirements: the operator must return the same forward values and gradients as the implementation it replaces, it may only change the order of computation, and its bound must survive checkpoint recomputation (which rules out retaining anything sized by the routing matrix).

Each operator is then built by taking one established component and replacing its schedule:

• For expert dispatch, the key insight is that a destination's buffer size is set by how many routes it receives, which the router decides. If each source rank is instead limited to putting a fixed cap of its own tokens into any one chunk, then no destination can receive more than the expert-parallel degree times the routing degree times that cap — a number that contains no dependence on router behavior. The chunks are then run through a three-stage pipeline (dispatch, grouped expert matmul, combine), and each chunk's expert matmul is wrapped in a reentrant gradient checkpoint nested inside the decoder layer's non-reentrant checkpoint, so one chunk's intermediates are freed before the next allocates. Because the chunking is applied to a relocation plan that groups routes by destination, tokens are assigned to chunks in strided order to keep each chunk balanced.

• For the vocabulary projection, the observation is that cross-entropy needs only a running maximum, an exponential sum, and the target logit per token — three scalars independent of vocabulary size. Each rank holds a distinct batch and a distinct vocabulary shard; over P ring rounds, every batch meets every shard exactly once, producing one logit strip that is folded into the per-token state and then released. The system dynamically chooses whether to move activations or weights, moving weights exactly when N > V/P, and picks the smallest P for which the strip fits.

• For checkpoint boundaries, each checkpointed layer's input normally lives on the device from forward until recomputation. SCO walks layers in forward order and offloads each boundary that still fits in the host budget, then in backward prefetches the next boundary on a separate stream while the current layer recomputes, so at most two restored boundaries are live at once.

• For the optimizer, offload normally means a slow serialized CPU update during which the GPU sits idle. OffloadStreamAdamW partitions each rank's state into buckets of at most β parameters and rotates them through staging slots and three streams (host-to-device transfer, fused GPU AdamW update plus bf16 working-weight refresh, and write-back), aiming to hide the update and write-back behind the host-link round trip of 12Θ/W bytes.

Each operator is measured in isolation against the baseline that solves the same problem, on identical input, on one eight-H200 node. The four are then composed inside the MoP rank layout, which assigns ZeRO-3 for dense weights, a sequence-to-head all-to-all for attention, and expert-parallel placement for expert weights over overlapping sub-groups of a single rank set, and compared against FSDP2-best — the highest-throughput configuration from a sweep over FSDP2 combined with expert, context, and tensor parallelism — on the same model at the same GPU count.

Why This Matters

Impact on research. The paper reframes long-context MoE training as a problem of bounding every peak rather than optimizing the largest one, and shows that each bound can be derived from launch configuration rather than realized runtime behavior. That distinction matters because a bound that depends on the routing matrix the router draws carries no guarantee for the next routing matrix. The per-rank budget equation makes a configuration checkable for feasibility before launch, and because the operators change only schedule and data movement — not numerics, parameterization, optimizer, precision, or approximation — the results are directly comparable to standard full-parameter BF16 training. Two separate savings are explicitly disentangled: host offload supplies the reduction from 16Θ/W to 4Θ/W of persistent state, while OffloadStreamAdamW reduces the time that offload costs.

Real-world applications.

• Pre-training and continued pre-training of large sparse MoE language models at very long context (the paper demonstrates one-million-token training at 120B, 241B, and 667B parameters).
• Increasing effective batch size in distributed training runs, which the paper shows is the other axis the saved memory can be spent on — up to 3M distinct tokens per forward and backward pass at the 667B scale.
• Retrofitting memory headroom onto existing fixed-size GPU clusters where adding GPUs "is rarely the available remedy" and memory must be traded for time.
• Serving as a drop-in memory-bounding layer within existing sharding frameworks (FSDP2, ZeRO-3, Megatron-style tensor parallelism, sequence parallelism) rather than requiring a new training stack.

Industry relevance. The work comes from Salesforce AI Research and is aimed at practitioners who train large MoE models on memory-constrained clusters with fast interconnect. Its value proposition is concrete: a tuned FSDP2 baseline runs out of memory between 32K and 128K tokens at these scales, while the composed stack trains at 1M, at up to 10.4× the throughput of the best FSDP2 configuration at the longest length FSDP2 reaches, and with up to 12× the largest global batch. Because the operators are independently enableable, a training run pays only for the peaks it actually has.

Future Directions

• Interconnect sensitivity. All-to-all and ring traffic assumes a fast interconnect; the paper notes that the launch-count trade which makes additional chunks inexpensive on the intra-node NVLink measured here could differ on a slower fabric.
• Token-budget selection. A smaller token budget tightens the receiver bound and raises the chunk count by the same factor, and forward time is flat in the chunk count only over part of that range, so c is selected from a measured curve rather than minimized.
• Host memory and link bandwidth limits. Checkpoint and optimizer streaming spend host capacity and link bandwidth, which the authors say binds on nodes with less host memory than the 2 TB used in their setup.
• Automatic configuration search. Topology sensitivity and automatic selection of (D, E_p, P, c, β) remain open.

Target Audience

This paper benefits most researchers and engineers who build or operate distributed training systems for large sparse models — particularly those working on memory-efficient training, expert parallelism, long-context scaling, and optimizer offloading. It is also relevant to practitioners who need to decide whether a given model, context length, and device count combination will fit before launching a run, and to systems researchers interested in bounding techniques that preserve exactness rather than trading numerical fidelity for memory. A working knowledge of sharding strategies, gradient checkpointing, and MoE routing is required; the paper is not an introductory treatment.

Authors’ abstract

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.05\times$ faster. Composed on MoE models from 120B to 667B parameters, they train at 1M context length, $8$--$32\times$ the reach of a tuned FSDP2 baseline, and up to $10.4\times$ its throughput.