Shape-bucketed compilation cache for training with a shared `TrainState`

[AzamatB]

Following up from Slack with @wsmoses and @avik-pal on training under AutoReactant() when one input axis has variable length across samples.

The workable approach is to choose N bucket sizes, pad each input up to its nearest bucket, and cache one compiled executable per bucket — bounding compilations to N and keeping padding overhead under the user's control. The Qwen3 inference example (examples/Qwen3/main.jl, cf. CachedReactantThunks) already does this for generation, keying Reactant.Compiler.Thunks on (config, size) while keeping ps/st in a separate struct so all compiled variants share the same parameter buffers.

The ask is a training-side equivalent. Lux.Training.single_train_step!(AutoReactant(), ...) currently specializes on the shapes of data and caches inside TrainState, so the naive bucketed version would allocate one TrainState per bucket — duplicating ps, optimizer state, and gradient scratch N times, which doesn't fit in VRAM for reasonable model sizes. What's needed is a supported pattern where one TrainState holds a single copy of parameters and optimizer state, and compiled train-step artifacts are cached separately, keyed on input shape (or a user-supplied bucket key).

[avik-pal]

What's needed is a supported pattern where one TrainState holds a single copy of parameters and optimizer state, and compiled train-step artifacts are cached separately, keyed on input shape (or a user-supplied bucket key).

Essentially all we need here is a fast way to convert input data to a key of structure of shapes and types. Then we check our cache to see if the corresponding thunk was already compiled. We can reuse the ps and dps across the compiled thunks without any issue.

The workable approach is to choose N bucket sizes, pad each input up to its nearest bucket, and cache one compiled executable per bucket

This is quite involved and I would rather have the user do it. Especially because how you pad the data is problem dependent. Sometimes you might want to pad the batch, sometimes the sequence, etc.

[AzamatB]

Here's how I've implemented this.

using Lux, Reactant, Optimisers, Enzyme
using Lux.Training: TrainState
using Static: False
using Accessors: @set

mutable struct CompilationCache{K}
    slots::Dict{K,Any}
end

CompilationCache{K}() where {K} = CompilationCache{K}(Dict{K,Any}())

function Base.get(cache::CompilationCache{K}, key::K, default) where {K}
    return get(cache.slots, key, default)
end

function Base.getindex(cache::CompilationCache{K}, key::K) where {K}
    return cache.slots[key]
end

function Base.setindex!(cache::CompilationCache{K}, value, key::K) where {K}
    cache.slots[key] = value
end

"""
    bucketed_single_train_step!(loss, data, train_state::TrainState, cache::CompilationCache{K}, key::K)

One real `TrainState`, one Reactant train-step cache per `key`. Assumes:
  - same model / loss type / optimizer / compile options across calls,
  - `data` is already padded to the bucket shape for `key`,
  - `key` captures every shape-varying axis of the batch.
"""
function bucketed_single_train_step!(
    loss, data, train_state::TrainState, cache::CompilationCache{K}, key::K
) where {K}
    # swap in the compiled cache for this bucket
    cache_k = get(cache, key, nothing)
    ts = @set train_state.cache = cache_k

    (grads, loss_value, stats, ts) = Lux.Training.single_train_step!(
        AutoReactant(), loss, data, ts; return_gradients=False()
    )
    # stash the (now-populated) cache back under this bucket
    cache[key] = ts.cache

    # hand back a TrainState that still owns the trained ps/st/opt_state but no per-bucket cache
    ts_out = @set ts.cache = nothing

    return (grads, loss_value, stats, ts_out)
end

At this point, I'd be happy if the authors simply kept TrainState.cache as a stable, writable field – i.e. treated it as a semi-public API.