Recompute converts retained-activation memory into repeated forward computation. Persistent training state—parameters, gradient buffers, and optimizer states—is fixed by the model and its sharding (Sections 3 and 4). Saved activations behave differently: they grow with the rank microbatch, the sequence length, and the number of blocks whose backward has not yet run, and they dominate the transient part of peak memory. Activation checkpointing reduces this component by retaining a checkpoint region’s boundary input instead of all of its internal forward activations (Chen et al. 2016). Before the region’s backward computation, the forward operations are recomputed to regenerate the values that backward needs. Figure 60 shows this trade for a mixed four-layer example.
Figure 60Interactive · conceptual accounting with the original’s illustrative 5 GB increments, not measured
Pick the checkpointed blocks, then step through forward and backward
- Checkpointed forward
- Regular forward
- Recompute
- Backward
- Boundary input
- Persistent memory
- Checkpointed forward’s short-lived interior
Conceptual accounting. The report’s example checkpoints F0 and F1 and keeps F2 and F3; the other schedules apply the same rule to every block or to none. Checkpointing any subset is the Selected-blocks policy of Table 28, and all four is Per-block. Activation steps are the original’s illustrative 5 GB; persistent memory and boundary inputs are drawn small and are not quantities. Op widths follow the original drawing, not measured time. Point at or tap the schedule, or drag the step slider (its arrow, Home and End keys also step).
Original report figure
The amount of recomputation also depends on the checkpoint implementation. PyTorch’s non-reentrant variant stops recomputing once it has regenerated every tensor needed by the pending backward computation, rather than necessarily executing the entire checkpointed region again (PyTorch Contributors 2026l). Figure 61 illustrates this early-stop behavior.
Figure 61Interactive · conceptual, the original’s three operators, no timing
Step through backward and see where recomputation stops
Conceptual. The three operators, the saved and discarded values and the stopping rule follow the original figure: “In non-reentrant mode, Op3 can be skipped if no future op requires x3.” The switch shows the other case, in which Op 3 backward needs x₃, so the same rule recomputes Op 3. The reentrant variant reruns the whole region. Point at or tap the diagram, or drag the step slider (its arrow, Home and End keys also step).
Original report figure
13.1 Choosing Recompute Granularity
The choice of checkpoint region determines which activations are retained and which operations repeat (Korthikanti et al. 2023; Jain et al. 2020). The DDP training stack selects among block-level and sub-block regions through its model configuration, summarized in Table 28. Per-block recompute, which the production recipes call per-layer recompute, wraps every transformer block, so backward holds only one block’s regenerated activations at a time. Chunk recompute instead wraps the entire block stack in a single region: only the embedding output is retained, and the stack recomputes from that boundary during backward. Selected-block recompute wraps only blocks named by key, and sub-block options wrap the attention region, the routed-expert compute and unpermute region, or, on the block all-to-all EP path, only the post-combine unpermute, inside an otherwise unwrapped block. The stack rejects chunk and per-block recompute together, and TBO (Section 15.3) rejects configurations that enable per-block recompute.
| Policy | Checkpointed region | Retained at the boundary |
|---|---|---|
| Per-block | Every transformer block | Each block’s input; used by the measured endpoint below and, per layer, by the Ultra production recipe of Section 14 |
| Chunk | The whole block stack as one region | Only the embedding output between forward and backward; backward regenerates the whole stack’s activations at once |
| Selected blocks | Blocks named by key | The named blocks’ inputs; other blocks save normally |
| Attention region | The attention sub-region of a block | The sub-region input; expert compute and EP transport save normally |
| MoE expert region | The routed-expert compute and the unpermute that follow dispatch, on the block all-to-all EP path | The permuted expert input; dispatch and combine transport save normally |
Recomputation must be deterministic to be correct. The DDP block rejects dropout by construction, router jitter is disabled in every recipe in this report, and the checkpoint wrappers restore the random-number-generator state, so random expert assignment, used by the headline runs of Section 14, replays the same routes. Routed blocks on the DeepEP v2 path use the reentrant checkpoint variant, which always recomputes the complete region; that path was validated under compilation.
13.2 Recomputing a Routed Block Repeats Its Communication
For a dense block, recomputation repeats local compute. For a routed block, the checkpointed region contains the router, dispatch, expert compute, and combine, so recomputation can also re-execute the expert-parallel communication during backward, as it does in the measured configuration below. The recomputed forward reproduces the same scores and the same top-\(K\) routes from the same boundary input, so the re-executed dispatch moves the same rows to the same destination ranks. The memory–time trade for a routed block therefore counts repeated transport, not only repeated FLOPs, and the measured cost below includes both.
Checkpointing also changes the rowwise path’s buffer policy. The lease system of Section 5.7 exists to keep a symmetric-memory payload alive from forward to its backward consumer. Under a checkpointed forward, the first pass discards its intermediates, and the recompute pass regenerates payloads immediately before they are consumed. The rowwise path therefore detects that it is running inside a checkpoint region and switches from leased symmetric buffers to per-call scratch buffers for combine outputs, rather than holding pool blocks whose saved contents no backward will read.
Metrics are accumulated only once. Router statistics, load-balancing metrics, and auxiliary-loss accumulators are forward side effects, and a recomputed forward would count every routed microbatch twice. The checkpoint context therefore distinguishes the original forward from recomputation, and the router and rowwise paths accumulate metrics and auxiliary-loss statistics only in the original pass. This keeps the balance-scope accounting of Section 10 intact under every recompute policy in eager execution; under compilation, the recompute detection falls back to a heuristic check.
One recompute mechanism in the router is independent of block policy. The router computes its logits from an FP32 cast of the block’s BF16 activation. Saving that cast would retain a four-byte copy of every routed token solely for the router linear’s backward, so the router materializes the FP32 input through an output-discard checkpoint: the storage is released after the logits are produced and regenerated immediately before the router linear’s backward consumes it. The BF16 activation remains the only persistent copy.
13.3 Composition with Pipeline Schedules
In the pipeline schedules of Section 6, a stage retains activations for every in-flight microbatch, so its activation footprint is the product of the in-flight count and the per-microbatch footprint. Recompute shrinks the second factor: a stage that checkpoints its blocks holds boundary inputs for each in-flight microbatch instead of full activation sets, at the cost of recomputing those blocks inside the backward phase of the schedule. This is the regime in which the production Ultra operating points run: at 58B active parameters with PP8, the recipe of Section 14 combines MXFP8 with per-layer recompute to stay within the memory budget. Recomputation can also repeat communication: regenerating the needed activations may require rerunning collectives, such as context-parallel attention or a routed block’s dispatch and combine. Sub-block policies can keep dispatch and combine outside the recomputed region, but the attention region still contains any context-parallel attention collectives.
13.4 Measured Memory–Throughput Trade-Off
We measured full per-block recompute in a matched four-GPU training run whose configuration is summarized in Table 29. Throughput is total useful tokens divided by total elapsed time over the timed window, repeated in three fresh processes with the execution order balanced across variants.
| Setting | Value |
|---|---|
| Model | One dense block followed by seven MoE blocks; \(d_{\text{model}}=4096\); expert intermediate width 4096 |
| Experts and routing | 32 routed experts, top-4, one shared expert; uniform routes and zero token drops |
| Hardware and parallelism | Four B300 GPUs, EP\({}=4\) |
| Batch | Two 8192-token sequences per rank microbatch; 524,288 tokens per optimizer step |
| Matched across variants | BF16 execution, identical normally initialized weights, identical input batches |
| Timing window | Steps 1–14 untimed (compilation and caches); throughput over steps 15–24 |
| Repetitions | Three fresh processes per variant |
Figure 62Interactive · measured, three fresh-process repetitions per variant · four B300 GPUs, EP = 4
What per-block recompute frees, and what it costs
Projected block-forward ranges per train_batch
- Block forward
- Block forward recomputed in backward
- Routed blocks 1–7: each recompute repeats dispatch and combine
Eight blocks (one dense, seven MoE) across eight rank-local microbatches; 8 = 524,288 tokens per step ÷ 4 GPUs ÷ (2 × 8,192 tokens), from Table 29.
“These duration sums are work indicators rather than an additive wall-time decomposition because kernels on different streams can overlap.”
Section 13.4
Plotted from the report’s data: the figure’s hash-validated metadata, checked against the steady-state CSVs, and the rank-0 profile summary of the same campaign (copies with provenance in /interactions/data/s13/). Bars are means of three ten-step windows; dots are the windows. Activations are peak minus persistent memory. Point at, tap or tab to a bar for exact values; the setup is Table 29.
Original report figure
The saving comes entirely out of activations. Persistent state—the parameters, gradient buffers, optimizer states, and persistent buffers resident between steps—is identical in both variants, so the 26.13 GiB reduction removes nearly two thirds of the transient part of the peak.
The profile shows which work repeats. A normal train_batch contains 64 projected block-forward ranges: eight blocks across eight rank-local microbatches. Per-block recompute raises that count to 128 because each block forward is recomputed during backward, and for the seven routed blocks each recomputation includes dispatch and combine. The summed GEMM and attention kernel duration rises by 24.8%, while the projected rank-0 GPU span rises by 26.7%, closely matching the independently measured 26.5% increase in time per unit of useful work. These duration sums are work indicators rather than an additive wall-time decomposition because kernels on different streams can overlap.
Full per-block recompute is useful when the 26.13 GiB reduction enables a larger microbatch, a deeper pipeline stage, or a model that would otherwise not fit, and this is how the Ultra production recipe uses it. When memory already fits, the measured 20.95% throughput cost favors retaining activations or selecting a narrower checkpoint region from Table 28. The report makes no timing claim for the chunk, selected-block, or sub-block policies; their checkpoint mechanism is the same, but their region boundaries and retained inputs differ.
13.5 Choosing a Recompute Policy
Transient memory can also be reduced through microbatch size and activation precision. A smaller rank microbatch reduces every saved activation at once but pushes grouped GEMMs toward the throughput crossover of Section 9, since fewer rows reach each expert. MXFP8 activation saving (Section 11) stores selected saved activations as qdata and scales instead of BF16, shrinking them without repeating any forward work. Recompute removes saved activations entirely, at the price of the recomputed region’s forward time. The Ultra recipe of Section 14 combines MXFP8 compute with per-layer recompute to keep the model within device memory.
Choosing a region starts with how much memory must be freed and which operations would repeat. A useful starting point is the smallest region that meets the memory budget, preferably excluding expert-parallel transport. The sub-block regions of Table 28 repeat attention or the routed-expert compute and unpermute, and neither repeats expert-parallel transport. Selected blocks cover a moderate shortfall at a cost proportional to the share of blocks named. Per-block recompute can accommodate a deep pipeline stage with many in-flight microbatches, at about a quarter more time per unit of useful work in this shape. Chunk recompute, which retains only the embedding output between forward and backward, helps a pipeline stage with very many in-flight microbatches; its backward regenerates the whole stack’s activations at once, so for a single microbatch its peak approaches that of no recompute.
Section takeaway
Checkpoint-region selection determines both memory savings and repeated work. For a complete routed block, that work can include dispatch and combine; metrics, however, are accumulated only on the original forward. At the measured operating point, per-block recompute removes nearly two thirds of activation memory, or one quarter of the peak, at a one-fifth throughput cost. The Ultra recipes accept this cost because the model otherwise exceeds the memory budget.