Gradient Checkpointing with jax.checkpoint / jax.remat

Deep neural networks consume enormous amounts of memory during training because the backward pass needs to access intermediate activations from the forward pass. Gradient checkpointing (also called rematerialization) trades computation for memory: instead of storing all intermediate values, it recomputes them during the backward pass.

import jax
import jax.numpy as jnp

# Without checkpointing: stores all intermediate activations
def forward(params, x):
    for layer in params:
        x = jax.nn.relu(x @ layer['w'] + layer['b'])
    return x

# With checkpointing: recomputes activations during backward pass
@jax.checkpoint  # or equivalently, @jax.remat
def forward_checkpointed(params, x):
    for layer in params:
        x = jax.nn.relu(x @ layer['w'] + layer['b'])
    return x

# Same output, same gradients, but uses much less memory
grads = jax.grad(lambda p: jnp.sum(forward_checkpointed(p, x)))(params)

Selective Checkpointing

You don't have to checkpoint everything. The most effective strategy is to checkpoint each Transformer block, which reduces memory from O(layers) to O(1) for intermediate activations:

from flax import nnx
import jax

class TransformerBlock(nnx.Module):
    def __init__(self, d_model, num_heads, d_ff, rngs):
        self.attention = nnx.MultiHeadAttention(
            num_heads=num_heads, in_features=d_model,
            qkv_features=d_model, out_features=d_model, rngs=rngs)
        self.ff1 = nnx.Linear(d_model, d_ff, rngs=rngs)
        self.ff2 = nnx.Linear(d_ff, d_model, rngs=rngs)
        self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
        self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)

    @nnx.jit
    def __call__(self, x):
        x = x + self.attention(self.ln1(x))
        x = x + self.ff2(nnx.gelu(self.ff1(self.ln2(x))))
        return x

# Checkpoint individual blocks
class Transformer(nnx.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, rngs):
        self.blocks = [
            TransformerBlock(d_model, num_heads, d_ff, rngs)
            for _ in range(num_layers)
        ]

    def __call__(self, x):
        for block in self.blocks:
            # Remat each block: activations are recomputed in backward
            x = jax.checkpoint(lambda b, x: b(x), block, x)
        return x

# You can also use remat with a policy for fine-grained control
from jax.ad_checkpoint import checkpoint_policies

# Only save certain operations (e.g., dots but not norms)
policy = checkpoint_policies.save_only_these_names('dot_general')

@jax.remat(policy=policy)
def block_with_policy(params, x):
    # Only the results of matrix multiplications are saved;
    # everything else is recomputed
    return transformer_block(params, x)