Optionally keep gradients in the training step state

[seanmor5]

Closes #577.

Axon.Loop.train_step/4 computes the gradients of every batch and then drops them once the optimizer has consumed them. Nothing downstream of the step can see them, so there is no way for a metric or an event handler to report a gradient norm, spot layers whose gradients have vanished, or log the raw gradients while debugging a model that does not train. The usual workaround is to write a custom step function, which means giving up trainer/4, loss scaling, and the rest of the supervised loop.

Design

train_step/4 and trainer/4 gain a keep_gradients? option (default false). When it is set, the step state carries one extra key, :gradients, holding a container with exactly the structure of Axon.ModelState.trainable_parameters/1: nested string-keyed maps of tensors, frozen parameters excluded. The value is the gradient of the loss with respect to each trainable parameter for the batch that was just processed, after loss-scale unscaling and before any optimizer transformation, which is exactly what the optimizer's update_fn receives. init_fn initializes the key to zeros of the same shapes and types.

When the option is off, the step state is byte-for-byte what it was before: the key is absent rather than nil, so existing pattern matches, checkpoints, and compiled-function templates are unaffected.

grad_norm = fn grads ->
  grads
  |> Nx.Defn.Composite.reduce(Nx.tensor(0.0), fn g, acc -> Nx.add(acc, Nx.sum(Nx.pow(g, 2))) end)
  |> Nx.sqrt()
end

model
|> Axon.Loop.trainer(:mean_squared_error, :adam, keep_gradients?: true)
|> Axon.Loop.metric(grad_norm, "gradient norm", :running_average, [:gradients])
|> Axon.Loop.handle_event(:iteration_completed, fn state ->
  IO.inspect(state.step_state.gradients["dense_0"]["kernel"], label: "kernel gradient")
  {:continue, state}
end)
|> Axon.Loop.run(data, Axon.ModelState.empty(), epochs: 5)

The metric transform has to be the list [:gradients] rather than the bare atom, because a list of fields is what metric/5 applies as the argument list of the metric function. The docs call this out.

Why the init traces the backward pass

Axon.Loop.run/4 compiles the batch function once, with the step state returned by init_fn as the template, and then feeds each step's output back in. Strict compilation requires the :gradients entry the init produces to have the same shapes and types as the one the step produces. zeros_like(trainable_parameters) is not a safe template: Nx.Defn.Grad seeds the backward pass with an f32 constant and never casts back to the parameter type, and the static and dynamic loss scales unscale by multiplying with an f32 scalar, so for bf16 or f16 parameters the gradient type can differ from the parameter type. Instead the init does what it already does for y_pred: it traces the computation (here value_and_grad plus unscale_grads) and reads only shape and type off the result with zeros_like/1. The traced expression is never referenced by the returned state, so it is never lowered. The cost is a little extra Elixir-side tracing at init time, only when the option is on.

Tradeoffs and limitations

• Keeping the gradients costs a parameter-sized buffer between iterations and stops the compiler from fusing the gradients into the optimizer update, which is why the option is opt-in and off by default.
• Checkpoints written by Axon.Loop.checkpoint/2 serialize the whole step state, so they grow by the size of the trainable parameters when the option is on.
• :gradients is not added to the donatable step-state keys: the step never reads the previous gradients back, so donating them would be pointless. donate_state?: true keeps working with the option on.
• Gradient accumulation across micro-batches is a different feature and is out of scope. This only exposes the per-iteration gradients.
• Unrelated, but noticed while writing the frozen-parameter test: a %Axon.ModelState{} with frozen_parameters set by Axon.ModelState.freeze/2 loses that field when passed back through a model's init_fn, because merge_model_state!/2 in Axon.Compiler only merges :data. That is pre-existing on main and not touched here; the test freezes from inside a {init_fn, apply_fn} model tuple so the frozen state actually reaches the step.

Tests

test/axon/loop_test.exs gains a describe "keep_gradients?" block:

• the default step state has no :gradients key, before and after a step;
• for a single bias-free dense layer with sgd(learning_rate: 0.1), the kept gradient equals both the analytic gradient (5w for mean((w * x)^2) with x = [1, 2]) and (w - w') / lr, so it is exactly what the optimizer received; the init template is zeros with the parameter's shape and type, and the init template and step output have identical shapes and types for every leaf;
• with dense_0 frozen, only dense_1 gradients are kept and dense_0 is unchanged by the step;
• the gradients kept under :static and :dynamic loss scaling match those under :identity, and in each case the init template matches the step output;
• under a bf16 mixed-precision policy the parameters are bf16 while the gradients are f32, and the init template still matches the step output for each loss scale, which guards the backward-pass tracing in init_fn (zeros of the parameter type would fail this);
• trainer/4 with the option exposes the gradients to a [:gradients] metric and to an :iteration_completed handler through a full Axon.Loop.run/4, which exercises the strict compilation path;
• donate_state?: true with the option on converges to the same model state and gradients as the non-donating run, and the returned state is not donatable.

All six substantive tests fail on main without the library change. mix test passes on the default backend (883 tests) and USE_EXLA=1 mix test test/axon/loop_test.exs passes as well, including the exla_only buffer-donation tests.

🤖 Generated with Claude Code

[polvalente]

Suggested change

	i: Nx.tensor(0),
	i: Nx.u64(0),