Support checkpoints in gradients

[josevalim]

Introduce Nx.Defn.checkpoint (or Nx.Defn.remat) to control the memory-computation tradeoff during reverse-mode AD in defn. Users and libraries (e.g. Axon) can mark computation regions whose intermediates are recomputed during the backward pass rather than stored, reducing peak memory for large models.

Background

Nx.Defn.grad performs reverse-mode AD. The current implementation (Nx.Defn.Grad) retains all intermediate activations as residuals for the backward pass. For an n-layer network this means O(n) activation memory — the dominant cost when training large models.

defn forward(params, input) do
  input
  |> dense(params.layer1)  # activation stored
  |> Nx.relu()             # activation stored
  |> dense(params.layer2)  # activation stored
  |> Nx.relu()             # activation stored
  # ... n layers
end

All 2n intermediates must be live simultaneously during backprop, quickly exhausting GPU/TPU memory for transformers and similar architectures.

Proposed solution: Nx.Defn.checkpoint

Nx.Defn.checkpoint(fun)

Only the inputs to fun are saved; intermediates inside fun are discarded and recomputed on demand during the backward pass.

API

defn forward(params, input) do
  input
  |> Nx.Defn.checkpoint(&dense_block(params.layer1, &1))
  |> Nx.Defn.checkpoint(&dense_block(params.layer2, &1))
  |> Nx.Defn.checkpoint(&dense_block(params.layer3, &1))
  |> softmax_cross_entropy(labels)
end

defnp dense_block(params, x) do
  x |> Nx.dot(params.weights) |> Nx.add(params.bias) |> Nx.relu()
end

Activation memory drops from O(n) to O(sqrt(n)) with optimal checkpoint placement (Chen et al., 2016).

Axon integration (axon#372):

Axon.input("input")
|> Axon.dense(512, activation: :relu)
|> Axon.checkpoint()
|> Axon.dense(256, activation: :relu)
|> Axon.checkpoint()
|> Axon.dense(10)

Semantics

1. Outside grad: no-op — checkpoint(fun) ≡ fun.()
2. Inside grad: backward pass re-executes fun with saved inputs instead of using stored intermediates
3. Composability: must work with while/4, cond/3, containers, nested checkpoints, and custom_grad/2

Implementation

Expression node

New :checkpoint node in Nx.Defn.Expr, following the :optional/:while pattern:

%Nx.Defn.Expr{
  op: :checkpoint,
  args: [input_container_expr, body_expr, body_fun]
}

• input_container_expr — inputs saved for recomputation
• body_expr — traced forward (for shape/type inference)
• body_fun — re-executed during backprop

Gradient transform

In Nx.Defn.Grad, when update_grads/6 encounters :checkpoint:

1. Re-trace body_fun with saved inputs
2. Compute gradients through the re-traced body
3. Propagate input gradients to the parent graph

This mirrors JAX's remat_call primitive.

Backend support

• Nx.Defn.Evaluator: pass-through (optionally remat for memory savings in eager mode)
• EXLA: must prevent XLA's CSE from defeating rematerialization; use rematerialization hints or custom-call boundaries to signal recomputation regions to HLO
• Torchx: delegate to torch.utils.checkpoint.checkpoint

Interaction with Nx.block (#946)

The two are complementary — Nx.block provides extensible named computation with backend dispatch; checkpoint provides memory optimization for gradients. A checkpoint can wrap a block and vice versa.

Future extensions

Rematerialization policies

An initial implementation should recompute everything. A future version could add selective policies (à la JAX's checkpoint_policies):

Nx.Defn.checkpoint(fun, policy: fn op, _args -> op in [:dot, :conv] end)

This saves expensive-to-recompute ops (matmuls, convolutions) while recomputing cheap elementwise ops. The initial design should accommodate this.

Comparison

Aspect	Today (store all)	checkpoint	Full recompute
Activation memory	O(n)	O(sqrt(n))	O(1)
Compute overhead	1×	~1.33×	2×
User control	None	Explicit	N/A
CSE prevention needed	No	Yes	N/A

Implementation plan

• Add :checkpoint node to Nx.Defn.Expr
• Tracing support: Nx.Defn.checkpoint/1 builds correct expression
• Nx.Defn.Tree traversals (apply_args/4, traverse/3, etc.)
• Gradient rule in Nx.Defn.Grad (recompute forward, then differentiate)
• Nx.Defn.Evaluator support
• EXLA compiler support with CSE prevention
• Torchx backend support
• Tests: gradient correctness, nested checkpoints, while/cond interaction, custom_grad interaction, container inputs/outputs
• Documentation and examples
• (Downstream) Axon per-layer/per-block checkpointing (axon#372)

References

• Nx source: nx/lib/nx/defn/grad.ex, expr.ex, tree.ex, evaluator.ex; exla/lib/exla/defn.ex; torchx/lib/torchx/backend.ex
• Related issues: #946 — Nx.block, axon#372
• Prior art: JAX jax.checkpoint/jax.remat, JAX PR #1749, PyTorch torch.utils.checkpoint, PyTorch min-cut recomputation blog, Chen et al. "Training Deep Nets with Sublinear Memory Cost" (2016)

[josevalim]

https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467

[blasphemetheus]

I'd like an axon feature blocked by this, so will start to look into this one, will try and post progress here

[blasphemetheus]

right, so I did a first try at this over at https://github.com/blasphemetheus/nx/tree/feat/gradient-checkpointing

things to note:

• Nx.Defn.checkpoint(input, fun) — explicit inputs API, pipe-friendly
• :checkpoint expression node following the :optional pattern
• Evaluator backend support (pass-through with correct parameter resolution)
• Gradient via parents_args re-tracing (same as :optional)
• 52 tests
• design doc in nx/CHECKPOINT_DESIGN.md

not done: EXLA backend (needs optimization barrier [thing says don't optimize across this boundary] readded for Common Subexpression Elimination prevention - we want to recompute the backward pass from the forward pass not reuse the computed value from the forward pass), any sort of axon integration (#372), rematerialization policies (checkpoint recomputes everything inside the body rn, policy lets you decide, the future optimization that was mentioned)

design recipe made this a lot faster holy cow (did this while waiting for someone at an ER overnight :) )

[polvalente]

@blasphemetheus I'll take a look at that diff later in the week, but just so you know, this issue was expanded by me yesterday and hasn't been discussed yet with @josevalim.

So the specification might change.

Also, an open question: should :checkpoint be implemented in a way that Nx.Defn.Evaluator never caches its results, or should this somehow depend on whether grad was applied to it? I think the former is fine and would be useful outside of grad. Name might need to change though.

[polvalente]

Also, for EXLA we probably want to compile :checkpoint as https://openxla.org/stablehlo/spec#optimization_barrier

[blasphemetheus]

gotcha! what does expanded mean in this context? (you're referring to edit of original issue description)

Probably the work on block has ramifications here
, is that what you mean?

And I was mistaken on what the implementation on the fork does.

Currently

currently:
forward: evals body, caches all intermediates (no memory benefit)
backward: retraces the function (correct gradients)

the version on the feat/gradient-checkpointing branch of my fork

currently caches everything -> checkpoint there rn is only a boundary for gradient re-tracing. so you're saying as an alternative:

# Alternative:
Forward: evaluates body, discards intermediates (memory benefits, even when eager)
Backward: re-evaluates body from saved input (recompute cost)

so the alternative current approach would make this useful any time you want to trade compute for memory

But yeah, whatever spec seems appropriate can adapt the implementation to. The tests define the correctness contract (so it shouldn't change anything to run checkpoint semantically), but we can switch around the internals.

And naming - JAX uses remat (rematerialize), then there's checkpoint. I am not the best at coining terms so I'll leave that to y'all lol

[blasphemetheus]

also when the desired spec is figured out lmk and I can make a draft PR based on that

[polvalente]

Expanded means that the issue was effectively empty before I updated it yesterday

[blasphemetheus]

Right! I see the edit now, gotcha I didn't see that before