Gradient checkpointing & memory maths

01 · Concept

Concept

Lesson 5.7 got persistent state down to about 6.75 GB per rank on a 64-way ZeRO-3 job. Lesson 5.8 partitioned the layers. On paper, a 27-billion-parameter model now fits an 80 GB accelerator with enormous headroom. Then you launch with an 8,192-token sequence and batch size one, and the job dies at the first backward pass. Nothing about your sharding was wrong. You budgeted the wrong category.

Backpropagation needs values from the forward pass. Autograd therefore retains activations, normalized copies, projection outputs, and attention intermediates until their backward functions run, and in a deep transformer those saved values routinely exceed the weights. Activation checkpointing — usually called gradient checkpointing, though nothing about a file on disk is involved — reduces that by deliberately forgetting and recomputing.

Get the size right first, because the intuition people carry is badly wrong. The tempting estimate counts the residual stream: batch times sequence times hidden times layers times dtype bytes.

Five gibibytes. Comfortable. Now count what autograd actually keeps in one layer. Take a uniform full-attention model with this model’s dimensions — the hybrid’s real bill needs the DeltaNet layers costed separately, which is its own exercise; what follows is the method and the order of magnitude. The residual input and its normalized copy are 5120 wide each. In a full-attention layer the projections are non-square: the query projection goes 5120 to 12288, because it carries the output gate alongside the 24 heads at dimension 256 (lesson 4.2); keys and values are 5120 to 1024 each (4 KV heads), the attention output is 6144, and the output projection returns 5120. That is roughly 35,800 values per token. The feed-forward block is larger still: a 5120-wide normalized input, then gate, up, and gated-product tensors of width 17,408 apiece, then a 5120-wide down-projection output — about 62,500 values. Together:

That is 19 times the residual stream alone, and the FFN’s three 17,408-wide tensors are most of the difference. Scale it up:

Ninety-six gibibytes of saved activations for one sequence at batch size one — more than ten times the persistent state your sharding worked so hard to reduce, and more than an 80 GB card holds. That is the figure for the uniform-attention model we posited, and it is the number to reason with. For Qwen3.8-27B specifically it is an approximation of unknown sign: lesson 4.16 showed only 16 of its 64 layers are full attention, and the other 48 are Gated DeltaNet, whose training activations are a separate accounting — a fixed-size inference state does not imply cheap backward storage, because the recurrent mixer still materialises per-token intermediates that autograd must keep. The FFN half, which is most of the total, is identical in all 64 layers either way. The residual-only estimate was off by roughly a factor of nineteen, and it is off in the direction that gets a job scheduled and then killed.

Now apply checkpointing, and let the architecture choose the boundary. Divide the network into segments; during the forward pass save only each segment’s boundary tensor rather than everything inside it; during backward, rerun that segment’s forward from the saved boundary, rebuild its intermediates, compute gradients, and release them. This model’s 64 layers form 16 repeating super-blocks of three Gated DeltaNet layers plus one full-attention layer, which is a natural and perfectly balanced segmentation. Each boundary is one hidden state:

During backward, one super-block is rematerialized at a time, costing four layers’ worth of intermediates — still priced at the uniform-attention rate, but that rate is only a hypothetical comparison point for the hybrid, not a known ceiling or a measured peak:

Within that full-attention teaching model, peak saved-activation memory becomes roughly $1.25+6.0\approx 7.3$ GiB instead of 96 GiB. The roughly thirteenfold ratio belongs only to that shape-based count; it does not transfer to the real hybrid until its DeltaNet autograd graph, fused kernels, allocator, and recomputation policy are measured.

The tradeoff is direct: lower peak activation memory for extra computation. It is not simply double training time, because backward, communication, and non-checkpointed work continue regardless and recomputation may overlap differently. Measure step time and realized throughput, not the theoretical FLOP increase.

A memory budget should keep the categories separate:

Checkpointing attacks $M_{activations}$ and nothing else. It does not shard parameters and does not change persistent optimizer state. The exact dtypes and components of that state are recipe-dependent, and lesson 9.3 owns the Qwen memory budget. A job dominated by persistent state needs ZeRO or FSDP; one dominated by long-sequence activations needs checkpointing.

Sequence length is where this model gets interesting. The estimate above used 8,192 tokens; native context is 262,144. Activations in the position-wise parts of the network grow linearly with $T$, so the same arithmetic at full context is thirty-two times larger and utterly infeasible without both checkpointing and the sequence parallelism of lesson 5.8. Ordinary attention would add a term quadratic in $T$ on top, which is precisely why memory-efficient attention kernels that never materialize the full probability matrix are mandatory at these lengths — and why only 16 of the 64 layers face that term at all, the other 48 scanning positions with a fixed-size recurrent state.

Two correctness traps deserve naming. Stochastic operations must reproduce their randomness: if dropout drew one mask in the original forward and a different one during recomputation, backward would differentiate a different function, so framework checkpoint utilities preserve and restore RNG state at a small cost. And stateful modules, side effects, mutable caches, and data-dependent external calls can make recomputation silently incorrect rather than loudly broken.

Composition adds the rest of the complexity. Pipeline parallelism makes recomputation compete with other microbatches for the same devices. FSDP may need to all-gather weights again during recompute unless parameters stayed materialized, trading communication against memory. Compilers fuse and reorder regions in ways that change what a segment even is. The effective unit is the whole distributed schedule, never an isolated function.

The durable mental model is a time–space exchange: save enough boundary state to reconstruct the forward graph, discard the bulky middle, and pay compute to recreate it just before its gradients are needed. The decision follows from a measured breakdown of where memory actually went — which, for a model with a 17,408-wide feed-forward block, is almost never where a first estimate puts it.