Trade Compute for Memory with Gradient Checkpointing

Hardpytorch

Recompute Instead of Remember

Implement run_blocks(blocks, x, use_checkpoint=False).

It applies each module in blocks to x in sequence and returns the result. When use_checkpoint is True, every block must run under torch.utils.checkpoint.checkpoint, so its intermediate activations are dropped after the forward pass and recomputed during backward.

The two modes must be numerically indistinguishable: same output, same gradients. Only peak memory differs.

run_blocks(blocks, x, use_checkpoint=False)   # stores every activation
run_blocks(blocks, x, use_checkpoint=True)    # recomputes them in backward

Hints

from torch.utils.checkpoint import checkpoint, then x = checkpoint(block, x, use_reentrant=False).
You must pass use_reentrant explicitly. Omitting it emits a UserWarning on stderr, and the grader treats any stderr output as a failure - which mirrors how this bites you in a real training log.
The non-reentrant implementation (use_reentrant=False) is the one to use in new code; it handles inputs that do not require grad and works with torch.no_grad().

Example:

Input:

blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]); run_blocks(blocks, x, use_checkpoint=True)

Output:

The same tensor run_blocks(blocks, x, use_checkpoint=False) produces, with identical gradients after .backward()

Reasoning:

Checkpointing changes only what is kept in memory between forward and backward. Each block's activations are dropped and recomputed on the backward pass, so the numerical result is unchanged while peak activation memory falls from O(depth) to O(1) segments.

Learn About topic

Contributors:

Deep-ML

Contribute