Gradient Checkpointing: Recompute Activations in Backward Instead of Storing Them

The problem, from first principles

01-kv-cache-autoregressive-generation cached results forward in time (reuse past K/V instead of recomputing them) to save compute at the cost of memory. Gradient checkpointing makes the exact opposite trade, on the training side of the same network: a standard backward pass needs every layer's intermediate activation kept in memory (there are num_layers + 1 of them, all alive at once), which becomes the actual memory bottleneck for very deep networks — gradient checkpointing throws those intermediates away after the forward pass and recomputes them, on demand, during backward, trading extra compute for far less peak memory.

From theory to code

Implement linear_backward, forward_full/backward_full (the standard, memory-heavy approach), forward_checkpointed/backward_checkpointed (which keeps nothing but the original input, recomputing everything else when backward needs it), and count_stored_activations, quantifying the memory difference.

Constraints

• forward_full returns (final_output, activations, pre_activations), storing every intermediate value.
• forward_checkpointed returns only final_output — nothing else is kept.
• backward_checkpointed must produce results numerically identical to backward_full on the same computation — checkpointing changes when things are computed, never what the answer is.
• count_stored_activations(num_layers, use_checkpointing) returns num_layers + 1 without checkpointing, 1 with it.

Hints

Open one at a time. Each gives away a little more than the last.