Salesforce bounds four MoE memory peaks and trains 120B–667B models at 1M context

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

Shrey Pandit, Xuan-Phi Nguyen, Yiran Zhao, Shafiq Joty

cs.DC

2026-09-13

Four exact schedules bound MoE dispatch, vocab, checkpoints, and optimizer peaks at launch. 120B–667B models train at 1M context, 8–32× a tuned FSDP2 baseline, with exact loss.

What problem this solves

MoE training at long context or large batch dies when any component's peak exceeds HBM. Average footprint is the wrong target. Common parallelism shards persistent weights and still leaves four live sets that grow with the workload, each on a different axis: expert dispatch with the routing matrix, the vocabulary projection with tokens times vocabulary, checkpoint boundaries with depth times sequence length, and AdamW state with parameter count. Shrink the current largest and the next one appears. Salesforce AI Research bounds all four at once, with each operator separately switchable.

Method

All four change only order and granularity. The model, precision, optimizer, and loss stay put, so the step remains exact full-parameter BF16.

• PipelinedLLEP, on top of least-loaded expert parallelism, caps how many tokens each source rank may put in a dispatch chunk and pipelines dispatch, expert matmul, and combine. Tokens inside a chunk are taken in stride, so contiguous positions do not dump an entire chunk on one destination. Receiver buffers are fixed at launch by Ep, k, and c; skew lengthens the schedule instead of the buffers. Each chunk's expert compute sits in a reentrant checkpoint so the backward peak obeys the same cap.
• Ring-DTP lets every rank keep a distinct local batch. Vocabulary weights are column-sharded; activations or shards circulate on a ring until every batch has met every shard. Logits fold into an online log-sum-exp per block, so the N×V tensor is never formed. Unlike Megatron vocabulary parallel, ranks do not replicate the same batch.
• SCO: one input tensor per checkpointed layer still lives from forward until recompute. A host-budget subset is offloaded to CPU and prefetched one layer ahead.
• OffloadStreamAdamW keeps optimizer state on the host, streams parameter buckets through a fused GPU AdamW, and writes back, overlapping H2D, update, and D2H. Device staging is O(sβ) rather than a resident 12Θ/W.

The four live sets are disjoint. Placed in a Mixture-of-Parallelisms rank layout (ZeRO-3 on dense weights, sequence parallel on attention, expert parallel on experts), they yield a per-rank budget that can be checked before launch.

Results

Component tests use eight H200s and identical inputs.

Strided chunks run 1.03–1.35× faster than contiguous ones and save up to 1.37 GiB. Composed on 120B, 241B, and 667B MoEs at 16/32/64 H200s, against the best FSDP2 mix at the same GPU count: all three scales train at 1M context, 8–32× the reach. FSDP2 OOMs past 128K (120B), 32K (241B), and 64K (667B). At the longest length FSDP2 still fits, throughput is 7.6× at 120B/128K and 10.4× at 667B/64K. Largest global batches are 1.5M, 1.8M, and 3M distinct tokens, 12×, 7×, and 3× the baseline. Per-GPU FLOPs rise from 91–110 TFLOP/s at 128K to 213–233 at 1M because attention work per token grows; FSDP2 stays below 40. An appendix reports unchanged training quality.

Why it matters

This is a scheduling paper for stacks that already use FSDP, expert parallel, checkpointing, and optimizer offload, not a new MoE architecture. The relative height of the four peaks moves with model, context, and device count, so a usable system needs a bound on every term, known at launch, rather than a large saving on one. Exactness keeps the numbers comparable to standard full-parameter BF16. If GPU count cannot grow, the composition turns memory into streamed time at 1M context on 100B–600B-class MoEs.

Limitations

All-to-all and ring traffic assume a fast fabric; the measurements are intra-node NVLink, and extra chunks may cost more on a slower network. A tighter token cap raises the chunk count, and forward time is flat in that count only over part of the range, so c is chosen from a measured curve. Checkpoint and optimizer streaming spend host RAM and host-link bandwidth; the nodes here have 2 TB of host memory. Automatic selection of (D, Ep, P, c, β) is open. The end-to-end baseline is the best point in an FSDP2 sweep, not every production MoE stack. "Unchanged training quality" is an appendix claim; the main text has no loss curves.

Terms

• Live set: tensors that must be resident in HBM at the same time; peak is their maximum sum.
• LLEP: expert parallelism that moves overflow routes onto idle ranks without changing expert choice.
• Online log-sum-exp: running max and exponential sum over vocabulary blocks, without materializing a full logit row.
• MoP: overlapping parallelism degrees on one rank set, rather than a multiplicative device grid.
• ZeRO-Offload: Adam state on CPU; the classic update runs serially on the host.

Source

• arXiv 2609.14306v1
• PDF
• Hugging Face

Related papers

• Orbit PAM cuts Cosmos DB nodes 35%, with 47 replica moves in 6 days
• Perplexity fabric-lib hits 400 Gbps on EFA and ConnectX, 1.3s for 1T-param weight sync
• Three-layer Redis cache hits 89.3% on 8-vCPU LLM search, not at query level
• OD-MoE drops the expert cache and still hits 75% of full-GPU Mixtral decode with <1 GB per worker
• SP-MoE prefetches MoE experts during speculative drafting, cutting TPOT by up to 3.5x
• Azure's Slasher sheds 15% power with near-zero hit to high-priority cloud services

All paper explainers