checkpoint(function: Callable[..., Tensor | tuple[Tensor, ...]], args: Tensor = (), use_reentrant: bool = False, kwargs: object = {})Run function without saving its intermediates; recompute them on backward.
Memory-for-compute tradeoff used to fit larger models / longer
sequences in fixed VRAM. In a normal forward pass autograd stashes
every intermediate activation needed by the backward pass — for a
deep transformer block this is the dominant memory cost. Wrapping
the block in checkpoint runs forward under no_grad,
saves only the inputs, and re-executes the block from scratch
during backward to recompute the intermediates on demand.
Net effect: peak memory drops from "all activations" to "inputs + one block's intermediates at backward time"; wall-clock grows by roughly the cost of one extra forward pass per checkpointed block.
Parameters
functioncallable(args, kwargs) — if it consumes randomness (dropout,
random init), seed it explicitly inside the callable, otherwise
forward and re-forward will disagree.function on backward.use_reentrantbool= FalseFunction subclass with manual recompute).
Default False.**kwargsobject= {}function. They
are not differentiated through.Returns
Notes
Best applied to a sequence of homogeneous blocks (transformer layers, ResNet stages) where the per-block memory saving compounds. Checkpointing every layer roughly doubles training time; the usual recipe is to checkpoint every other layer or once per stage.
Examples
>>> import lucid
>>> from lucid.utils.checkpoint import checkpoint
>>>
>>> def block(x):
... return (x @ x.T).relu().sum(dim=-1)
>>>
>>> x = lucid.randn(4, 8, requires_grad=True)
>>> y = checkpoint(block, x) # forward under no_grad; saves only x
>>> y.sum().backward() # re-runs `block(x)` to populate gradsSee Also
- lucid.autograd.Function—the autograd primitive underneath.