Skip to content

← Flax step 82 of 100

Medium Primitives

nn.checkpoint (gradient checkpointing)

Jump to the task

During reverse-mode differentiation, JAX retains values needed to compute derivatives. These saved values are called residuals. They can include inputs and intermediate results, but are not necessarily every activation produced by the forward pass.

Gradient checkpointing lets JAX recompute selected values during the backward pass instead of keeping them all in memory. It is also called rematerialization, or remat. This is different from saving model weights to a checkpoint file.

Wrap a Flax module

flax.linen.checkpoint, also available as nn.remat, adapts JAX's checkpoint transformation to Flax modules and their variables and RNGs. For this exercise, transform a module class and then construct an instance:

import jax
import flax.linen as nn

class MLP(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.features)(x)
        x = nn.relu(x)
        return nn.Dense(self.features)(x)

CheckpointMLP = nn.checkpoint(MLP)
model = CheckpointMLP(features=4)
variables = model.init(jax.random.PRNGKey(0), x)
output = model.apply(variables, x)

Pass the class MLP to this form of the transformation, not MLP(...). Initialization and application then use the usual Flax calls. The default Dense layers include biases and are named Dense_0 and Dense_1.

What changes when differentiating

The wrapped model computes the same mathematical forward function with the same parameters. Checkpointing changes which derivative residuals are stored and which are recomputed. It does not add a layer or change the loss. Floating-point results should be compared with a tolerance rather than a promise of byte-for-byte equality across compilation settings.

A forward-only call does not demonstrate backward-pass memory savings. To explore the transformation, compare gradients of the same scalar loss for wrapped and unwrapped models, using the same parameter tree. Inspect the differentiated computation or profile a training step to study what is recomputed and how much memory it uses.

Choose where to recompute

Savings and runtime cost depend on the operations, checkpoint boundaries, compiler and hardware. There is no fixed percentage improvement or slowdown. A small Dense layer is a valid target, but whether wrapping it helps is a measurement question.

Checkpoint policies can allow selected intermediate results to be saved. A useful tradeoff may retain expensive-to-recompute results and recompute cheaper ones, taking their sizes into account. Nested checkpoint regions can express different schedules; repeatedly wrapping the same region is not a general recipe for saving more memory.

Problem

Implement nn_checkpoint_forward(seed, x, features):

  1. Define the compact MLP above, with two default nn.Dense(features) layers and ReLU only between them. Do not add explicit layer names.
  2. Wrap the class with nn.checkpoint and instantiate it with width F.
  3. Initialize with jax.random.PRNGKey(int(seed)) and x.
  4. Apply the initialized model to x and return its flattened output.

x is a finite float32 JAX vector of shape (D,), with D >= 1. features is a positive integer-valued scalar F; seed is a nonnegative integer-valued scalar. Either scalar may arrive as a float, so convert it to an integer. The result is a JAX vector of shape (F,). Leave x unchanged. Use the default Dense initializers and biases.

The output tests check the seeded forward computation. They cannot prove that a submission used checkpointing, and they do not measure memory savings. Follow the wrapping requirement even though an unwrapped model can produce the same returned values.