Merge branch 'main' into main

LuxDL · Nov 6, 2024 · c8fa306 · c8fa306
2 parents 03b0d68 + 4079372
commit c8fa306
Show file tree

Hide file tree

Showing 22 changed files with 42 additions and 63 deletions.
diff --git a/.buildkite/benchmarks.yml b/.buildkite/benchmarks.yml
@@ -1,5 +1,6 @@
 steps:
   - group: ":racehorse: Benchmarks"
+    if: build.message !~ /\[skip benchmarks\]/ && build.message !~ /\[skip ci\]/ && !build.pull_request.draft
     steps:
       - label: "CPU: Run Benchmarks with {{matrix.threads}} thread(s)"
         matrix:

diff --git a/.buildkite/testing.yml b/.buildkite/testing.yml
@@ -1,7 +1,7 @@
 steps:
   - group: ":julia: (Lux) CUDA GPU"
     steps:
-      - label: ":julia: Julia {{matrix.julia}} + CUDA GPU"
+      - label: ":julia: Julia {{matrix.julia}} + CUDA GPU + {{matrix.group}}"
         plugins:
           - JuliaCI/julia#v1:
               version: "{{matrix.julia}}"
@@ -45,13 +45,16 @@ steps:
               include(joinpath(dir, "../test/runtests.jl"))'
         env:
           BACKEND_GROUP: "CUDA"
+          LUX_TEST_GROUP: "{{matrix.group}}"
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip ci\]/
         timeout_in_minutes: 120
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
+            group:
+              - "!reactant"
+              - "reactant"
 
   - group: ":julia: (Lux) AMD GPU"
     steps:
@@ -105,7 +108,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
 env:

diff --git a/.buildkite/testing_luxcuda.yml b/.buildkite/testing_luxcuda.yml
@@ -23,7 +23,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
 env:

diff --git a/.buildkite/testing_luxlib.yml b/.buildkite/testing_luxlib.yml
@@ -45,7 +45,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
   - group: ":julia: (LuxLib) AMD GPU"
@@ -95,7 +94,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
 env:

diff --git a/.buildkite/testing_luxtestutils.yml b/.buildkite/testing_luxtestutils.yml
@@ -25,7 +25,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
 env:

diff --git a/.buildkite/testing_mldatadevices.yml b/.buildkite/testing_mldatadevices.yml
@@ -26,7 +26,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
             group:
               - "CUDA"
@@ -60,7 +59,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
   - group: ":julia: (MLDataDevices) Metal GPU"
@@ -91,7 +89,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
   - group: ":julia: (MLDataDevices) oneAPI GPU"
@@ -122,7 +119,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
 env:

diff --git a/.buildkite/testing_weightinitializers.yml b/.buildkite/testing_weightinitializers.yml
@@ -26,7 +26,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
               - "1"
 
   - group: ":julia: (WeightInitializers) AMD GPU"

diff --git a/.github/workflows/CI_LuxCUDA.yml b/.github/workflows/CI_LuxCUDA.yml
@@ -23,7 +23,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.10"
           - "1"
     steps:
       - uses: actions/checkout@v4

diff --git a/.github/workflows/CI_LuxCore.yml b/.github/workflows/CI_LuxCore.yml
@@ -25,7 +25,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.10"
           - "1"
         os:
           - ubuntu-latest

diff --git a/.github/workflows/CI_LuxLib.yml b/.github/workflows/CI_LuxLib.yml
@@ -41,21 +41,6 @@ jobs:
         loopvec:
           - "true"
         include:
-          - version: "1.10"
-            os: ubuntu-latest
-            test_group: "dense"
-            blas_backend: "blis"
-            loopvec: "true"
-          - version: "1.10"
-            os: ubuntu-latest
-            test_group: "dense"
-            blas_backend: "mkl"
-            loopvec: "true"
-          - version: "1.10"
-            os: macos-latest
-            test_group: "dense"
-            blas_backend: "appleaccelerate"
-            loopvec: "true"
           - version: "1.10"
             os: ubuntu-latest
             test_group: "dense"

diff --git a/.github/workflows/CI_LuxTestUtils.yml b/.github/workflows/CI_LuxTestUtils.yml
@@ -24,7 +24,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.10"
           - "1"
         os:
           - ubuntu-latest

diff --git a/.github/workflows/CI_MLDataDevices.yml b/.github/workflows/CI_MLDataDevices.yml
@@ -24,7 +24,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.10"
           - "1"
         os:
           - ubuntu-latest

diff --git a/.github/workflows/CI_WeightInitializers.yml b/.github/workflows/CI_WeightInitializers.yml
@@ -24,7 +24,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.10"
           - "1"
         os:
           - ubuntu-latest

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

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -60,11 +60,6 @@ if BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda"
     @info sprint(CUDA.versioninfo)
 end
 
-if BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu"
-    using AMDGPU
-    @info sprint(AMDGPU.versioninfo)
-end
-
 using Lux
 
 @testset "Load Tests" begin
-Original file line number
+Diff line change
@@ Expand Up / @@ -23,7 +23,6 @@ steps: @@
             matrix:
               setup:
                 julia:
-                  - "1.10"
                   - "1"
     env:
@@ Expand Down @@