Add tests for the distributed backends

.buildkite/pipeline.yml

@@ -17,6 +17,7 @@ steps:
           cuda: "*"
         env:
           GROUP: "CUDA"
+          JULIA_MPI_TEST_NPROCS: 2  # Needs to be same as number of GPUs for NCCL
         if: build.message !~ /\[skip tests\]/
         timeout_in_minutes: 240
         matrix:

Project.toml

@@ -120,6 +120,9 @@ LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+NCCL = "3fe64909-d7a1-4096-9b7d-7a0f12cf0f6b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -135,4 +138,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Adapt", "Aqua", "ChainRules", "ChainRulesCore", "ComponentArrays", "ExplicitImports", "Flux", "Functors", "Logging", "LuxAMDGPU", "LuxCUDA", "LuxCore", "LuxLib", "LuxTestUtils", "Optimisers", "Random", "ReTestItems", "Reexport", "ReverseDiff", "Setfield", "SimpleChains", "StableRNGs", "Statistics", "Test", "Tracker", "Zygote"]
+test = ["ADTypes", "Adapt", "Aqua", "ChainRules", "ChainRulesCore", "ComponentArrays", "ExplicitImports", "Flux", "Functors", "Logging", "LuxAMDGPU", "LuxCUDA", "LuxCore", "LuxLib", "LuxTestUtils", "MLUtils", "MPI", "NCCL", "Optimisers", "Pkg", "Random", "ReTestItems", "Reexport", "ReverseDiff", "Setfield", "SimpleChains", "StableRNGs", "Statistics", "Test", "Tracker", "Zygote"]

docs/src/manual/distributed_utils.md

@@ -3,7 +3,7 @@
 !!! tip
 
     For a fully functional example, see the
-    [ImageNet Training Example](https://github.com/LuxDL/Lux.jl/tree/main/examples/ImageNet)
+    [ImageNet Training Example](https://github.com/LuxDL/Lux.jl/tree/main/examples/ImageNet).
 
 DDP Training using `Lux.DistributedUtils` is a spiritual successor to
 [FluxMPI.jl](https://github.com/avik-pal/FluxMPI.jl), but has some key differences.

examples/ImageNet/Project.toml

@@ -5,7 +5,7 @@ Configurations = "5218b696-f38b-4ac9-8b61-a12ec717816d"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
-Formatting = "59287772-0a20-5a39-b81b-1366585eb4c0"
+Format = "1fa38f19-a742-5d3f-a2b9-30dd87b9d5f8"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -30,17 +30,24 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Augmentor = "0.6"
 Boltz = "0.1, 0.2, 0.3"
 Configurations = "0.17"
-Formatting = "0.4"
+FLoops = "0.2"
+FileIO = "1.16"
+Format = "1.3"
 Functors = "0.2, 0.3, 0.4"
+Images = "0.26"
+JLD2 = "0.4.46"
 JpegTurbo = "0.1"
 Lux = "0.4, 0.5"
 LuxAMDGPU = "0.1, 0.2"
 LuxCUDA = "0.2, 0.3"
 MLUtils = "0.2.10, 0.3, 0.4"
+MPI = "0.20.19"
 Metalhead = "0.9"
+NCCL = "0.1.1"
 OneHotArrays = "0.1, 0.2"
 Optimisers = "0.2, 0.3"
-Setfield = "0.8.2, 1"
+ParameterSchedulers = "0.4"
+Setfield = "1"
 SimpleConfig = "0.1"
 Statistics = "1"
 Zygote = "0.6"

examples/ImageNet/main.jl

@@ -6,8 +6,7 @@ import FLoops: ThreadedEx
 import Metalhead
 import MPI, NCCL
 using LuxAMDGPU, LuxCUDA
-
-using Formatting # TODO: Get rid of this
+using Format
 
 # Distributed Training: NCCL for NVIDIA GPUs and MPI for anything else
 if LuxCUDA.functional()
@@ -79,10 +78,12 @@ function construct(cfg::OptimizerConfig)
     end
 
     if cfg.scheduler.name == "cosine"
-        scheduler = CosineAnnealSchedule(cfg.learning_rate, cfg.learning_rate / 100,
-            cfg.scheduler.cycle_length; dampen=cfg.scheduler.damp_factor)
+        l0 = cfg.learning_rate
+        l1 = cfg.learning_rate / 100
+        scheduler = ComposedSchedule(CosAnneal(l0, l1, cfg.scheduler.cycle_length),
+            Step(l0, cfg.scheduler.damp_factor, cfg.scheduler.cycle_length))
     elseif cfg.scheduler.name == "constant"
-        scheduler = ConstantSchedule(cfg.learning_rate)
+        scheduler = Constant(cfg.learning_rate)
     elseif cfg.scheduler.name == "step"
         scheduler = Step(
             cfg.learning_rate, cfg.scheduler.lr_step_decay, cfg.scheduler.lr_step)

examples/ImageNet/utils.jl

@@ -77,51 +77,6 @@ function load_checkpoint(fname::String)
     end
 end
 
-# Parameter Scheduling
-## Copied from ParameterSchedulers.jl due to its heavy dependencies
-struct CosineAnnealSchedule{restart, T, S <: Integer}
-    range::T
-    offset::T
-    dampen::T
-    period::S
-
-    function CosineAnnealSchedule(
-            lambda_0, lambda_1, period; restart::Bool=true, dampen=1.0f0)
-        range = abs(lambda_0 - lambda_1)
-        offset = min(lambda_0, lambda_1)
-        return new{restart, typeof(range), typeof(period)}(range, offset, dampen, period)
-    end
-end
-
-function (s::CosineAnnealSchedule{true})(t)
-    d = s.dampen^div(t - 1, s.period)
-    return (s.range * (1 + cos(pi * mod(t - 1, s.period) / s.period)) / 2 + s.offset) / d
-end
-
-function (s::CosineAnnealSchedule{false})(t)
-    return s.range * (1 + cos(pi * (t - 1) / s.period)) / 2 + s.offset
-end
-
-struct Step{T, S}
-    start::T
-    decay::T
-    step_sizes::S
-
-    function Step(start::T, decay::T, step_sizes::S) where {T, S}
-        _step_sizes = (S <: Integer) ? Iterators.repeated(step_sizes) : step_sizes
-
-        return new{T, typeof(_step_sizes)}(start, decay, _step_sizes)
-    end
-end
-
-(s::Step)(t) = s.start * s.decay^(searchsortedfirst(s.step_sizes, t - 1) - 1)
-
-struct ConstantSchedule{T}
-    val::T
-end
-
-(s::ConstantSchedule)(t) = s.val
-
 # Tracking
 @kwdef mutable struct AverageMeter
     fmtstr
@@ -132,7 +87,7 @@ end
 end
 
 function AverageMeter(name::String, fmt::String)
-    fmtstr = Formatting.FormatExpr("$name {1:$fmt} ({2:$fmt})")
+    fmtstr = FormatExpr("$name {1:$fmt} ({2:$fmt})")
     return AverageMeter(; fmtstr=fmtstr)
 end
 
@@ -160,7 +115,7 @@ function reset_meter!(meter::AverageMeter)
 end
 
 function print_meter(meter::AverageMeter)
-    return Formatting.printfmt(meter.fmtstr, meter.val, meter.average)
+    return printfmt(meter.fmtstr, meter.val, meter.average)
 end
 
 # ProgressMeter
@@ -171,10 +126,9 @@ end
 
 function ProgressMeter(num_batches::Int, meters::NTuple{N}, prefix::String="") where {N}
     fmt = "%" * string(length(string(num_batches))) * "d"
+    fmt2 = "{1:" * string(length(string(num_batches))) * "d}"
     prefix = prefix != "" ? endswith(prefix, " ") ? prefix : prefix * " " : ""
-    batch_fmtstr = Formatting.generate_formatter("$prefix[$fmt/" *
-                                                 Formatting.sprintf1(fmt, num_batches) *
-                                                 "]")
+    batch_fmtstr = FormatExpr("$prefix[$fmt2/" * cfmt(fmt, num_batches) * "]")
     return ProgressMeter{N}(batch_fmtstr, meters)
 end
 
@@ -184,8 +138,7 @@ function reset_meter!(meter::ProgressMeter)
 end
 
 function print_meter(meter::ProgressMeter, batch::Int)
-    base_str = meter.batch_fmtstr(batch)
-    print(base_str)
+    printfmt(meter.batch_fmtstr, batch)
     foreach(x -> (print("\t"); print_meter(x)), meter.meters[1:end])
     println()
     return nothing

ext/LuxComponentArraysExt.jl

@@ -14,8 +14,8 @@ Lux.__named_tuple(ca::ComponentArray) = NamedTuple(ca)
 # Distributed Functionality
 function DistributedUtils.synchronize!!(
         backend::Lux.AbstractLuxDistributedBackend, ps::ComponentArray; root::Int=0)
-    ps_synced = DistributedUtils.synchronize!!(backend, getdata(ps); root)
-    return ComponentArray(ps_synced, getaxes(ps))
+    ps_synced = DistributedUtils.synchronize!!(backend, ComponentArrays.getdata(ps); root)
+    return ComponentArray(ps_synced, ComponentArrays.getaxes(ps))
 end
 
 end

ext/LuxMPIExt.jl

@@ -39,7 +39,6 @@ DistributedUtils.local_rank(backend::MPIBackend) = MPI.Comm_rank(backend.comm)
 DistributedUtils.total_workers(backend::MPIBackend) = MPI.Comm_size(backend.comm)
 
 # Broadcast
-
 function DistributedUtils.__bcast!(
         backend::MPIBackend, sendrecvbuf, dev::AbstractLuxDevice; root=0)
     MPI.Bcast!(sendrecvbuf, backend.comm; root)
@@ -78,20 +77,21 @@ for (aware, dType) in ((MPI_CUDA_AWARE, LuxCUDADevice), (MPI_ROCM_AWARE, LuxAMDG
 end
 
 # Allreduce
-
 function DistributedUtils.__allreduce!(
         backend::MPIBackend, sendrecvbuf, op::F, dev::AbstractLuxDevice) where {F}
-    MPI.Allreduce!(sendrecvbuf, op, backend.comm)
-    if op === typeof(DistributedUtils.avg)
+    mpiop = ifelse(op === DistributedUtils.avg, +, op)
+    MPI.Allreduce!(sendrecvbuf, mpiop, backend.comm)
+    if op === DistributedUtils.avg
         sendrecvbuf ./= DistributedUtils.total_workers(backend)
     end
     return sendrecvbuf
 end
 
 function DistributedUtils.__allreduce!(
         backend::MPIBackend, sendbuf, recvbuf, op::F, dev::AbstractLuxDevice) where {F}
-    MPI.Allreduce!(sendbuf, recvbuf, op, backend.comm)
-    if op === typeof(DistributedUtils.avg)
+    mpiop = ifelse(op === DistributedUtils.avg, +, op)
+    MPI.Allreduce!(sendbuf, recvbuf, mpiop, backend.comm)
+    if op === DistributedUtils.avg
         recvbuf ./= DistributedUtils.total_workers(backend)
     end
     return recvbuf
@@ -123,20 +123,21 @@ for (aware, dType) in ((MPI_CUDA_AWARE, LuxCUDADevice), (MPI_ROCM_AWARE, LuxAMDG
 end
 
 # Reduce
-
 function DistributedUtils.__reduce!(backend::MPIBackend, sendrecvbuf, op::F,
         dev::AbstractLuxDevice; root::Int) where {F}
-    MPI.Reduce!(sendrecvbuf, op, backend.comm; root)
-    if op === typeof(DistributedUtils.avg)
+    mpiop = ifelse(op === DistributedUtils.avg, +, op)
+    MPI.Reduce!(sendrecvbuf, mpiop, backend.comm; root)
+    if op === DistributedUtils.avg
         sendrecvbuf ./= DistributedUtils.total_workers(backend)
     end
     return sendrecvbuf
 end
 
 function DistributedUtils.__reduce!(backend::MPIBackend, sendbuf, recvbuf, op::F,
         dev::AbstractLuxDevice; root::Int) where {F}
-    MPI.Reduce!(sendbuf, recvbuf, op, backend.comm; root)
-    if op === typeof(DistributedUtils.avg)
+    mpiop = ifelse(op === DistributedUtils.avg, +, op)
+    MPI.Reduce!(sendbuf, recvbuf, mpiop, backend.comm; root)
+    if op === DistributedUtils.avg
         recvbuf ./= DistributedUtils.total_workers(backend)
     end
     return recvbuf

ext/LuxMPINCCLExt.jl

@@ -33,7 +33,6 @@ DistributedUtils.total_workers(backend::NCCLBackend) = NCCL.size(backend.comm)
 
 # For non-CUDA Arrays, fallback to MPI
 # Broadcast
-
 function DistributedUtils.__bcast!(
         backend::NCCLBackend, sendrecvbuf, ::LuxCUDADevice; root=0)
     NCCL.Broadcast!(sendrecvbuf, backend.comm; root)
@@ -57,7 +56,6 @@ function DistributedUtils.__bcast!(
 end
 
 # Allreduce
-
 function DistributedUtils.__allreduce!(
         backend::NCCLBackend, sendrecvbuf, op::F, ::LuxCUDADevice) where {F}
     op = ifelse(op === DistributedUtils.avg, NCCL.avg, op)
@@ -83,7 +81,6 @@ function DistributedUtils.__allreduce!(
 end
 
 # Reduce
-
 function DistributedUtils.__reduce!(
         backend::NCCLBackend, sendrecvbuf, op::F, ::LuxCUDADevice; root::Int) where {F}
     op = ifelse(op === DistributedUtils.avg, NCCL.avg, op)

ext/LuxOptimisersExt.jl

@@ -49,7 +49,7 @@ end
 end
 
 function Optimisers.apply!(opt::DistributedOptimizer, state, x, y)
-    y_avg = allreduce!(opt.backend, y, DistributedUtils.avg)
+    y_avg = DistributedUtils.allreduce!(opt.backend, y, DistributedUtils.avg)
     return Optimisers.apply!(opt.opt, state, x, y_avg)
 end
 

test/distributed/common_distributedtest.jl

@@ -0,0 +1,15 @@
+using Lux, MPI, NCCL, Test
+
+const input_args = length(ARGS) == 3 ? ARGS[1:3] :
+                   (length(ARGS) == 2 ? (ARGS[2], "mpi") : ("CPU", "mpi"))
+const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+const dev = input_args[1] == "CPU" ? LuxCPUDevice() :
+            (input_args[1] == "CUDA" ? LuxCUDADevice() : LuxAMDGPUDevice())
+
+DistributedUtils.initialize(backend_type)
+backend = DistributedUtils.get_distributed_backend(backend_type)
+
+@test DistributedUtils.initialized(backend_type)
+
+# Should always hold true
+@test DistributedUtils.local_rank(backend) < DistributedUtils.total_workers(backend)

test/distributed/data_distributedtest.jl

@@ -0,0 +1,28 @@
+using Lux, MLUtils, MPI, NCCL, Random, Test
+
+const input_args = length(ARGS) == 3 ? ARGS[1:3] :
+                   (length(ARGS) == 2 ? (ARGS[2], "mpi") : ("CPU", "mpi"))
+const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+const dev = input_args[1] == "CPU" ? LuxCPUDevice() :
+            (input_args[1] == "CUDA" ? LuxCUDADevice() : LuxAMDGPUDevice())
+
+rng = Xoshiro(1234)
+
+DistributedUtils.initialize(backend_type)
+backend = DistributedUtils.get_distributed_backend(backend_type)
+
+data = randn(rng, Float32, 10)
+dcontainer = DistributedUtils.DistributedDataContainer(backend, data)
+
+rank = DistributedUtils.local_rank(backend)
+tworkers = DistributedUtils.total_workers(backend)
+
+if rank != tworkers - 1
+    @test length(dcontainer) == ceil(length(data) / tworkers)
+else
+    @test length(dcontainer) ==
+          length(data) - (tworkers - 1) * ceil(length(data) / tworkers)
+end
+
+dsum = sum(Base.Fix1(MLUtils.getobs, dcontainer), 1:MLUtils.numobs(dcontainer))
+@test DistributedUtils.allreduce!(backend, [dsum], +)[1] ≈ sum(data)

test/distributed/optimizer_distributedtest.jl

@@ -0,0 +1,30 @@
+using Lux, MPI, NCCL, Optimisers, Random, Test
+
+const input_args = length(ARGS) == 3 ? ARGS[1:3] :
+                   (length(ARGS) == 2 ? (ARGS[2], "mpi") : ("CPU", "mpi"))
+const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+const dev = input_args[1] == "CPU" ? LuxCPUDevice() :
+            (input_args[1] == "CUDA" ? LuxCUDADevice() : LuxAMDGPUDevice())
+
+DistributedUtils.initialize(backend_type)
+backend = DistributedUtils.get_distributed_backend(backend_type)
+
+opt = Adam(0.001f0)
+ps = (a=zeros(4), b=zeros(4)) |> dev
+st_opt = Optimisers.setup(opt, ps)
+
+dopt = DistributedUtils.DistributedOptimizer(backend, opt)
+st_dopt = Optimisers.setup(dopt, ps)
+
+@test st_dopt.a.state == st_opt.a.state
+@test st_dopt.b.state == st_opt.b.state
+
+@test_nowarn DistributedUtils.synchronize!!(backend, st_dopt)
+
+gs = (a=ones(4), b=ones(4)) |> dev
+
+_, ps_dopt = Optimisers.update(st_dopt, ps, gs)
+_, ps_opt = Optimisers.update(st_opt, ps, gs)
+
+@test ps_dopt.a≈ps_opt.a atol=1.0e-5 rtol=1.0e-5
+@test ps_dopt.b≈ps_opt.b atol=1.0e-5 rtol=1.0e-5