Bottom line: Using jax.checkpoint with the checkpoint_dots policy on TPU v5e hardware cuts peak activation memory by roughly 40% during large-model training, at a recomputation cost of around 12% in per-step wall-clock time. For transformer-class models running on v5e pods, this tradeoff is almost always worth taking when HBM capacity is the constraint blocking larger batch sizes or longer sequences.

• API: jax.checkpoint (canonical alias: jax.remat), stable across JAX 0.4.x series
• Target hardware: TPU v5e — 16 GB HBM2 per chip, available in pod slices from v5e-4 through v5e-256
• Default policy: jax.checkpoint_policies.checkpoint_dots — saves only matrix multiplication inputs, discards all other intermediates
• Reported memory saving: ~40% reduction in peak activation memory for multi-layer transformer models
• Speed overhead: ~12% increase in per-step wall-clock time from recomputation during the backward pass
• JAX version: Tested configurations consistent with JAX 0.4.25+ behavior

How does JAX gradient checkpointing reduce memory on TPU v5e?

Gradient checkpointing (also called rematerialization in XLA terminology) works by discarding intermediate activations during the forward pass and recomputing them on-demand during the backward pass. Instead of holding every layer’s output in HBM for the full forward sweep, JAX’s jax.checkpoint decorator instructs XLA to rebuild those tensors from saved checkpoints when the gradient computation needs them. On TPU v5e, where each chip carries 16 GB of HBM2, this distinction between “store everything” and “recompute selectively” is often the difference between a run that fits and one that OOMs.

The documentation screenshot above shows the jax.checkpoint API surface: the function takes a fun argument (the forward function to wrap) and an optional policy callable that decides, at the XLA level, which intermediate values are worth saving and which should be dropped and recomputed. Without a policy, JAX’s default behavior is to save nothing — every intermediate is discarded and recomputed. That’s the maximum memory reduction, but it also pays the highest recomputation tax. The policy argument is how you dial in the 40%/12% sweet spot.

The underlying mechanism is part of JAX’s autodiff system. When you call jax.grad on a checkpointed function, the traced computation is split into a “residuals” set (values that cross the forward-backward boundary and must be stored) and an “unreachable” set (values that can be recovered by rerunning the forward computation from the residuals). The checkpoint_dots policy specifically marks the inputs to lax.dot_general calls — the primitive underlying all matrix multiplications — as residuals, while everything else gets recomputed. For transformer attention and FFN blocks, where dot_general dominates compute, this is a well-calibrated default.

What is the checkpoint policy that drives the 40% memory saving?

The checkpoint_dots policy saves the inputs to all dot products (matrix multiplications) and discards all other activations. For a standard transformer block with self-attention and a two-layer FFN, this typically means keeping four tensors per block (Q, K, V projections and FFN intermediate input) instead of the full activation stack, which is why the memory footprint drops by roughly 40% on multi-layer models. Shallower networks with fewer repeated blocks see less benefit.

Here is the minimal setup for wrapping a forward function with this policy:

import jax
import jax.numpy as jnp
import functools

# Wrap the function you want to checkpoint
@functools.partial(
    jax.checkpoint,
    policy=jax.checkpoint_policies.checkpoint_dots
)
def transformer_block(params, x):
    # Self-attention
    qkv = jnp.dot(x, params['qkv_weight'])  # saved: x
    # ... attention computation ...
    ffn_out = jnp.dot(x, params['ff1'])      # saved: x
    # ... all non-dot intermediates are recomputed during backward
    return out

JAX ships several built-in policies under jax.checkpoint_policies. The ones most relevant for TPU v5e training are:

• checkpoint_dots — saves dot product inputs; good default for transformers
• checkpoint_dots_with_no_batch_dims — like checkpoint_dots but excludes batch-dimension contractions, reducing residuals further at the cost of more recomputation
• everything_saveable — saves all intermediates (equivalent to no checkpointing, maximum speed)
• nothing_saveable — saves nothing (maximum memory reduction, maximum recomputation)

You can also write a custom policy as a Python callable that receives an Primitive and returns a boolean. This is useful when you know your model has expensive-to-recompute operations (like softmax with large sequence lengths) that you’d rather keep in memory:

from jax._src import core as jax_core

def custom_policy(prim, *args, **params):
    # Save softmax outputs, let everything else be recomputed
    if prim is jax.lax.reduce_window_sum_p:
        return True  # save
    # Default: use checkpoint_dots behavior for dot products
    return jax.checkpoint_policies.checkpoint_dots(prim, *args, **params)

@functools.partial(jax.checkpoint, policy=custom_policy)
def attention_block(params, x, mask):
    ...

The benchmark chart above illustrates the memory-versus-speed curve across checkpoint policies on a TPU v5e-8 pod slice running a 3B parameter decoder model. The checkpoint_dots policy sits near the knee of the curve — the point where you recover the most memory per unit of added compute. Policies that recompute more aggressively (like nothing_saveable) push further left on memory but the speed penalty accelerates non-linearly once you start recomputing expensive operations like layer norm statistics. The 40%/12% point corresponds to checkpoint_dots at batch size 32, sequence length 2048.

What is the actual speed overhead on TPU v5e?

The 12% wall-clock overhead from gradient checkpointing on TPU v5e is lower than the equivalent configuration on NVIDIA A100s for two structural reasons: TPU v5e’s matrix units are optimized for exactly the class of operations being recomputed (dense matrix multiplications), and the TPU’s on-chip SRAM is fast enough that the recomputation latency is partially absorbed by compute overlap. The overhead you’ll see in practice depends on your batch size and sequence length, but 10–15% is the typical reported range for standard transformer workloads.

The specific reason the overhead stays contained is that XLA’s scheduler can overlap the recomputation of layer N’s activations with the gradient computation of layer N+1. This pipeline overlap means you’re not paying the full cost of recomputation sequentially — the TPU’s matrix units stay busy. If you disable XLA’s overlap (not something you’d normally do, but possible for debugging), the overhead can climb to 25–30%, which shows how much of the “free” speed you’re getting from the hardware’s ability to parallelize compute and memory traffic.

One thing that does blow up the overhead: recomputing operations that are memory-bandwidth-bound rather than compute-bound. Layer norm, element-wise activations, and large embedding lookups are all in this category. If your model has these on the critical path and you checkpoint aggressively, the recomputation hits bandwidth rather than FLOPS, and the overhead grows. The custom policy approach described above is the right fix for this — save the outputs of bandwidth-bound ops, recompute only the compute-bound ones.

How do you apply jax.checkpoint across multi-layer Flax models?

Flax exposes gradient checkpointing through nn.remat, which wraps jax.checkpoint with Flax’s module system. You can apply it at the module class level with a decorator, or use nn.remat inline when scanning over layers. The decorator approach is cleaner for fixed-depth models; the scan approach is standard for transformers with many identical blocks.

from flax import linen as nn
import functools
import jax

# Option 1: Decorate the entire module class
RematTransformerBlock = nn.remat(
    TransformerBlock,
    policy=jax.checkpoint_policies.checkpoint_dots,
)

# Option 2: Apply remat inside nn.scan for N-layer transformers
class Transformer(nn.Module):
    num_layers: int

    @nn.compact
    def __call__(self, x):
        # Use nn.scan with remat for memory-efficient multi-layer forward pass
        ScanBlock = nn.remat(
            TransformerBlock,
            policy=jax.checkpoint_policies.checkpoint_dots,
        )
        x, _ = nn.scan(
            ScanBlock,
            variable_axes={'params': 0},
            split_rngs={'params': True},
            length=self.num_layers,
        )(x)
        return x

The scan-plus-remat pattern is particularly effective on TPU v5e because nn.scan compiles to an XLA while_loop, and the rematerialization pass can reason about the loop body’s memory requirements holistically. The activation memory for an N-layer model goes from O(N) to O(1) with respect to depth — each layer’s activations are recomputed as needed rather than accumulated across the full depth.

The diagram above shows the activation memory profile for a 24-layer transformer using nn.scan with and without nn.remat. Without rematerialization, peak memory grows linearly with layer count and hits the 16 GB HBM ceiling around layer 18 for the sequence lengths tested. With checkpoint_dots, the peak memory stays roughly flat after the first few layers because each scan iteration reuses the same HBM footprint rather than accumulating new activations. This is the core reason jax.checkpoint on TPU v5e enables fitting models that would otherwise require per-layer tensor parallelism.

Is jax gradient checkpointing the right tool for every TPU v5e memory problem?

Not always. Gradient checkpointing trades memory for recompute time, but other memory pressure sources have better-targeted solutions. If your HBM is dominated by optimizer states (Adam’s first and second moments), checkpointing your forward pass won’t help — consider Adafactor or 8-bit optimizer states instead. If the bottleneck is model parameters themselves for very large models, tensor parallelism across chips (via jax.sharding with PartitionSpec) is the right lever, not checkpointing.

Where checkpointing consistently earns its keep is when the bottleneck is specifically intermediate activations — the tensors produced during the forward pass that are needed by the backward pass. For transformer training with sequence lengths above 1024, this is usually the dominant memory consumer, and jax.checkpoint with checkpoint_dots is the lowest-friction way to address it. The 12% speed cost is predictable and bounded, which matters for long training runs where you’re budgeting compute carefully.

JAX’s PyPI download trajectory reflects growing adoption across the ML infrastructure community — the volume of downloads tracks with increasing use of JAX for production training workloads, including on TPU hardware. As more teams move training to Google Cloud TPU v5e pods, the combination of jax.checkpoint, nn.scan, and jax.sharding is becoming the standard toolkit for fitting large models within per-chip HBM budgets without sacrificing training throughput.

References

• JAX Gradient Checkpointing documentation — Official JAX docs covering the jax.checkpoint API, the policy argument, built-in policies including checkpoint_dots, and custom policy functions. The primary reference for all API behavior described in this article.
• google/jax on GitHub — The JAX source repository. The jax.checkpoint_policies module and the underlying rematerialization pass in jax/_src/interpreters/ are authoritative for understanding how policies are evaluated against XLA primitives.
• Google Cloud TPU v5e documentation — Covers TPU v5e hardware specifications including HBM capacity per chip (16 GB), pod slice configurations (v5e-4 through v5e-256), and memory bandwidth figures that inform the recomputation overhead analysis.
• Flax linen transformations — nn.remat — Flax’s official API reference for nn.remat, which wraps jax.checkpoint with module-aware behavior. Covers the interaction between nn.remat and nn.scan for multi-layer models.