fix: make enzyme testing opt-in for now (#1041)

* fix: make enzyme testing opt-in for now * test: mark certain tests as broken and unbroken
LuxDL · Nov 5, 2024 · 4079372 · 4079372 · avik-pal · Nov 5, 2024
1 parent 900c21c
commit 4079372
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 26 deletions.
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -48,7 +48,7 @@ Lux = "1"
 LuxCUDA = "0.3.2"
 LuxCore = "1"
 LuxLib = "1.3.4"
-LuxTestUtils = "1.4"
+LuxTestUtils = "1.5"
 MLDataDevices = "1.4"
 Optimisers = "0.3.3"
 Pkg = "1.10"

diff --git a/lib/LuxLib/test/Project.toml b/lib/LuxLib/test/Project.toml
@@ -46,7 +46,7 @@ Hwloc = "3.2"
 InteractiveUtils = "<0.0.1, 1"
 JLArrays = "0.1.5"
 LoopVectorization = "0.12.171"
-LuxTestUtils = "1.4"
+LuxTestUtils = "1.5"
 MKL = "0.7"
 MLDataDevices = "1.0.0"
 NNlib = "0.9.21"

diff --git a/lib/LuxTestUtils/Project.toml b/lib/LuxTestUtils/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <[email protected]>"]
-version = "1.4.0"
+version = "1.5.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/lib/LuxTestUtils/src/LuxTestUtils.jl b/lib/LuxTestUtils/src/LuxTestUtils.jl
@@ -37,17 +37,14 @@ try
     using Enzyme: Enzyme
     __ftest(x) = x
     Enzyme.autodiff(Enzyme.Reverse, __ftest, Enzyme.Active, Enzyme.Active(2.0))
+    # XXX: Enzyme has been causing some issues lately. Let's just disable it for now.
+    #      We still have opt-in testing available for Enzyme.
     # XXX: Lift this once Enzyme supports 1.11 properly
-    global ENZYME_TESTING_ENABLED = v"1.10-" ≤ VERSION < v"1.11-"
+    global ENZYME_TESTING_ENABLED = false # v"1.10-" ≤ VERSION < v"1.11-"
 catch err
     global ENZYME_TESTING_ENABLED = false
 end
 
-if !ENZYME_TESTING_ENABLED
-    @warn "`Enzyme.jl` is currently not functional on $(VERSION) either because it errored \
-           or the current version is a prerelease. Enzyme tests will be skipped..."
-end
-
 include("test_softfail.jl")
 include("utils.jl")
 include("autodiff.jl")

diff --git a/lib/LuxTestUtils/src/autodiff.jl b/lib/LuxTestUtils/src/autodiff.jl
@@ -98,7 +98,7 @@ Test the gradients of `f` with respect to `args` using the specified backends.
 | ReverseDiff.jl | `AutoReverseDiff()` | ✔   | ✖   |                   |
 | ForwardDiff.jl | `AutoForwardDiff()` | ✔   | ✖   | `len ≤ 100`       |
 | FiniteDiff.jl  | `AutoFiniteDiff()`  | ✔   | ✖   | `len ≤ 100`       |
-| Enzyme.jl      | `AutoEnzyme()`      | ✔   | ✖   | Only Reverse Mode |
+| Enzyme.jl      | `AutoEnzyme()`      | ✖   | ✖   | Only Reverse Mode |
 
 ## Arguments
 
@@ -115,6 +115,7 @@ Test the gradients of `f` with respect to `args` using the specified backends.
     overrides any `broken` kwargs. Alternatively, a list of backends can be passed to
     `soft_fail` to allow soft_fail tests for only those backends.
   - `enzyme_set_runtime_activity`: If `true`, then activate runtime activity for Enzyme.
+  - `enable_enzyme_reverse_mode`: If `true`, then enable reverse mode for Enzyme.
   - `kwargs`: Additional keyword arguments to pass to `check_approx`.
 
 ## Example
@@ -131,6 +132,7 @@ julia> test_gradients(f, 1.0, x, nothing)
 function test_gradients(f, args...; skip_backends=[], broken_backends=[],
         soft_fail::Union{Bool, Vector}=false,
         enzyme_set_runtime_activity::Bool=false,
+        enable_enzyme_reverse_mode::Bool=false,
         # Internal kwargs start
         source::LineNumberNode=LineNumberNode(0, nothing),
         test_expr::Expr=:(check_approx(∂args, ∂args_gt; kwargs...)),
@@ -148,7 +150,7 @@ function test_gradients(f, args...; skip_backends=[], broken_backends=[],
         total_length ≤ 100 && push!(backends, AutoForwardDiff())
         total_length ≤ 100 && push!(backends, AutoFiniteDiff())
         # TODO: Move Enzyme out of here once it supports GPUs
-        if ENZYME_TESTING_ENABLED
+        if enable_enzyme_reverse_mode || ENZYME_TESTING_ENABLED
             mode = enzyme_set_runtime_activity ?
                    Enzyme.set_runtime_activity(Enzyme.Reverse) :
                    Enzyme.Reverse

diff --git a/test/Project.toml b/test/Project.toml
@@ -59,7 +59,7 @@ Logging = "1.10"
 LoopVectorization = "0.12.171"
 LuxCore = "1.0"
 LuxLib = "1.3.4"
-LuxTestUtils = "1.4"
+LuxTestUtils = "1.5"
 MLDataDevices = "1.3"
 MLUtils = "0.4.3"
 NNlib = "0.9.24"

diff --git a/test/contrib/debug_tests.jl b/test/contrib/debug_tests.jl
@@ -15,18 +15,24 @@
         model_debug = Lux.Experimental.@debug_mode model
 
         @test_throws DimensionMismatch model_debug(x, ps, st)
-        @test_logs (:info,) (:error,
-            "Layer Dense(1 => 1) failed!! This layer is present at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2).") match_mode=:any try
-            model_debug(x, ps, st)
-        catch
+        # XXX this is a bit flaky in CI on 1.11+
+        if VERSION < v"1.11-"
+            @test_logs (:info,) (:error,
+                "Layer Dense(1 => 1) failed!! This layer is present at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2).") match_mode=:any try
+                model_debug(x, ps, st)
+            catch
+            end
         end
 
         model_debug = Lux.Experimental.@debug_mode model error_check=false
 
         @test_throws DimensionMismatch model_debug(x, ps, st)
-        @test_logs min_level=Logging.Error try
-            model_debug(x, ps, st)
-        catch
+        # XXX this is a bit flaky in CI on 1.11+
+        if VERSION < v"1.11-"
+            @test_logs min_level=Logging.Error try
+                model_debug(x, ps, st)
+            catch
+            end
         end
 
         model_fixed = Chain(
@@ -38,8 +44,11 @@
 
         model_fixed_debug = Lux.Experimental.@debug_mode model_fixed
 
-        @test_logs min_level=Logging.Error Zygote.gradient(
-            sum ∘ first ∘ model_fixed_debug, x, ps, st)
+        # XXX this is a bit flaky in CI on 1.11+
+        if VERSION < v"1.11-"
+            @test_logs min_level=Logging.Error Zygote.gradient(
+                sum ∘ first ∘ model_fixed_debug, x, ps, st)
+        end
     end
 end
 

diff --git a/test/helpers/loss_tests.jl b/test/helpers/loss_tests.jl
@@ -274,11 +274,13 @@ end
 
             __f = Base.Fix2(FocalLoss(), y)
             # FD will lead to out of domain errors
-            @test_gradients(__f, ŷ;
-                atol=1.0f-3,
-                rtol=1.0f-3,
-                skip_backends=[AutoFiniteDiff()],
-                broken_backends=ongpu ? [AutoTracker()] : [])
+            broken_backends = if VERSION ≥ v"1.11-"
+                []
+            else
+                ongpu ? [AutoTracker()] : []
+            end
+            @test_gradients(__f, ŷ; atol=1.0f-3, rtol=1.0f-3,
+                skip_backends=[AutoFiniteDiff()], broken_backends)
         end
     end
 end