Long-context MoE training rarely dies from "average memory being too small" — it dies when any single component's peak allocation blows past device memory first. A preprint from Salesforce AI Research (Shrey Pandit, Xuan-Phi Nguyen, Yiran Zhao, Shafiq Joty), filed September 13, takes the problem down to component level: four memory peaks in common parallelism plans are left unbounded, each grows on a different curve, and which one runs out first depends on model size, context length, and device count. Lower the largest, and the next one is simply exposed.

Four peaks, four growth curves

The paper names them explicitly: expert dispatch grows with the routing matrix; the vocabulary projection grows with tokens times vocabulary size; gradient checkpoint boundaries grow with depth times sequence length; optimizer state grows with parameter count. None of the four has an upper bound under the parallelism plans in common use, and any one exceeding device memory kills the run. Hence the core claim: the target is every peak at once, not the average footprint.

Four schedules with a fixed GPU working set

Against the four peaks, the paper fields four schedules whose shared property is that the 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, overlapping communication with computation chunk-wise;
  • 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, never materializing the full projection at once;
  • SCO (Selective checkpoint offload) keeps the one long-lived tensor of each checkpoint boundary in CPU memory;
  • OffloadStreamAdamW turns the serial CPU Adam update of optimizer offload into a bucket pipeline.

The binding design constraint: all four change only the order and granularity of computation and data movement, so loss and gradients stay exact — this is scheduling, not approximation.

The numbers

In matched component tests: the MoE dispatch peak drops by up to 59.3% with no throughput loss; the vocabulary projection peak drops 86.6%; the offloaded optimizer step gets 2.05x faster. Composed on MoE models from 120B to 667B parameters, the stack trains at 1M context length — 8 to 32 times the reach of a tuned FSDP2 baseline, and up to 10.4x its throughput.

Caveats

Three things worth flagging. First, this is a v1 preprint from September 13; no external replication exists yet. Second, the paper page ships no code repository link — the engineering details currently live only in the PDF. Third, the 8-32x reach figure is relative to a tuned FSDP2 baseline; it is a ratio, not an absolute capability claim. For production users: wait for code or third-party reproduction.

So what

The value here is not any single benchmark number. It is the decomposition of "long-context MoE training runs out of memory" from a vague engineering complaint into four separately governable peaks, plus a demonstration that scheduling-layer changes can be mathematically exact. In the long-context race, attention innovations (sparse, linear, hybrid) take the headlines — but boring training infrastructure is often the plank that decides who can actually run 1M context.

Reference: arXiv:2609.14306