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

> Four bounded-streaming operators cut MoE memory peaks by up to 86.6%, enabling exact 1M-token training 8–32× longer than FSDP2 with up to 10.4× throughput.

- **Source:** [arXiv](https://arxiv.org/abs/2609.14306)
- **Published:** 2026-09-19
- **Permalink:** https://picx.dev/p/RbUoz4
- **Whiteboard:** https://picx.dev/p/RbUoz4/image

## Summary

## Summary (Overview)

- **Problem**: Training Mixture-of-Experts (MoE) models at long context or large batch size fails when any single component's peak memory allocation exceeds device memory. Four memory peaks are left unbounded by common parallelism plans: expert dispatch, vocabulary projection, gradient checkpoint boundaries, and optimizer state.
- **Solution**: The paper introduces four "bounded-streaming" operators that replace the schedule of each problematic component, keeping the GPU working set fixed at launch time without changing the model, loss, or gradients (all exact).
- **Key results**: The operators cut MoE dispatch peak by up to 59.3%, vocabulary projection peak by 86.6%, and speed up the offloaded optimizer step by 2.05×. Composed on MoE models from 120B to 667B parameters, the stack trains at 1M context length—8–32× the reach of a tuned FSDP2 baseline—with up to 10.4× its throughput.
- **Core principle**: All four operators change only the *order and granularity* of computation and data movement, leaving the mathematical results bitwise identical to the standard implementations they replace.

## Introduction and Theoretical Foundation

### Background and Motivation

MoE models are often trained at long context or large batch size at large scale. When adding GPUs is not an option, memory is traded for time via techniques like gradient checkpointing (recomputation), state sharding (communication), and optimizer offload (serialized host updates). A training step fails when *any* component's peak memory exceeds device capacity—so a plan that reduces three bottlenecks but leaves the fourth unbounded buys nothing at the point where the workload grows.

### The Four Unbounded Live Sets

Peak memory is a maximum over **live sets**—groups of tensors that must be resident simultaneously. The paper identifies four live sets left unbounded by common parallelism plans:

| Live set | Unbounded size driver | Grows with |
|---|---|---|
| Expert dispatch | Routing matrix realized size | Batch size, top-k, routing imbalance |
| Vocabulary projection | $N \times V$ logit tensor | Tokens × vocabulary size |
| Checkpoint boundaries | Retained layer inputs | Depth × sequence length |
| Offloaded AdamW state | Serial host update | Parameter count |

These grow at different rates, so the dominant bottleneck depends on configuration: logits dominate at large vocabulary and long context, MoE dispatch at high routing imbalance, checkpointing boundaries at high depth, and optimizer state at large parameter count on few devices.

### Design Requirements

Each operator must:
1. Return the **exact** forward values and gradients of the implementation it replaces
2. Only change the *order* of computation—results must be identical
3. Have a memory bound that **survives checkpoint recomputation** (ruling out retaining anything sized by the routing matrix)

This excludes low-rank adapters, quantized state, and approximate routing or attention—all memory and throughput figures are comparable with standard full-parameter BF16 training.

## Methodology

### 1. PipelinedLLEP: Receiver-Bounded Expert Dispatch

**Problem**: Under expert parallelism, each rank owns only some experts. The number of routes $R_d$ to a destination rank $d$ determines its dispatch buffer size, which grows with batch size, routing top-k $k$, and routing imbalance. Least-loaded expert parallelism (LLEP) removes the imbalance cost but still allocates the entire routed batch at once.

**Method**: PipelinedLLEP splits the batch into $K$ chunks with a maximum token budget $c$ per source rank. The number of chunks is:

$$K = \min(\lceil N / c \rceil, K_{\text{max}}), \qquad c_{\text{eff}} = \lceil N / K \rceil, \qquad R_d^{(i)} \le E_p k c_{\text{eff}} \tag{1}$$

where $N$ is tokens per rank, $k$ is the router's top-k, $E_p$ is the expert-parallel degree, and $K_{\text{max}}$ is a predetermined chunk limit. The bound contains only $E_p$, $k$, and $c_{\text{eff}}$—none depend on router behavior, so a skewed router cannot push any buffer past the limit.

**Key implementation details**:
- **Nested gradient checkpointing**: Each chunk's expert matmul is wrapped in a reentrant gradient checkpoint, nested inside the decoder layer's non-reentrant checkpoint, freeing one chunk's intermediates before the next chunk allocates.
- **Strided chunk membership**: Chunk $i$ gets positions $i, i+K, i+2K, \ldots$ to spread each chunk's routes over destinations (consecutive chunks would send nearly all routes to one rank, reintroducing imbalance).

**Memory bound**: Buffers take at most $2E_p k c H b$ bytes (hidden width $H$, $b$ bytes per value).

### 2. Ring-DTP: Exact Vocabulary Projection over Distinct Batches

**Problem**: The vocabulary projection materializes a tokens-by-vocabulary logit tensor, which dominates memory when both factors are large. Fused kernels avoid this on a single device; Megatron's vocabulary-parallel requires replicated batches, reducing effective batch size.

**Method**: Each rank $r$ holds a distinct local batch $X_r \in \mathbb{R}^{N \times H}$ and owns vocabulary interval $\mathcal{V}_r$ of size $V/P$. Over $P$ ring-like rounds, every batch $X_j$ meets every weight shard $W_r$ exactly once. Each round forms only the logit strip $Y_{j,r} = X_j W_r \in \mathbb{R}^{N \times V/P}$, folds its partial normalizer into a running per-token state $S_j = (\bar{m}_j, z_j, y_{t_j})$, and releases the strip.

**Peak logit memory**:

$$M_{\text{logit}}^{\text{Ring-DTP}} = \mathcal{O}(N V b / P) \tag{2}$$

**Online softmax recurrence** (standard form from FlashAttention):

$$m' = \max(m, \max_v y_v), \qquad z' = e^{m - m'} z + \sum_v e^{y_v - m'} \tag{3}$$

**Dynamic payload choice**: Either activations $X_j$ or weight shards $W_r$ can travel. Ring-DTP moves weights when $N > V/P$ (since bytes per hop are $\mathcal{O}(NH)$ for activations vs. $\mathcal{O}(HV/P)$ for weights). The move-weights schedule needs no return hop in forward.

### 3. SCO: Selective Checkpoint Offload

**Problem**: Under gradient checkpointing, one tensor per checkpointed layer (its input) lives on device from forward until that layer is recomputed in backward.

**Method**: SCO walks checkpointed layers in forward order and offloads each boundary that fits in a host budget (set $\mathcal{S}$). Backward visits layers in reverse; while layer $\ell$ recomputes from $h_\ell$, the next boundary $h_{\ell-1}$ is copied back asynchronously. At most two restored boundaries are live on device at once.

**Memory bounds**:

$$M_{\text{HBM}}^{\text{SCO}} \leq 2 N_{\text{max}} H b, \qquad M_{\text{host}} \leq |\mathcal{S}| N_{\text{max}} H b \tag{4}$$

### 4. OffloadStreamAdamW: Bounded GPU Updates over Host CPU State

**Problem**: CPU AdamW updates are slow; during the update the GPU is idle and (having released activations) largely empty.

**Method**: Partition each rank's state into buckets of at most $\beta$ parameters, rotated through $s$ staging slots and three streams: (i) host-to-device transfer, (ii) fused GPU AdamW update + bf16 working-weight refresh, (iii) write-back of updated fp32 state.

**Timing and memory**:

$$T_{\text{stream}} = G \max(T_{\text{H2D}}, T_{\text{update}}, T_{\text{D2H}}) + \mathcal{O}(T_{\text{H2D}} + T_{\text{update}} + T_{\text{D2H}}), \quad M_{\text{stage}} = \mathcal{O}(s\beta) \tag{5}$$

The floor is the host-link round trip of $12\Theta/W$ bytes; the goal is to hide update and write-back behind that transfer.

### Composition: Per-Rank Budget

The integrated per-rank device budget is:

$$M_{\text{peak}}^{\text{integrated}} = \underbrace{\frac{4\Theta}{W}}_{\text{weights, grads}} + \underbrace{\mathcal{O}(s\beta)}_{\text{opt. staging}} + \underbrace{\mathcal{O}(NHb)}_{\text{attention}} + \underbrace{\mathcal{O}(E_p q H b)}_{\text{dispatch}} + \underbrace{\mathcal{O}(NVb/P)}_{\text{strips}} + \underbrace{\mathcal{O}(NHb)}_{\text{boundaries}} \tag{6}$$

Every term is fixed by model or launch configuration once $N \leq N_{\text{max}}$—so feasibility can be checked *before* launch.

## Empirical Validation / Results

### Isolated Component Benchmarks (8× H200 GPUs)

**PipelinedLLEP** (65K tokens/rank, $H=7168$, top-8, $c=6554$, $K=10$):

| Shape | LLEP (GiB) | PipelinedLLEP (GiB) | Peak saved | Speedup vs. LLEP |
|---|---|---|---|---|
| 65K tokens/rank, H=7168, top-8 | 52.7–53.0 | 21.4–22.9 | 56.9–59.3% | 1.01–1.10× |

**Ring-DTP** ($P=8$, $H=7168$, $V=200{,}000$, FP32):

| N | Schedule | Standard (GiB) | Ring-DTP (GiB) | Saved | Latency cost |
|---|---|---|---|---|---|
| 16,384 | move-activations | 42.462 | 7.322 | 82.8% | +5.1% |
| 32,768 | move-weights | 79.521 | 10.622 | 86.6% | +4.4% |

**SCO** (gpt-oss-20b, 8×H200, global batch 556,432 tokens):

| Host budget | Boundaries offloaded | Peak HBM (GiB) | Node RAM (GiB) | Throughput (tok/s/GPU) | Largest batch |
|---|---|---|---|---|---|
| Off | 0/47 | 139.790 | 402.517 | 2,641 | 557,056 |
| 8 GiB | 21/47 | 133.546 | 487.990 | 2,692 | 589,824 |
| 16 GiB | 42/47 | 125.677 | 573.473 | 2,687 | 622,592 |
| Full | 47/47 | 123.728 | 593.410 | 2,687 | 655,360 |

**OffloadStreamAdamW** (gpt-oss-20b, 8×H200):

| Configuration | Step time (s) | Speedup | Staging (GiB/GPU) |
|---|---|---|---|
| CPU Adam (ZeRO-Offload) | 3.95 | — | 0 |
| Streamed, 2 slots | 1.93 | 2.05× | 4.234 |

### End-to-End Integration (MoP rank layout, 120B/241B/667B models, 16/32/64 H200s)

- **Context reach**: The composed stack trains at **1M tokens** at all three scales; FSDP2-best exhausts memory past 128K (120B), 32K (241B), and 64K (667B)—an **8–32× reach** improvement.
- **Throughput**: The composed stack is **7.6× faster** at 128K on 120B and **10.4× faster** at 64K on 667B than FSDP2-best at its longest feasible length.
- **Floating-point rate**: Per-GPU rate roughly doubles between 128K and 1M tokens (from 91–110 to 213–233 TFLOP/s), while FSDP2-best stays below 40 TFLOP/s.
- **Largest global batch**: 1.5M (120B), 1.8M (241B), and 3M (667B) distinct tokens per forward/backward pass—**12×, 7×, and 3×** the baseline's.
- **Training quality**: Loss and gradients remain exact (unchanged from standard training).

## Theoretical and Practical Implications

### Theoretical Significance

1. **Closed-form memory budget**: Equation (6) provides a per-rank budget where every term is fixed by launch configuration—enabling *a priori* feasibility checking before launching a training run.
2. **Exactness without approximation**: All four operators achieve their bounds without low-rank adapters, quantized state, or approximate routing/attention—extending the Pareto frontier of memory-efficient training without sacrificing fidelity.
3. **Orthogonal bounds**: The four operators bound *disjoint* live sets and can be enabled independently, so a training run pays only for the peaks it actually has.

### Practical Implications

1. **Long-context MoE training**: At 1M context length, the logit tensor alone would exceed HBM; Ring-DTP's move-weights branch scales best with context (per-hop payload $\mathcal{O}(HV/P)$ does not grow with token count).
2. **Interconnect dependence**: All-to-all and ring traffic assume fast interconnect (NVLink intra-node measured here); slower fabrics could change the launch-count trade.
3. **Host memory requirements**: Checkpoint and optimizer streaming require substantial host RAM (2 TB used here); nodes with less host memory may be constrained.
4. **Hyperparameter selection**: The token budget $c$ and bucket size $\beta$ are selected from measured curves, not minimized—topology sensitivity and automatic selection of $(D, E_p, P, c, \beta)$ remain open.

## Conclusion

MoE training at long context or large batch size fails for four unrelated reasons whose relative heights shift with configuration. The paper provides a schedule for each:

- **PipelinedLLEP**: A per-source token budget for expert dispatch
- **Ring-DTP**: A ring of vocabulary meetings for the projection head  
- **SCO**: A host-budget checkpoint-boundary offload
- **OffloadStreamAdamW**: A bucket pipeline for the offloaded optimizer update

Together they yield a closed-form per-rank budget where every term is fixed once the token ceiling is set. Matched component tests show each bound holding at a measurable throughput cost. Composed inside the MoP rank layout at 120B–667B scale, the stack trains at one-million-token context where a tuned FSDP2 baseline runs out of memory between 32K and 128K, spending the memory it saves on larger batches as readily as on longer contexts.

**Future directions**: Topology sensitivity analysis and automatic selection of parallelism degrees and streaming hyperparameters $(D, E_p, P, c, \beta)$ remain open problems.

---

_Markdown view of https://picx.dev/p/RbUoz4, served by PicX — AI-generated visual whiteboard summaries of research papers._