From 876416ebd447e5a8da9edba5a622ea6206d235fe Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 4 Aug 2024 14:40:55 -0700 Subject: [PATCH] refactor: update to latest versions --- .github/workflows/CI.yml | 5 ----- Project.toml | 40 +++++++---------------------------- src/basis.jl | 4 ++-- src/layers/Layers.jl | 2 +- src/layers/tensor_product.jl | 4 ++-- test/Project.toml | 41 ++++++++++++++++++++++++++++++++++++ test/layer_tests.jl | 5 +++-- test/runtests.jl | 22 +++++++++++++------ test/shared_testsetup.jl | 11 +++++----- 9 files changed, 79 insertions(+), 55 deletions(-) create mode 100644 test/Project.toml diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index b7df668..a7a0714 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -77,9 +77,6 @@ jobs: - uses: julia-actions/julia-downgrade-compat@v1 - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - env: - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 - uses: julia-actions/julia-processcoverage@v1 with: directories: src,ext @@ -121,5 +118,3 @@ jobs: env: BACKEND_GROUP: "CPU" - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 diff --git a/Project.toml b/Project.toml index 555af74..b8a3a50 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Boltz" uuid = "4544d5e4-abc5-4dea-817f-29e4c205d9c8" authors = ["Avik Pal and contributors"] -version = "0.3.10" +version = "0.3.11" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -15,7 +15,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" -LuxDeviceUtils = "34f89e08-e1d5-43b4-8944-0b49ac560553" +MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -34,49 +34,25 @@ BoltzMetalheadExt = "Metalhead" BoltzZygoteExt = "Zygote" [compat] -ADTypes = "1.3" -Aqua = "0.8.7" +ADTypes = "1.5" ArgCheck = "2.3" -Artifacts = "1.10" +Artifacts = "1.10, 1" ChainRulesCore = "1.24" -ComponentArrays = "0.15.13" ConcreteStructs = "0.2.3" DataInterpolations = "< 5.3" -ExplicitImports = "1.9" ForwardDiff = "0.10.36" GPUArraysCore = "0.1.6" JLD2 = "0.4.48" LazyArtifacts = "1.10" -Lux = "0.5.60" -LuxCore = "0.1.16" -LuxDeviceUtils = "0.1.26" -LuxLib = "0.3.26" -LuxTestUtils = "0.1.18" +Lux = "0.5.62" +LuxCore = "0.1.21" +MLDataDevices = "1.0.0" Markdown = "1.10" Metalhead = "0.9" -NNlib = "0.9.17" -Pkg = "1.10" +NNlib = "0.9.21" Random = "1.10" -ReTestItems = "1.24.0" Reexport = "1.2.2" Statistics = "1.10" -Test = "1.10" WeightInitializers = "0.1.7, 1" Zygote = "0.6.70" julia = "1.10" - -[extras] -Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" -DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" -ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" -LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" -LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531" -Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" -Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[targets] -test = ["Aqua", "ComponentArrays", "DataInterpolations", "ExplicitImports", "LuxLib", "LuxTestUtils", "Metalhead", "Pkg", "ReTestItems", "Test", "Zygote"] diff --git a/src/basis.jl b/src/basis.jl index 1d71e1e..f6d8e83 100644 --- a/src/basis.jl +++ b/src/basis.jl @@ -4,8 +4,8 @@ using ArgCheck: @argcheck using ..Boltz: _unsqueeze1 using ChainRulesCore: ChainRulesCore, NoTangent using ConcreteStructs: @concrete -using LuxDeviceUtils: get_device, LuxCPUDevice using Markdown: @doc_str +using MLDataDevices: get_device, CPUDevice const CRC = ChainRulesCore @@ -34,7 +34,7 @@ end x_new = reshape(x, new_x_size) if grid isa AbstractRange dev = get_device(x) - grid = dev isa LuxCPUDevice ? collect(grid) : dev(grid) + grid = dev isa CPUDevice ? collect(grid) : dev(grid) end grid_shape = ntuple(i -> i == basis.dim ? basis.n : 1, ndims(x) + 1) grid_new = reshape(grid, grid_shape) diff --git a/src/layers/Layers.jl b/src/layers/Layers.jl index 8c7b003..7bbe33b 100644 --- a/src/layers/Layers.jl +++ b/src/layers/Layers.jl @@ -8,8 +8,8 @@ using ChainRulesCore: ChainRulesCore using ForwardDiff: ForwardDiff using Lux: Lux, StatefulLuxLayer using LuxCore: LuxCore, AbstractExplicitLayer, AbstractExplicitContainerLayer -using LuxDeviceUtils: get_device, LuxCPUDevice, LuxCUDADevice using Markdown: @doc_str +using MLDataDevices: get_device_type, CPUDevice, CUDADevice using NNlib: NNlib using Random: AbstractRNG using WeightInitializers: zeros32, randn32 diff --git a/src/layers/tensor_product.jl b/src/layers/tensor_product.jl index a320e04..7346bf6 100644 --- a/src/layers/tensor_product.jl +++ b/src/layers/tensor_product.jl @@ -24,8 +24,8 @@ function TensorProductLayer(basis_fns, out_dim::Int; init_weight::F=randn32) whe prod(Base.Fix2(getproperty, :n), basis_fns) => out_dim; use_bias=false, init_weight) return Lux.@compact(; basis_fns=Tuple(basis_fns), dense, out_dim, dispatch=:TensorProductLayer) do x::AbstractArray # I1 x I2 x ... x T x B - dev = get_device(x) - @argcheck dev isa LuxCPUDevice || dev isa LuxCUDADevice # kron is not widely supported + dev = get_device_type(x) + @argcheck dev <: CPUDevice || dev <: CUDADevice # kron is not widely supported x_ = Lux._eachslice(x, Val(ndims(x) - 1)) # [I1 x I2 x ... x B] x T @argcheck length(x_) == length(basis_fns) diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 0000000..0a92649 --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,41 @@ +[deps] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" +DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" +LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531" +MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" +Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +Aqua = "0.8.7" +ComponentArrays = "0.15.16" +DataInterpolations = "< 5.3" +ForwardDiff = "0.10.36" +ExplicitImports = "1.9.0" +GPUArraysCore = "0.1.6" +Hwloc = "3.2.0" +InteractiveUtils = "<0.0.1, 1" +Lux = "0.5.62" +LuxLib = "0.3.39" +LuxTestUtils = "1.1.2" +MLDataDevices = "1.0.0" +Metalhead = "0.9" +Random = "1.10" +Reexport = "1.2.2" +Pkg = "1.10" +ReTestItems = "1.24.0" +Test = "1.10" +Zygote = "0.6.70" diff --git a/test/layer_tests.jl b/test/layer_tests.jl index 9f361d4..e12c21d 100644 --- a/test/layer_tests.jl +++ b/test/layer_tests.jl @@ -14,7 +14,7 @@ @jet model(x, ps, st) __f = (x, ps) -> sum(abs2, first(model(x, ps, st))) - @eval @test_gradients $(__f) $x $ps gpu_testing=$(ongpu) atol=1e-3 rtol=1e-3 + test_gradients(__f, x, ps; atol=1e-3, rtol=1e-3, soft_fail=[AutoFiniteDiff()]) end end end @@ -104,7 +104,8 @@ end # @jet tensor_project(x, ps, st) __f = (x, ps) -> sum(abs2, first(tensor_project(x, ps, st))) - @eval @test_gradients $(__f) $x $ps gpu_testing=$(ongpu) atol=1e-3 rtol=1e-3 skip_tracker=true + test_gradients(__f, x, ps; atol=1e-3, rtol=1e-3, + skip_backends=[AutoTracker(), AutoEnzyme()]) end end end diff --git a/test/runtests.jl b/test/runtests.jl index 3fd7ef2..c7d4513 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,6 @@ -using ReTestItems, Pkg +using ReTestItems, Pkg, InteractiveUtils, Hwloc + +@info sprint(io -> versioninfo(io; verbose=true)) const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "all")) const EXTRA_PKGS = String[] @@ -14,10 +16,18 @@ if !isempty(EXTRA_PKGS) Pkg.instantiate() end -const BOLTZ_TEST_GROUP = lowercase(get(ENV, "BOLTZ_TEST_GROUP", "all")) -@info "Running tests for group: $BOLTZ_TEST_GROUP" -const RETESTITEMS_NWORKERS = parse(Int, get(ENV, "RETESTITEMS_NWORKERS", "0")) +using Boltz + +const BOLTZ_TEST_GROUP = get(ENV, "BOLTZ_TEST_GROUP", "all") +const RETESTITEMS_NWORKERS = parse( + Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 16)))) +const RETESTITEMS_NWORKER_THREADS = parse(Int, + get(ENV, "RETESTITEMS_NWORKER_THREADS", + string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1)))) + +@info "Running tests for group: $BOLTZ_TEST_GROUP with $RETESTITEMS_NWORKERS workers" ReTestItems.runtests( - @__DIR__; tags=(BOLTZ_TEST_GROUP == "all" ? nothing : [Symbol(BOLTZ_TEST_GROUP)]), - nworkers=ifelse(BACKEND_GROUP ∈ ("cuda", "amdgpu"), 0, RETESTITEMS_NWORKERS)) + Boltz; tags=(BOLTZ_TEST_GROUP == "all" ? nothing : [Symbol(BOLTZ_TEST_GROUP)]), + nworkers=ifelse(BACKEND_GROUP ∈ ("cuda", "amdgpu"), 0, RETESTITEMS_NWORKERS), + nworker_threads=RETESTITEMS_NWORKER_THREADS, testitem_timeout=3600) diff --git a/test/shared_testsetup.jl b/test/shared_testsetup.jl index 7aaeea2..717dcbe 100644 --- a/test/shared_testsetup.jl +++ b/test/shared_testsetup.jl @@ -3,6 +3,7 @@ import Reexport: @reexport @reexport using Boltz, Lux, GPUArraysCore, LuxLib, LuxTestUtils, Random import Metalhead +using MLDataDevices LuxTestUtils.jet_target_modules!(["Boltz", "Lux", "LuxLib"]) @@ -21,18 +22,18 @@ end cpu_testing() = BACKEND_GROUP == "all" || BACKEND_GROUP == "cpu" function cuda_testing() return (BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda") && - LuxDeviceUtils.functional(LuxCUDADevice) + MLDataDevices.functional(CUDADevice) end function amdgpu_testing() return (BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu") && - LuxDeviceUtils.functional(LuxAMDGPUDevice) + MLDataDevices.functional(AMDGPUDevice) end const MODES = begin modes = [] - cpu_testing() && push!(modes, ("cpu", Array, LuxCPUDevice(), false)) - cuda_testing() && push!(modes, ("cuda", CuArray, LuxCUDADevice(), true)) - amdgpu_testing() && push!(modes, ("amdgpu", ROCArray, LuxAMDGPUDevice(), true)) + cpu_testing() && push!(modes, ("cpu", Array, CPUDevice(), false)) + cuda_testing() && push!(modes, ("cuda", CuArray, CUDADevice(), true)) + amdgpu_testing() && push!(modes, ("amdgpu", ROCArray, AMDGPUDevice(), true)) modes end