Repository: LuxDL/Lux.jl
Pull request #1273: feat: support for ForwardDiff training
Author: login: hstrey
Created: 2025-03-20T20:14:51Z | Merged: 2025-03-26T18:23:41Z

added LuxForwardDiffExt

It is working, but I may not have implemented the gradient and loss eval in the most efficient way



Diff:
diff --git a/Project.toml b/Project.toml
index afc2b117c6..2842f403ac 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.10.1"
+version = "1.11.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -11,6 +11,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
@@ -82,6 +83,7 @@ ChainRulesCore = "1.25"
 Compat = "4.16"
 ComponentArrays = "0.15.22"
 ConcreteStructs = "0.2.3"
+DiffResults = "1.1"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.35"
 EnzymeCore = "0.8.8"
diff --git a/src/Lux.jl b/src/Lux.jl
index 5255db7f93..f8aaaa3fa2 100644
--- a/src/Lux.jl
+++ b/src/Lux.jl
@@ -82,6 +82,7 @@ include("extended_ops.jl")
 # Training Helpers
 include("helpers/optimizers.jl")
 include("helpers/training.jl")
+include("helpers/forwarddiff_training.jl")
 
 # Experimental
 include("contrib/contrib.jl")
@@ -155,7 +156,8 @@ export Training
 
 export jacobian_vector_product, vector_jacobian_product
 export batched_jacobian
-export AutoEnzyme, AutoForwardDiff, AutoReverseDiff, AutoTracker, AutoZygote
+export AutoEnzyme,
+    AutoForwardDiff, AutoReverseDiff, AutoTracker, AutoZygote, AutoForwardDiff
 
 export BinaryCrossEntropyLoss,
     BinaryFocalLoss,
diff --git a/src/helpers/forwarddiff_training.jl b/src/helpers/forwarddiff_training.jl
new file mode 100644
index 0000000000..ef7b9ddab6
--- /dev/null
+++ b/src/helpers/forwarddiff_training.jl
@@ -0,0 +1,91 @@
+using ADTypes: AutoForwardDiff
+using DiffResults: DiffResults
+using ForwardDiff: ForwardDiff
+using Setfield: @set!
+using Static: True, False
+
+function Training.compute_gradients_impl(
+    ad::AutoForwardDiff, obj_fn::F, data, ts::Training.TrainState
+) where {F}
+    @assert ts.parameters isa AbstractArray "AutoForwardDiff only supports AbstractArray \
+                                             parameters, not $(typeof(ts.parameters)). To \
+                                             convert the parameter structure to an array \
+                                             use `ComponentArray(ps)`."
+
+    obj_fn_wrap, st_wrap, stats_wrap = Training.wrap_objective_function(
+        obj_fn, ts.model, ts.parameters, ts.states, data, True()
+    )
+
+    gradient_result = DiffResults.GradientResult(ts.parameters)
+    ForwardDiff.gradient!(
+        gradient_result, ps -> obj_fn_wrap(ts.model, ps, ts.states, data), ts.parameters
+    )
+
+    cache = Training.TrainingBackendCache(
+        ad, False(), gradient_result, (; obj_fn=obj_fn_wrap, st_wrap, stats_wrap)
+    )
+    @set! ts.cache = cache
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = st_wrap[]
+    return (
+        DiffResults.gradient(gradient_result),
+        DiffResults.value(gradient_result),
+        stats_wrap[],
+        ts,
+    )
+end
+
+const FORWARDDIFF_CACHE_TYPE = Training.TrainingBackendCache{
+    <:AutoForwardDiff,False,PS,<:NamedTuple{(:obj_fn, :st_wrap, :stats_wrap)}
+} where {PS}
+
+function Training.compute_gradients_impl(
+    ::AutoForwardDiff, obj_fn::F, data, ts::Training.TrainState{<:FORWARDDIFF_CACHE_TYPE,F}
+) where {F}
+    gradient_result = ts.cache.dparameters
+
+    ForwardDiff.gradient!(
+        gradient_result,
+        ps -> ts.cache.extras.obj_fn(ts.model, ps, ts.states, data),
+        ts.parameters,
+    )
+
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = ts.cache.extras.st_wrap[]
+
+    return (
+        DiffResults.gradient(gradient_result),
+        DiffResults.value(gradient_result),
+        ts.cache.extras.stats_wrap[],
+        ts,
+    )
+end
+
+function Training.compute_gradients_impl(
+    ::AutoForwardDiff,
+    obj_fn::F,
+    data,
+    ts::Training.TrainState{<:Training.TrainingBackendCache{<:AutoForwardDiff,False}},
+) where {F}
+    @warn "Detected calls to `compute_gradients(::AutoForwardDiff, ...)` with objective \
+           function that is changing across function calls. This can lead to the \
+           generation of slow code" maxlog = 1
+    gradient_result = ts.cache.dparameters
+
+    # We do exactly same thing as the first case but without caching the function
+    obj_fn_wrap, st_wrap, stats_wrap = Training.wrap_objective_function(
+        obj_fn, ts.model, ts.parameters, ts.states, data, False()
+    )
+
+    ForwardDiff.gradient!(
+        gradient_result, ps -> obj_fn_wrap(ts.model, ps, ts.states, data), ts.parameters
+    )
+
+    @set! ts.states = st_wrap[]
+    return (
+        DiffResults.gradient(gradient_result),
+        DiffResults.value(gradient_result),
+        stats_wrap[],
+        ts,
+    )
+end
diff --git a/src/helpers/training.jl b/src/helpers/training.jl
index 7a9438904b..537547a594 100644
--- a/src/helpers/training.jl
+++ b/src/helpers/training.jl
@@ -160,6 +160,7 @@ Compute the gradients of the objective function wrt parameters stored in `ts`.
 | `AutoReverseDiff(; compile)` | `ReverseDiff.jl` |
 | `AutoTracker`                | `Tracker.jl`     |
 | `AutoEnzyme`                 | `Enzyme.jl`      |
+| `AutoForwardDiff`            |                  |
 
 ## Arguments
 
@@ -185,6 +186,8 @@ A 4-Tuple containing:
   - `AutoReverseDiff(; compile=true)` is not supported for Lux models with non-empty state
     `st`. Additionally the returned stats must be empty (`NamedTuple()`). We catch these
     issues in most cases and throw an error.
+  - AutoForwardDiff only works with parameters that are AbstractArrays
+    (e.g. ps=ComponentVector(ps))
 
 !!! danger "Aliased Gradients"
 
diff --git a/test/helpers/training_tests.jl b/test/helpers/training_tests.jl
index a8a67532ad..338af8aabd 100644
--- a/test/helpers/training_tests.jl
+++ b/test/helpers/training_tests.jl
@@ -139,6 +139,56 @@ end
     end
 end
 
+@testitem "Training API ForwardDiff" setup = [SharedTestSetup] tags = [:misc] begin
+    using ADTypes, Optimisers, ComponentArrays
+
+    mse = MSELoss()
+
+    rng = StableRNG(12345)
+
+    x_data = randn(rng, Float32, 4, 32)
+    y_data = evalpoly.(x_data, ((1, 2, 3),)) .- evalpoly.(x_data, ((5, 2),))
+    y_data = (y_data .- minimum(y_data)) ./ (maximum(y_data) - minimum(y_data))
+    dataset = [(x_data[:, i], y_data[:, i]) for i in Iterators.partition(1:32, 8)]
+
+    model = Chain(
+        Dense(4, 32, tanh), BatchNorm(32), Dense(32, 32, tanh), BatchNorm(32), Dense(32, 4)
+    )
+
+    dataset_ = [(x, y) for (x, y) in dataset]
+    opt = Adam(0.001f0)
+
+    ps, st = Lux.setup(rng, model)
+    tstate = Training.TrainState(model, ComponentVector(ps), st, opt)
+
+    initial_loss = first(
+        mse(model, tstate.parameters, Lux.testmode(tstate.states), dataset_[1])
+    )
+
+    for epoch in 1:100, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.compute_gradients(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+        tstate = Training.apply_gradients!(tstate, grads)
+    end
+
+    for epoch in 1:100, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.single_train_step!(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+    end
+
+    for epoch in 1:100, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.single_train_step(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+    end
+
+    final_loss = first(mse(model, tstate.parameters, tstate.states, dataset_[1]))
+
+    @test final_loss * 50 < initial_loss
+end
+
 @testitem "Enzyme: Invalidate Cache on State Update" setup = [SharedTestSetup] tags = [
     :misc
 ] skip = :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin


Reviews:
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T14:06:32Z
user: login: hstrey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T18:37:47Z
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T19:05:29Z
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T19:05:57Z
user: login: avik-pal
body: Overall looks good, lets add a test for this
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-25T19:06:27Z
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T00:16:44Z
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T00:20:03Z
user: login: hstrey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T13:25:57Z
user: login: avik-pal
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T17:07:42Z

Review comments:
diff_hunk: @@ -73,6 +73,7 @@ LuxMPIExt = "MPI"
 LuxMPINCCLExt = ["CUDA", "MPI", "NCCL"]
 LuxReactantExt = ["Enzyme", "Reactant"]
 LuxReverseDiffExt = ["FunctionWrappers", "ReverseDiff"]
+LuxForwardDiffExt = "ForwardDiff"
path: Project.toml
user: login: avik-pal
body: ForwardDiff is already a main dependency, you can add the training dispatch to the main codebase without the extensions
created_at: 2025-03-25T14:06:31Z
author_association: MEMBER
diff_hunk: @@ -73,6 +73,7 @@ LuxMPIExt = "MPI"
 LuxMPINCCLExt = ["CUDA", "MPI", "NCCL"]
 LuxReactantExt = ["Enzyme", "Reactant"]
 LuxReverseDiffExt = ["FunctionWrappers", "ReverseDiff"]
+LuxForwardDiffExt = "ForwardDiff"
path: Project.toml
user: login: hstrey
body: @avik-pal Should this training dispatch go into src/helpers/training.jl?
created_at: 2025-03-25T18:37:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -160,6 +161,7 @@ Compute the gradients of the objective function wrt parameters stored in `ts`.
 | `AutoReverseDiff(; compile)` | `ReverseDiff.jl` |
 | `AutoTracker`                | `Tracker.jl`     |
 | `AutoEnzyme`                 | `Enzyme.jl`      |
+| `AutoForwardDiff`            | `ForwardDiff.jl` |
path: src/helpers/training.jl
user: login: avik-pal
body: ```suggestion
| `AutoForwardDiff`            |                  |
```

we already have it loaded so no extra packages are needed
created_at: 2025-03-25T19:05:28Z
author_association: MEMBER
diff_hunk: @@ -1,13 +1,14 @@
 module Training
 
-using ADTypes: AbstractADType, AutoEnzyme, AutoReverseDiff, AutoTracker, AutoZygote
+using ADTypes: AbstractADType, AutoEnzyme, AutoReverseDiff, AutoTracker, AutoZygote, AutoForwardDiff
 using Compat: @compat
 using ConcreteStructs: @concrete
 using FastClosures: @closure
 using Functors: Functors, fmap
 using Optimisers: Optimisers
 using Setfield: @set!
 using Static: StaticBool, Static, False, True, static
+using ForwardDiff
path: src/helpers/training.jl
user: login: avik-pal
body: ```suggestion
using ForwardDiff: ForwardDiff
```
created_at: 2025-03-25T19:05:57Z
author_association: MEMBER
diff_hunk: @@ -139,6 +139,58 @@ end
     end
 end
 
+@testitem "Training API ForwardDiff" setup = [SharedTestSetup] tags = [:misc] begin
+    using ADTypes, Optimisers, ComponentArrays
+
+    mse = MSELoss()
+
+    rng = StableRNG(12345)
+
+    x_data = randn(rng, Float32, 4, 32)
+    y_data = evalpoly.(x_data, ((1, 2, 3),)) .- evalpoly.(x_data, ((5, 2),))
+    y_data = (y_data .- minimum(y_data)) ./ (maximum(y_data) - minimum(y_data))
+    dataset = [(x_data[:, i], y_data[:, i]) for i in Iterators.partition(1:32, 8)]
+
+    model = Chain(
+        Dense(4, 32, tanh),
+        BatchNorm(32),
+        Dense(32, 32, tanh),
+        BatchNorm(32),
+        Dense(32, 4),
+    )
+
+    dataset_ = [(x, y) for (x, y) in dataset]
+    opt = Adam(0.001f0)
+
+    ps, st = Lux.setup(rng, model)
+    tstate = Training.TrainState(model, ComponentVector(ps), st, opt)
+
+    initial_loss = first(mse(model, tstate.parameters, tstate.states, dataset_[1]))
+
+    for epoch in 1:1000, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.compute_gradients(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+        tstate = Training.apply_gradients!(tstate, grads)
+    end
+
+    for epoch in 1:1000, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.single_train_step!(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+    end
+
+    for epoch in 1:1000, (x, y) in dataset_
+        grads, loss, _, tstate = allow_unstable() do
+            Training.single_train_step(AutoForwardDiff(), mse, (x, y), tstate)
+        end
+    end
+
+    final_loss = first(
+        mse(model, tstate.parameters, tstate.states, dataset_[1])
+    )
path: test/helpers/training_tests.jl
user: login: avik-pal
body: add a check for final_loss < ...
created_at: 2025-03-26T00:16:43Z
author_association: MEMBER
diff_hunk: @@ -208,6 +212,15 @@ function maybe_wrap_adtype(
                          Enzyme.jl (`AutoEnzyme`)."))
 end
 
+function compute_gradients_impl(
+    ::AutoForwardDiff, objective_function::F, data, ts::Lux.Training.TrainState
+) where {F}
+    grads = ForwardDiff.gradient(ps -> objective_function(ts.model,ps,ts.states,data)[1], ts.parameters);
+    loss, st, stats = objective_function(ts.model,ts.parameters,ts.states,data)
+    @set! ts.states = st
path: src/helpers/training.jl
user: login: avik-pal
body: This is incorrect, we shouldn't do 2 evals of `objective_function` here (apart from being inefficient, it will actually give incorrect results for normalization layers). The easiest way would be to use a closure approach similar to how it is done in the enzyme extension
created_at: 2025-03-26T00:20:03Z
author_association: MEMBER
diff_hunk: @@ -208,6 +212,15 @@ function maybe_wrap_adtype(
                          Enzyme.jl (`AutoEnzyme`)."))
 end
 
+function compute_gradients_impl(
+    ::AutoForwardDiff, objective_function::F, data, ts::Lux.Training.TrainState
+) where {F}
+    grads = ForwardDiff.gradient(ps -> objective_function(ts.model,ps,ts.states,data)[1], ts.parameters);
+    loss, st, stats = objective_function(ts.model,ts.parameters,ts.states,data)
+    @set! ts.states = st
path: src/helpers/training.jl
user: login: hstrey
body: @avik-pal I did what you suggested, but I am not completely sure whether my implementation is sound.  I used the DiffResults interface for ForwardDiff to get both the value and the grad at the same time.  I also used the wrapper that is used by the Enzyme extension. 
On another note, when I benchmarked the different AD interfaces, for small NNs ForwardDiff can be significantly faster.  I found that for our use case, we are 4x faster.  Maybe it makes sense to implement ForwardDiff in a way that takes care of the named tuple issue by applying ComponentVector when necessary.  But that would mean that ComponentVector becomes a dependency.  One could also fix it at the ForwardDiff level.
The tests should pass now.
created_at: 2025-03-26T13:25:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -160,6 +161,8 @@ Compute the gradients of the objective function wrt parameters stored in `ts`.
 | `AutoReverseDiff(; compile)` | `ReverseDiff.jl` |
 | `AutoTracker`                | `Tracker.jl`     |
 | `AutoEnzyme`                 | `Enzyme.jl`      |
+| `AutoForwardDiff`            |                  |
+
path: src/helpers/training.jl
user: login: avik-pal
body: ```suggestion
```
created_at: 2025-03-26T17:07:07Z
author_association: MEMBER
diff_hunk: @@ -8,6 +8,7 @@ using Functors: Functors, fmap
 using Optimisers: Optimisers
 using Setfield: @set!
 using Static: StaticBool, Static, False, True, static
+using ForwardDiff: ForwardDiff, DiffResults
path: src/helpers/training.jl
user: login: avik-pal
body: ```suggestion
```
created_at: 2025-03-26T17:07:17Z
author_association: MEMBER
diff_hunk: @@ -185,6 +188,7 @@ A 4-Tuple containing:
   - `AutoReverseDiff(; compile=true)` is not supported for Lux models with non-empty state
     `st`. Additionally the returned stats must be empty (`NamedTuple()`). We catch these
     issues in most cases and throw an error.
+  - AutoForwardDiff only works with parameters that are converted to ComponentVectors (e.g. ps=ComponentVector(ps))
path: src/helpers/training.jl
user: login: avik-pal
body: ```suggestion
  - AutoForwardDiff only works with parameters that are AbstractArrays
    (e.g. ps=ComponentVector(ps))
```
created_at: 2025-03-26T17:07:35Z
author_association: MEMBER

PR comments:
user: login: hstrey
created_at: 2025-03-24T22:24:18Z
body: @avik-pal do you want me to add tests?
author_association: CONTRIBUTOR
user: login: hstrey
created_at: 2025-03-25T19:33:27Z
body: @avik-pal I added AutoForwardDiff to the training test.  I was not sure whether ForwardDiff works on gpus, so I added AutoForwardDiff to the list of adtypes that are not run on gpus - at least this is how I understood the code.
author_association: CONTRIBUTOR
user: login: hstrey
created_at: 2025-03-25T22:08:27Z
body: @avik-pal I checked the logs.  The ForwardDiff test fails because ForwardDiff.gradient chokes on the fact that ts.parameters is a named tuple.  When I tested it, I had used ComponentVector.  Should I do that here too?

author_association: CONTRIBUTOR
user: login: avik-pal
created_at: 2025-03-25T22:23:18Z
body: I think it is reasonable to only allow ComponentVectors for ForwardDiff, so you can just skip the NamedTuple tests for forwarddiff.
author_association: MEMBER
user: login: hstrey
created_at: 2025-03-25T22:43:06Z
body: @avik-pal, all the tests I added involved named tuples. Should I write a separate test to test just the ForwardDiff training interface?
author_association: CONTRIBUTOR
user: login: avik-pal
created_at: 2025-03-25T22:44:15Z
body: > @avik-pal, all the tests I added involved named tuples. Should I write a separate test to test just the ForwardDiff training interface?

yeah
author_association: MEMBER