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

• PipelinedLLEP chunks routes and caps each source’s tokens per chunk, making receiver buffers follow the cap while routing skew lengthens the schedule.Source: §1 · p. 2
• Ring-DTP circulates activations or weight shards so each rank’s distinct batch meets every shard without forming the large logits tensor.Source: §1 · p. 2
• SCO offloads selected checkpoint inputs to host memory, while OffloadStreamAdamW streams parameters to GPUs for updates and writes them back.Source: §1 · p. 2
• PipelinedLLEP cuts dispatch peak up to 59.3% versus LLEP without noticeable slowdown; Ring-DTP removes 86.6% of projection peak for under 5% more time.SCO changes throughput by under 2% and raises the largest completed batch by 17.7%; OffloadStreamAdamW makes the offloaded optimizer step 2.05× faster.Source: §1 · p. 2
• At one-million-token context, the composed stack trains 120B, 241B, and 667B MoE models at 8–32× tuned FSDP2 reach, up to 10.4× throughput, and up to 12× its largest global batch.Source: §1 · p. 2

• AdamW ordinarily requires 16Θ bytes for working weights, gradients, master weights, and moments; ZeRO-3, sequence parallelism, and expert parallelism shard or reduce other memory but do not bound the four targeted live sets.Source: §2 · p. 2
• The four targeted live sets are expert-dispatch routes, vocabulary projection, SCO checkpoint layers, and optimizer buckets, with parameters including R_d, E_p, k, c_eff, P, S, N_max, and β.Source: Table 1 · p. 3
• Each operator must return the replaced implementation’s forward values and gradients, change only computation order, and keep its memory bound through checkpoint recomputation.A forward-only bound is insufficient because a checkpointed step reaches its peak inside backward.Source: §2 · p. 3

3.1 PIPELINEDLLEP: RECEIVER-BOUNDED EXPERT DISPATCH

• 56.9–59.3% lower peak memory versus LLEP at 1.01–1.10× its speed.PipelinedLLEP caps tokens per sending-rank chunk and uses strided token selection to spread routes across destinations; inner reentrant checkpoints release each chunk’s intermediate tensors before the next allocation.Source: §3.1 · pp. 3–5, Figure 2 · p. 5

3.2 RING-DTP: EXACT VOCABULARY PROJECTION OVER DISTINCT BATCHES

• 82.8% and 86.6% lower vocabulary-projection peaks for at most 5.1% more time.Ring-DTP circulates distinct batches and vocabulary shards, folds each logit strip into a running per-token state, and releases it; backward recomputes gradients strip by strip without storing full logits.Source: §3.2 · pp. 5–6, Table 2 · p. 6

3.3 SCO: EXACT GRADIENT CHECKPOINTING BOUNDARY OFFLOAD

• 35.71% more tokens fit in the largest batch that completes; device peak falls monotonically as the host budget grows.SCO offloads selected checkpoint boundaries to pinned host memory and asynchronously prefetches the next boundary during backward; at most two restored boundaries are live on-device.Source: §3.3 · pp. 6–7, Table 3 · p. 7

3.4 OFFLOADSTREAMADAMW: BOUNDED GPU UPDATES OVER HOST CPU STATE

• 2.05× faster optimizer step than the CPU AdamW kernel replaced by OffloadStreamAdamW.It divides each rank’s optimizer state into buckets and overlaps host-to-device transfer, GPU AdamW updates, and write-back; two staging slots suffice to keep transfers busy.Source: §3.4 · p. 7, Table 4 · p. 8

3.5 COMPOSITION AND THE PER-RANK BUDGET

• Once N ≤ Nmax, every quantity in the combined per-rank budget is fixed by the model or launch configuration, allowing feasibility to be checked before launch.The four operators bound disjoint live sets and can be enabled independently; host offload reduces persistent state from 16Θ/W to 4Θ/W, while streaming reduces its time cost.Source: §3.5 · pp. 7–8

• At 120B, 241B, and 667B, the stack trains at 1M tokens (8–32× the reach), while FSDP2-best exhausts memory past 128K, 32K, and 64K, respectively.Training quality remains unchanged.Source: §4 · pp. 8, 10, Figure 6 · p. 9
• 7.6× at 128K on 120B and 10.4× at 64K on 667B: the composed stack is faster than FSDP2-best at the longest lengths that baseline reaches.Source: §4 · p. 8, Figure 6 · p. 9
• At 128K to 1M tokens, the stack’s per-GPU rate roughly doubles across all three scales, from 91–110 to 213–233 TFLOP/s; FSDP2-best stays below 40 throughout its sweep.At those same lengths, the stack delivers 2.2× the baseline’s rate at 667B scale.Source: §4 · pp. 8–9, Figure 6 · p. 9
• 1.5M, 1.8M, and 3M distinct tokens per forward and backward pass are the largest global batches that run, or 12×, 7×, and 3× the baseline’s, respectively; the longest context is 1M at every scale.The batch and context frontiers are reported in Figure 7.Source: Figure 7 · p. 9, §4 · p. 9

• PipelinedLLEP differs from dispatch methods that partition for bandwidth by bounding what each destination may receive with a per-source token budget, chunk membership policy, and nested checkpoint contract.The paper positions its operator definitions, bounds, exactness conditions, and mechanism-level evaluation as details not specified by parallelism layout alone.Source: §5 · p. 10
• Ring-DTP addresses distinct per-rank batches, unlike fused cross-entropy kernels that assume a batch resident on one device and vocabulary-parallel reduction that needs a replicated batch.Source: §5 · p. 10
• SCO offloads a host-budget subset of checkpoint boundaries and prefetches them one layer ahead, while quantization and block-shared learning rates shrink optimizer state by changing numerics, which this setting excludes.Source: §5 · p. 10

• Four schedules bound dispatch, vocabulary logits, checkpoint boundaries, and the offloaded optimizer update, with GPU working sets fixed by launch configuration.The bounds can be enabled separately, and the closed-form per-rank budget can be checked before launch.Source: §6 · p. 10
• Matched component tests show each bound holding at a measured throughput cost.Source: §6 · p. 10
• At 120B, 241B, and 667B parameters, the composed stack trains at one-million-token context, while tuned FSDP2 runs out of memory between 32K and 128K.The saved memory can be spent on a larger batch or a longer context.Source: §6 · p. 10

• All-to-all and ring traffic assume a fast interconnect; the launch-count trade measured on intra-node NVLink could differ on a slower fabric.Source: §7 · p. 10
• Reducing the token budget tightens the receiver bound but raises chunk count proportionally, and forward time stays flat only over part of that range, so c is chosen from a measured curve.Source: §7 · p. 11
• Checkpoint and optimizer streaming can be constrained by host capacity and link bandwidth on nodes with less host memory than the 2 TB used; topology sensitivity and automatic selection remain open.Source: §7 · p. 11

THE USE OF LARGE LANGUAGE MODELS

• Language models assisted with drafting, editing, caption tightening, and reference formatting; authors checked every claim, number, and citation against source measurements and cited papers.The authors state that the models did not generate research ideas, experimental designs, or results.Source: THE USE OF LARGE LANGUAGE MODELS · p. 11
• For implementation, language models helped with routine code editing and test scaffolding, while the authors designed, reviewed, and executed operators, correctness tests, and measurements.Source: THE USE OF LARGE LANGUAGE MODELS · p. 11
• The authors take full responsibility for the paper’s content.Source: THE USE OF LARGE LANGUAGE MODELS · p. 11

A EXPERIMENTAL SETUP

• Controlled component benchmarks use byte-identical inputs on one node with eight H200 GPUs, measuring operator-specific peaks and latencies.Dispatch latency is averaged over ten measured forward calls; peak memory uses maximum rank allocation for dispatch and projection, and warmed-up NVML samples for SCO.Source: A EXPERIMENTAL SETUP · p. 15
• In the end-to-end pre-training workload, models span 120B, 241B, and 667B parameters across 16, 32, and 64 GPUs, with FSDP2-best selected by a joint parallelism sweep.The models share an architecture and differ in depth; the sweep selects the highest-throughput configuration that fits at each scale.Source: A EXPERIMENTAL SETUP · p. 15
• MoP places sequence, expert, and vocabulary work in independently chosen overlapping rank groups, while dense weights are sharded across all ranks.The groups are views of the same rank set rather than factors of a device-grid equation.Source: A EXPERIMENTAL SETUP · p. 15
• At 128K context, composed-stack fine-tuning reaches 59.8% Avg@8 on AIME 2025, versus 59.6% for FSDP2 and 45.2% before fine-tuning.The comparison uses the same data, schedule, and hyperparameters for the composed stack and FSDP2 baseline.Source: A EXPERIMENTAL SETUP · p. 16

B.1 BACKWARD AND THE NESTED CHECKPOINT CONTRACT

• The PipelinedLLEP token partition is recomputable from N and K, and the reentrant inner checkpoint lets the outer checkpoint intercept expert-matmul tensors instead of retaining stacked weights across layers.A non-reentrant inner checkpoint would keep strong references outside the saved-tensor mechanism; checkpointing only the expert matmul also permits dispatch prefetch of the next chunk.Source: §B.1 · p. 16
• Only the expert-parameter gradient sum changes summation order across chunks; hidden-state and router gradients are not split, and at K = 1 expert gradients are bitwise identical to LLEP’s.The checkpointed step’s peak falls inside backward, so the memory bound must survive recomputation.Source: §B.1 · p. 16

B.2 ROUTING PROFILES

• Routing profiles vary hot-token percentage and hot-expert count; skew increases as the hot-token share rises or the hot set shrinks, with h ≥ k marking the qualitative boundary.Source: §B.2 · p. 17

B.3 PER-PROFILE MEASUREMENTS

• At 128 tokens, deriving K = 1 from the token budget runs at 1.60–2.96×, while a fixed ten chunks run at 0.45–0.48×.The fixed chunk count pays ten collective rounds for roughly one token per piece; the budget-derived choice avoids that overhead.Source: §B.3 · p. 17
• Strided membership runs at 1.03–1.35× contiguous membership’s speed and uses 0.03–1.37 GiB less peak memory.The comparison changes membership schedule quality while preserving token sets and route destinations.Source: §B.3 · p. 17, Table 8 · p. 19
• Strided membership reaches 12.5% of the Equation (1) bound across chunk sizes and profiles, whereas contiguous membership reaches 17.1% on the 50%/4 profile.The bound is respected everywhere and attained nowhere.Source: §B.3 · p. 17
• Across the routing-skew sweep, Table 6 reports mean forward latency and maximum peak allocation, with speedup and peak reduction measured against LLEP.The caption states that standard expert parallelism runs out of memory on every skewed profile at the 65K shape.Source: Table 6 · p. 18
• The buffer bound is a guarantee rather than a description of traffic, because buffers use realized split counts and the peak follows kceff.Source: §B.3 · p. 18
• 56.9–59.3% peak reduction persists with forward and backward together at the reference shape.Table 9 covers balanced and 80%/16 routing profiles, with chunk settings chosen for K = 10 at each shape.Source: §B.3 · p. 18, Table 9 · p. 19

B.5 TWO LATENCY MODES AT THE 65K SHAPE

• At 65K, the slow mode adds 12–56% and never subtracts, while fast-mode profile timings reproduce within 0.5% across passes.The same routing profile can appear in either mode across passes, so the mode is not a property of routing profile.Source: §B.5 · p. 19
• K = 10 is slow in 2 of 56 profile measurements, compared with 21 of 84 for K = 16.The authors use Kmax = 10 in part because the slow mode is rarer there.Source: §B.5 · p. 19
• The host-side launch pattern determines the mode; adding one host synchronization per chunk reproducibly moves the slow mode between profiles.The diagnostic synchronization changes nothing about the computation.Source: §B.5 · p. 19

C RING-DTP

• Each rank keeps a distinct local batch and visits every batch–vocabulary-shard pair once; backward replays meetings without storing full logits or logit gradients.The benchmark uses N = 16,384, H = 7,168, and V = 200,000 in FP32.Source: C RING-DTP · p. 19

D SCO

• SCO asynchronously offloads selected checkpoint boundaries to capped pinned host memory, restores them during reverse-layer recomputation, and leaves allocations exceeding the cap on device.At most two restores occupy the device at once.Source: D SCO · pp. 20–21

C RING-DTP

• Matched forward latency is 989.8 against 1035.7 ms at P = 4 and 1001.2 against 1052.2 ms at P = 8.These comparisons are for the same shape and are reported relative to the standard dense projection.Source: Table 10 · p. 21

E OFFLOADSTREAMADAMW

• OffloadStreamAdamW groups each parameter shard, master value, moments, final gradient, and refreshed working weight into a bucket for the update.Its fused kernel preserves the reference update order and bias correction.Source: E OFFLOADSTREAMADAMW · p. 22
• Global gradient clipping permits backward-time bucket updates only once the global norm exists; per-bucket and predictive clipping change the update rule and fall outside the exactness claim.A bucket’s gradient must also be final before its update begins.Source: E OFFLOADSTREAMADAMW · p. 22
• 1.93 s versus the 3.95 s AVX CPU Adam baseline: the selected 100M-parameter, two-slot OffloadStreamAdamW configuration averages 1.93 s per optimizer step.The measurements use gpt-oss-20b on eight H200 GPUs.Source: §E.1 · p. 23

E.1 BUCKET AND SLOT SWEEP

• The 100M-parameter bucket wins on optimizer-step time and staging, while two slots remain within 1.1% of three at every bucket size using one-third less staging memory.Step time is flat from two to five slots, while each additional slot adds one bucket’s staging allocation.Source: Table 14 · p. 23