Implement a caching based version for the training

ext/LuxEnzymeExt.jl

@@ -1,26 +1,92 @@
 module LuxEnzymeExt
 
 using ADTypes: AutoEnzyme
-using Enzyme: Enzyme
+using ConcreteStructs: @concrete
+using Enzyme: Enzyme, Active, Const, Duplicated
 using Lux: Lux
 using Setfield: @set!
 
+@concrete struct CachedEnzymeExtras
+ dparameters
+ forward
+ reverse
+end
+
+# Case I: We have CachedEnzymeExtras and objective_function is unchanged.
+function Lux.Experimental.compute_gradients(::AutoEnzyme, objective_function::F, data,
+ ts::Lux.Experimental.TrainState{<:CachedEnzymeExtras, F}) where {F}
+ Lux.__recursive_make_zero!(ts.cache.dparameters)
+ loss, st_new, stats = __compute_gradients!(
+ ts.cache.forward, ts.cache.reverse, objective_function,
+ ts.model, ts.parameters, ts.cache.dparameters, ts.states, data)
+ ts_new = __construct_new_trainstate(
+ st_new, ts.states, ts.cache.forward, ts.cache.reverse,
+ ts, objective_function, ts.cache.dparameters)
+ return ts.cache.dparameters, loss, stats, ts_new
+end
+
+# Case II: We have CachedEnzymeExtras and objective_function is changed.
+function Lux.Experimental.compute_gradients(::AutoEnzyme, objective_function::F, data,
+ ts::Lux.Experimental.TrainState{<:CachedEnzymeExtras}) where {F}
+ forward, reverse = Enzyme.autodiff_thunk(
+ Enzyme.ReverseSplitWithPrimal, Const{typeof(objective_function)},
+ Active, Const{typeof(ts.model)}, Duplicated{typeof(ts.parameters)},
+ Const{typeof(ts.states)}, Const{typeof(data)})
+
+ Lux.__recursive_make_zero!(ts.cache.dparameters)
+ loss, st_new, stats = __compute_gradients!(
+ forward, reverse, objective_function, ts.model,
+ ts.parameters, ts.cache.dparameters, ts.states, data)
+
+ ts_new = __construct_new_trainstate(
+ st_new, ts.states, forward, reverse, ts, objective_function, ts.cache.dparameters)
+ return ts.cache.dparameters, loss, stats, ts_new
+end
+
+# Case III: Nothing is cached
 function Lux.Experimental.compute_gradients(::AutoEnzyme, objective_function::F, data,
  ts::Lux.Experimental.TrainState) where {F}
- dps = Enzyme.make_zero(ts.parameters)
- fwd, rev = Enzyme.autodiff_thunk(
- Enzyme.ReverseSplitWithPrimal, Enzyme.Const{typeof(objective_function)},
- Enzyme.Active, Enzyme.Const{typeof(ts.model)},
- Enzyme.Duplicated{typeof(ts.parameters)},
- Enzyme.Const{typeof(ts.states)}, Enzyme.Const{typeof(data)})
- tape, (loss, st_new, stats), shadow_result = fwd(
- Enzyme.Const(objective_function), Enzyme.Const(ts.model),
- Enzyme.Duplicated(ts.parameters, dps), Enzyme.Const(ts.states), Enzyme.Const(data))
- rev(Enzyme.Const(objective_function), Enzyme.Const(ts.model),
- Enzyme.Duplicated(ts.parameters, dps), Enzyme.Const(ts.states), Enzyme.Const(data),
- (one(loss), Enzyme.make_zero(st_new), Enzyme.make_zero(stats)), tape)
- @set! ts.states = st_new
- return dps, loss, stats, ts
+ dps = Lux.__recursive_make_zero(ts.parameters)
+ forward, reverse = Enzyme.autodiff_thunk(
+ Enzyme.ReverseSplitWithPrimal, Const{typeof(objective_function)},
+ Active, Const{typeof(ts.model)}, Duplicated{typeof(ts.parameters)},
+ Const{typeof(ts.states)}, Const{typeof(data)})
+
+ loss, st_new, stats = __compute_gradients!(
+ forward, reverse, objective_function, ts.model, ts.parameters, dps, ts.states, data)
+ ts_new = __construct_new_trainstate(
+ st_new, ts.states, forward, reverse, ts, objective_function, dps)
+ return dps, loss, stats, ts_new
+end
+
+function __compute_gradients!(
+ forward::F, reverse::R, obj_fn::O, model, ps, dps, st, data) where {F, R, O}
+ pps = Duplicated(ps, dps)
+ args = (Const(obj_fn), Const(model), pps, Const(st), Const(data))
+ tape, (loss, st_new, stats), shadow_result = forward(args...)
+ reverse(args...,
+ (one(loss), Lux.__recursive_make_zero(st_new), Lux.__recursive_make_zero(stats)),
+ tape)
+ return loss, st_new, stats
+end
+
+# If `st_new` is of a new type, we will have to recompute the cache anyway. Force it
+# my not storing the objective function.
+function __construct_new_trainstate(
+ st_new::S, ::S, forward::F, reverse::R, ts::Lux.Experimental.TrainState,
+ objective_fn::O, dps) where {S, F, R, O}
+ cache = CachedEnzymeExtras(dps, forward, reverse)
+ return Lux.Experimental.TrainState(
+ cache, ts.objective_function, ts.model, ts.parameters,
+ st_new, ts.optimizer_state, ts.step + 1)
+end
+
+function __construct_new_trainstate(
+ st_new, _, forward::F, reverse::R, ts::Lux.Experimental.TrainState,
+ objective_fn::O, dps) where {F, R, O}
+ cache = CachedEnzymeExtras(dps, nothing, nothing)
+ return Lux.Experimental.TrainState(
+ cache, nothing, ts.model, ts.parameters, st_new, ts.optimizer_state, ts.step + 1)
 end
 
 end

ext/LuxOptimisersExt.jl

@@ -36,10 +36,18 @@ function Lux.Experimental.TrainState(
  return Lux.Experimental.TrainState(nothing, nothing, model, ps, st, st_opt, 0)
 end
 
-function Lux.Experimental.apply_gradients(ts::Lux.Experimental.TrainState, grads)
- optimizer_state, ps = Optimisers.update(ts.optimizer_state, ts.parameters, grads)
- return Lux.Experimental.TrainState(ts.cache, ts.objective_function, ts.model,
- ps, ts.states, optimizer_state, ts.step + 1)
+function Lux.Experimental.apply_gradients(
+ ts::Lux.Experimental.TrainState, grads, update_inplace=false)
+ if update_inplace
+ optimizer_state, ps = Optimisers.update(ts.optimizer_state, ts.parameters, grads)
+ return Lux.Experimental.TrainState(ts.cache, ts.objective_function, ts.model,
+ ps, ts.states, optimizer_state, ts.step + 1)
+ else
+ Optimisers.update!(ts.optimizer_state, ts.parameters, grads)
+ return Lux.Experimental.TrainState(
+ ts.cache, ts.objective_function, ts.model, ts.parameters,
+ ts.states, ts.optimizer_state, ts.step + 1)
+ end
 end
 
 # DistributedUtils

src/contrib/training.jl

@@ -25,14 +25,15 @@ Internal fields:
 end
 
 """
- apply_gradients(ts::TrainState, grads)
+ apply_gradients(ts::TrainState, grads, update_inplace::Bool=false)
 
 Update the parameters stored in `ts` using the gradients `grads`.
 
 ## Arguments
 
  - `ts`: [`TrainState`](@ref) object.
  - `grads`: Gradients of the loss function wrt `ts.params`.
+ - `update_inplace`: Whether to update the parameters inplace or not.
 
 ## Returns
 
@@ -73,6 +74,17 @@ A 4-Tuple containing:
  - `loss`: Loss from the objective function.
  - `stats`: Any computed statistics from the objective function.
  - `ts`: Updated Training State.
+
+## Special Notes on Backends
+
+ - `AutoEnzyme`: `mode` is always ignored.
+
+!!! danger
+
+ `grads` returned by this function might be aliased by the implementation of the gradient
+ backend. For example, if you cache the `grads` from step `i`, the new gradients
+ returned in step `i + 1` might be aliased by the old gradients. If you want to prevent
+ this, simply use `copy(grads)` or `deepcopy(grads)` to make a copy of the gradients.
 """
 function compute_gradients(ad::ADTypes.AbstractADType, ::F, _, ::TrainState) where {F}
  return __maybe_implemented_compute_gradients(ad)

src/utils.jl

@@ -287,3 +287,20 @@ end
 
 @inline __size(x::AbstractArray) = size(x)
 @inline __size(x::T) where {T} = hasmethod(size, Tuple{T}) ? size(x) : nothing
+
+@inline __recursive_make_zero(x::AbstractArray{<:Number}) = zero(x)
+@inline __recursive_make_zero(x::AbstractArray) = map(__recursive_make_zero, x)
+@inline __recursive_make_zero(x::Tuple) = map(__recursive_make_zero, x)
+@inline __recursive_make_zero(x::NamedTuple{fields}) where {fields} = NamedTuple{fields}(map(
+ __recursive_make_zero, values(x)))
+@inline __recursive_make_zero(::Nothing) = nothing
+@inline __recursive_make_zero(v::Val) = v
+@inline __recursive_make_zero(x) = fmap(__recursive_make_zero, x)
+
+@inline __recursive_make_zero!(x::AbstractArray{<:Number}) = fill!(x, zero(eltype(x)))
+@inline __recursive_make_zero!(x::AbstractArray) = map(__recursive_make_zero!, x)
+@inline __recursive_make_zero!(x::Tuple) = map(__recursive_make_zero!, x)
+@inline __recursive_make_zero!(x::NamedTuple{fields}) where {fields} = NamedTuple{fields}(map(
+ __recursive_make_zero!, values(x)))
+@inline __recursive_make_zero!(::Nothing) = nothing
+@inline __recursive_make_zero!(x) = fmap(__recursive_make_zero!, x)