TopolearnLarge models run out of accelerator memory on activations long before parameters; remat is the standard fix.
Reverse-mode autodiff stores residuals, intermediate values from the forward pass that the backward pass needs. For deep networks these activations dominate memory, growing with depth, sequence length and batch size.
jax.checkpoint (alias jax.remat) wraps a function so that its internal residuals are not saved. During the backward pass JAX reruns that function's forward computation to regenerate them. You pay roughly one extra forward pass in compute for a large memory reduction, often enabling bigger batches or models.
The typical pattern is to checkpoint each transformer block or layer and run the stack with lax.scan, so only per-layer inputs are kept.
Policies give finer control over what is saved versus recomputed:
jax.checkpoint_policies.nothing_saveable: recompute everything (default behaviour).jax.checkpoint_policies.dots_with_no_batch_dims_saveable: keep matmul outputs, recompute cheap element-wise ops, a good balance for transformers.save_only_these_namestogether withjax.ad_checkpoint.checkpoint_nameto save specific tagged tensors.
You can verify what gets stored with jax.ad_checkpoint.print_saved_residuals(f, *args).
XLA itself also rematerializes under memory pressure, but explicit checkpointing is predictable. Profile memory before and after: remat helps when activations, not parameters or optimizer state, are the bottleneck.
Example
import jax
import jax.numpy as jnp
from jax.ad_checkpoint import print_saved_residuals
def block(x, w):
return jnp.sin(jnp.sin(x @ w))
def net(x, w):
for _ in range(3):
x = block(x, w)
return x.sum()
ckpt_block = jax.checkpoint(
block, policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)
def net_remat(x, w):
for _ in range(3):
x = ckpt_block(x, w)
return x.sum()
x, w = jnp.ones((4, 4)), jnp.eye(4) * 0.5
print(jnp.allclose(jax.grad(net, 1)(x, w), jax.grad(net_remat, 1)(x, w))) # True
print_saved_residuals(net, x, w) # many sin/cos residuals listed
print_saved_residuals(net_remat, x, w) # fewer: dot_general outputs + inputsRemember
- Checkpointing drops saved activations and recomputes them in the backward pass
- Policies choose what to keep, e.g. save matmuls and recompute element-wise ops
- Apply it per layer, usually combined with
lax.scanover layers
Watch out for
- Checkpointing tiny functions everywhere, adding recompute cost without meaningful memory savings.