Activation Checkpointing Trade-offs

Activation checkpointing saves device memory by discarding selected forward intermediates and rebuilding them during backward; the bill arrives as recomputation, runtime complexity, and possible semantic mistakes. This guide chooses boundaries from measured peaks and proves the checkpointed step matches its unmodified control.

The proposed standard asks for more than an out-of-memory exception disappearing. Its receipt contains a phase trace, segment plan, peak allocation, step-time distribution, RNG policy, paired losses and gradients, and the smallest selective plan that unlocks the target shape.

The method connects gradient checkpointing, activation memory, selective recomputation, and training memory once each so the implementation vocabulary stays natural and bounded. That vocabulary supports the article's activation checkpointing decision without creating a second intent.

1. Forward
2. Retain
3. Replay
4. Gradient

Activation checkpointing begins with a phase trace

The useful target is the peak that blocks a real batch shape, not a generic claim that activations are large. The PyTorch checkpoint documentation documents recomputation behavior, implementation variants, RNG handling, and important semantic warnings. The blocking peak should be located before checkpoint wrappers are added, because not every allocation belongs to retained forward intermediates. A proposed review of “Activation checkpointing begins with a phase trace” has four inspectable moves.

• Warm the allocator before measuring
• Mark forward, backward, optimizer, and save phases
• Reset and read peak statistics
• Store batch, sequence, dtype, and device context

The proposed evidence for activation checkpointing begins with a phase trace is two reproducible phase traces for the target and a boundary shape. The activation checkpointing begins with a phase trace receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is optimizing an idle snapshot or an unrepresentative short sequence. If optimizing an idle snapshot or an unrepresentative short sequence occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is no boundary work starts until the blocking phase and allocation are identified. Promotion under “no boundary work starts until the blocking phase and allocation are identified” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Choose boundaries by saved bytes

Equal layer counts rarely imply equal memory because attention, experts, residuals, and temporary workspaces have different shapes. The Sublinear memory paper formalizes trading extra computation for lower memory through checkpointed graph evaluation. Segment value depends on saved bytes and replay cost, not on making a visually symmetrical layer partition. A proposed review of “Choose boundaries by saved bytes” has four inspectable moves.

• Profile saved intermediates by block
• Estimate replay cost per candidate segment
• Respect stateful and side-effecting boundaries
• Start with the smallest high-value region

The proposed evidence for choose boundaries by saved bytes is bytes released per millisecond of added step time. The choose boundaries by saved bytes receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is checkpointing every module because the API makes it easy. If checkpointing every module because the API makes it easy occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is select the next segment from measured memory value rather than visual symmetry. Promotion under “select the next segment from measured memory value rather than visual symmetry” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Plot the Pareto frontier

Memory reduction and training speed must be read together at the same batch and accumulation plan. The PyTorch activation checkpointing techniques compares practical checkpoint approaches across the memory and runtime trade-off. The useful frontier pairs end-to-end step time with the batch or sequence shape unlocked by the lower peak. A proposed review of “Plot the Pareto frontier” has four inspectable moves.

• Keep data and optimizer settings fixed
• Measure warm p50 and p95 step time
• Include allocator retries and communication
• Retain the unmodified control

The proposed evidence for plot the pareto frontier is qualified batch capacity versus end-to-end step time. The plot the pareto frontier receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is quoting kernel recompute without data loading and synchronization. If quoting kernel recompute without data loading and synchronization occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is ship the least expensive plan that clears the required reserve. Promotion under “ship the least expensive plan that clears the required reserve” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Signal	Decision	Evidence
No checkpoints	Control	23.4 GB · 412 ms
4 segments	Ship	15.1 GB · 486 ms
Every block	Reject	11.8 GB · 731 ms

Prove gradient equivalence

A memory-saving step that changes the mathematical result without explanation is a new training method, not an optimization. Output, loss, and gradient comparison distinguishes a semantic optimization from a different training procedure that merely completes. A proposed review of “Prove gradient equivalence” has four inspectable moves.

• Seed control and checkpointed runs identically
• Compare outputs and scalar loss
• Compare selected and global gradient norms
• Declare dtype-aware tolerances

The proposed evidence for prove gradient equivalence is running a paired fixture with one intentional mismatch control. The prove gradient equivalence receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is checking only that the process finishes. If checking only that the process finishes occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is paired outputs, losses, and gradients stay within the stated numeric tolerance. Promotion under “paired outputs, losses, and gradients stay within the stated numeric tolerance” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Runnable artifact. Save this bounded specimen as activation-checkpointing-tradeoffs.test.mjs and run node --test activation-checkpointing-tradeoffs.test.mjs. Expected result: PASS: gradients remain within tolerance.

import assert from "node:assert/strict";
import test from "node:test";
const close = (a, b, tolerance = 1e-6) => Math.abs(a - b) <= tolerance;
test("compares checkpointed gradient", () => {
  assert.equal(close(.314159, .3141594), true);
  assert.equal(close(.314159, .32), false);
  console.log("PASS: gradients remain within tolerance");
});

Handle randomness deliberately

Dropout and other stochastic operations can diverge when forward work is replayed under a different random state. Random-state handling is part of equivalence whenever dropout or another stochastic operation is replayed during backward. A proposed review of “Handle randomness deliberately” has four inspectable moves.

• Choose and document RNG preservation
• Test dropout-heavy blocks
• Compare repeated seeded steps
• Measure the preservation overhead

The proposed evidence for handle randomness deliberately is gradient agreement across several fixed seeds. The handle randomness deliberately receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is accepting nondeterminism without locating its source. If accepting nondeterminism without locating its source occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is stochastic equivalence or an explicitly justified alternative must be proved. Promotion under “stochastic equivalence or an explicitly justified alternative must be proved” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Reject stateful replay hazards

A checkpointed function should behave like a pure computation over its inputs during both original and replayed execution. Mutable caches, counters, and device moves can execute twice inside a selected region and create silent stateful divergence. A proposed review of “Reject stateful replay hazards” has four inspectable moves.

• Inventory global counters and mutable caches
• Avoid device moves inside the region
• Test training and evaluation modes
• Fail when forward and replay paths diverge

The proposed evidence for reject stateful replay hazards is call counts and outputs from an instrumented boundary. The reject stateful replay hazards receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is mutating hidden state twice during a single optimization step. If mutating hidden state twice during a single optimization step occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is no selected region may produce an externally visible replay side effect. Promotion under “no selected region may produce an externally visible replay side effect” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

1. TraceTrace

   Measure phase peaks on the target batch shape.

2. SelectSelect

   Place boundaries around expensive saved intermediates.

3. CompareCompare

   Match outputs, losses, and gradients to a control.

4. StressStress

   Test RNG, stateful layers, mixed precision, and failure.

Measure distributed consequences

Longer backward recomputation can shift communication overlap and move the bottleneck outside the local device. Extra backward work can expose collectives or pipeline bubbles, so a local memory win needs a production-topology trace. A proposed review of “Measure distributed consequences” has four inspectable moves.

• Trace collective start and finish times
• Compare exposed communication
• Inspect pipeline bubbles
• Repeat at the target world size

The proposed evidence for measure distributed consequences is global samples per second beside per-device peaks. The measure distributed consequences receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is choosing from a single-device microbenchmark for a distributed job. If choosing from a single-device microbenchmark for a distributed job occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is the production topology must retain the memory win without unacceptable exposed communication. Promotion under “the production topology must retain the memory win without unacceptable exposed communication” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Publish the segment plan

Future model revisions need to know why each boundary exists and which check will catch stale assumptions. Each wrapper should point to its measured saving and equivalence fixture so future model revisions know why it remains. A proposed review of “Publish the segment plan” has four inspectable moves.

• Name modules and revision hashes
• Attach before-and-after traces
• Save equivalence tolerances
• List unsupported execution modes

The proposed evidence for publish the segment plan is a clean rerun of memory, speed, and gradient fixtures. The publish the segment plan receipt should bind the relevant input and configuration to the observed output, then retain a deliberately rejected control so the check can prove it distinguishes a bad result.

The falsification target for this activation checkpointing decision is leaving anonymous wrapper calls scattered through model code. If leaving anonymous wrapper calls scattered through model code occurs, a favorable aggregate can still conceal the exact cohort, queue state, position, geometry, or frame that invalidates the conclusion.

The article's proposed boundary is every boundary maps to a measured saving and a semantic test. Promotion under “every boundary maps to a measured saving and a semantic test” should wait until another reviewer can evaluate the boundary from named evidence without relying on the author's authority.

Keep the boundary visible

Choose replay boundaries from phase-labeled memory and runtime evidence, then compare the checkpointed step with an unmodified seeded control. RNG behavior, side effects, mixed precision, and distributed overlap are part of correctness.

The best plan is usually selective: enough saved memory to unlock the intended batch or sequence, with the smallest measured recomputation bill and a test attached to every boundary.

Continue through four related field notes: distributed checkpointing, FlashAttention IO-aware attention, tensor parallelism, model distillation. Each extends the activation checkpointing method without changing this article's single search intent.