Skip to content

← Flax step 84 of 100

Hard Primitives

nn.remat with checkpoint policies

Jump to the task

Checkpointing trades storage of derivative residuals for recomputation. A residual is a value needed later to evaluate a derivative. A checkpoint policy lets some intermediate values remain eligible for storage instead of requiring them to be recomputed when needed.

The useful tradeoff depends on both value size and computation cost. Retaining an expensive matrix-product result may avoid substantial work; recomputing a cheap elementwise operation may save storage. Neither choice guarantees lower peak memory or faster training for every model.

What a policy decides

Pass a policy through the policy= keyword of nn.remat. The standard saving policies inspect a primitive operation and its type-level arguments, including shapes and operation parameters. They return a boolean:

  • True: the result is allowed to be saved as a derivative residual.
  • False: it must be recomputed if the derivative needs it.

Permission to save is not an instruction to retain every permitted value. Unneeded results need not be stored. This is also not a callback over the actual array values at each training step.

nothing_saveable permits no internal results to be saved, matching the default checkpoint policy. Inputs can still be retained for recomputation. everything_saveable permits all results to be saved, following the usual residual-saving strategy without checkpointing. It does not force storage of every intermediate or promise an identical compiled execution plan.

The policy used here

This exercise uses jax.checkpoint_policies.dots_with_no_batch_dims_saveable. For a dot_general operation, it permits saving the result when the operation has no paired batch dimensions.

These are dimensions in the operation's dimension_numbers, not a test of whether an array contains multiple examples or has a small output:

  • Multiplying (B, D) by a shared (D, F) weight matrix can use no paired batch dimensions. The leading B axis is carried into the output.
  • Multiplying (B, N, D) by (B, D, F) as independent per-example products pairs the two B axes as batch dimensions.

The first result can be large and the second small. The policy is a structural rule, not a memory budget or a universal transformer default. Elementwise operations are not selected by this dot rule. Profile the actual differentiated computation before drawing performance conclusions.

Wrap the module class

RematMLP = nn.remat(
    MLP,
    policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable,
)
model = RematMLP(features=F)
variables = model.init(jax.random.PRNGKey(seed), x)
output = model.apply(variables, x)

Here MLP is a module class, not an instance. Define its compact method with three default Dense layers of width F, placing ReLU after each of the first two. The policy changes residual handling during differentiation, not the mathematical forward function. Use numerical tolerances when comparing results across compiled forms.

Problem

Implement nn_remat_policy(seed, x, features):

  1. Define MLP(nn.Module) with a features field and @nn.compact method: Dense(F) -> relu -> Dense(F) -> relu -> Dense(F).
  2. Keep default Dense initializers, biases and automatic names Dense_0, Dense_1 and Dense_2.
  3. Transform the class with nn.remat and the policy shown above, then construct the module with F = int(features).
  4. Initialize using jax.random.PRNGKey(int(seed)) and x.
  5. Apply the initialized model to x and return a flattened JAX array.

x is a finite float32 JAX vector of shape (D,), with D >= 1. features is a positive integer-valued scalar; seed is a nonnegative integer-valued scalar. Either scalar may arrive as a float. The output has shape (F,). Leave the input unchanged.

The output tests check the seeded forward computation. They cannot prove that the requested policy was used or measure memory savings. To study the policy itself, compare saved residuals or profile a differentiated training step while holding parameters, inputs and loss fixed.