Add training loop fusion

[seanmor5]

Closes #101.

By default, Axon.Loop.run/4 dispatches one computation per batch. Every iteration crosses the boundary between the BEAM and the compiler's runtime, and the step state crosses back with it. For small models that dispatch costs more than the training step itself.

fuse: true stacks a chunk of batches, moves it to the device at once, and steps through it in a device-side while loop which carries the step state and the metrics, so the loop synchronizes with the host once per chunk instead of once per batch:

model
|> Axon.Loop.trainer(:categorical_cross_entropy, :adam)
|> Axon.Loop.run(data, %{}, epochs: 10, compiler: EXLA, fuse: true)

The same number of bytes crosses to the device either way, they just cross in one transfer instead of one per batch. :fuse also takes a positive integer chunk size, which trades device memory for dispatch amortization, and defaults to 32.

Both paths run the same batch function, so a fused and an unfused loop compute the same thing — the tests assert that directly on parameters, metrics, and event counts.

Results

bench/loop_fusion.exs, EXLA CPU, batch size 32, 100 batches/epoch, steady state (compilation excluded):

model	params	unfused	fused	speedup
mlp-tiny	1.1K	266us/iter	68us/iter	3.9x
mlp-small	25K	306us/iter	115us/iter	2.7x
mlp-large	1.2M	1052us/iter	770us/iter	1.4x

Chunk size sweep for mlp-small: 1 → 0.99x, 4 → 1.23x, 8 → 1.51x, 16 → 2.02x, 32 → 2.55x, 64 → 2.67x.

Why the data is device resident

The first design pulled each batch from the host inside the loop with Nx.runtime_call/4, which works (including inside while, on EXLA :host and :cuda), but measures at roughly twice the cost of the dispatch it removes. Timing the three shapes of the same loop on a small model:

	per iteration
unfused	66us
fused, batch pulled through a host callback	141us
fused, chunk resident on the device	1.3us

So the fused path uses no host callbacks at all.

Callbacks and hooks

• Handlers on :started, :epoch_started, :epoch_completed, :epoch_halted, :halted and :completed are unaffected: host side, between epochs, complete loop state.
• :iteration_started and :iteration_completed fire on the host once the chunk containing their iteration is done, in the same order the unfused loop fires them, with real counters and the metrics as of their iteration, recorded into a device-side buffer during the chunk. Their step_state is nil, since it only exists on the device while a chunk runs; metrics and handler metadata they return are kept, a returned step_state is discarded.
• Halting from an iteration handler takes effect at the end of the chunk it was requested in. fuse: 1 halts on the exact iteration, as does halting from :epoch_completed.

A trainer(log: 5) loop with validate/4 and early_stop/3 produces byte-identical log output and identical parameters fused and unfused.

Limitations

• Every batch must have the same shape and type, since a chunk is stacked into one tensor. Ragged data raises a fusion-specific error rather than failing inside Nx.stack/2.
• Requires jit_compile?: true, which is the default.

Tests

test/axon/loop/fused_test.exs — 19 tests, most of which assert the fused loop matches the unfused one exactly (parameters, metrics, event counts, halting paths), plus chunk padding, data exhaustion, :iterations bounds, batch consumption, and stream halting. Suite is green under both Nx.Defn.Evaluator and USE_EXLA=1.

🤖 Generated with Claude Code