Handle NCCL.avg correctly (#54)

JuliaGPU · Feb 28, 2024 · 0ee0774 · 0ee0774 · avik-pal · Mar 2, 2024
1 parent 3a49978
commit 0ee0774
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 17 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "NCCL"
 uuid = "3fe64909-d7a1-4096-9b7d-7a0f12cf0f6b"
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"

diff --git a/src/base.jl b/src/base.jl
@@ -33,6 +33,8 @@ ncclRedOp_t(::typeof(+)) = ncclSum
 ncclRedOp_t(::typeof(*)) = ncclProd
 ncclRedOp_t(::typeof(max)) = ncclMax
 ncclRedOp_t(::typeof(min)) = ncclMin
+# Handles the case where user directly passed in the ncclRedOp_t (eg. `NCCL.avg`)
+ncclRedOp_t(x::ncclRedOp_t) = x
 
 """
     NCCl.avg

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -25,24 +25,49 @@ end
 @testset "Allreduce!" begin
     devs  = CUDA.devices()
     comms = NCCL.Communicators(devs)
-    recvbuf = Vector{CuVector{Float64}}(undef, length(devs))
-    sendbuf = Vector{CuVector{Float64}}(undef, length(devs))
-    N = 512
-    for (ii, dev) in enumerate(devs)
-        CUDA.device!(ii - 1)
-        sendbuf[ii] = CuArray(fill(Float64(ii), N))
-        recvbuf[ii] = CUDA.zeros(Float64, N)
-    end
-    NCCL.group() do
-        for ii in 1:length(devs)
-            NCCL.Allreduce!(sendbuf[ii], recvbuf[ii], +, comms[ii])
+
+    @testset "sum" begin
+        recvbuf = Vector{CuVector{Float64}}(undef, length(devs))
+        sendbuf = Vector{CuVector{Float64}}(undef, length(devs))
+        N = 512
+        for (ii, dev) in enumerate(devs)
+            CUDA.device!(ii - 1)
+            sendbuf[ii] = CuArray(fill(Float64(ii), N))
+            recvbuf[ii] = CUDA.zeros(Float64, N)
+        end
+        NCCL.group() do
+            for ii in 1:length(devs)
+                NCCL.Allreduce!(sendbuf[ii], recvbuf[ii], +, comms[ii])
+            end
+        end
+        answer = sum(1:length(devs))
+        for (ii, dev) in enumerate(devs)
+            device!(ii - 1)
+            crecv = collect(recvbuf[ii])
+            @test all(crecv .== answer)
         end
     end
-    answer = sum(1:length(devs))
-    for (ii, dev) in enumerate(devs)
-        device!(ii - 1)
-        crecv = collect(recvbuf[ii])
-        @test all(crecv .== answer)
+
+    @testset "NCCL.avg" begin
+        recvbuf = Vector{CuVector{Float64}}(undef, length(devs))
+        sendbuf = Vector{CuVector{Float64}}(undef, length(devs))
+        N = 512
+        for (ii, dev) in enumerate(devs)
+            CUDA.device!(ii - 1)
+            sendbuf[ii] = CuArray(fill(Float64(ii), N))
+            recvbuf[ii] = CUDA.zeros(Float64, N)
+        end
+        NCCL.group() do
+            for ii in 1:length(devs)
+                NCCL.Allreduce!(sendbuf[ii], recvbuf[ii], NCCL.avg, comms[ii])
+            end
+        end
+        answer = sum(1:length(devs)) / length(devs)
+        for (ii, dev) in enumerate(devs)
+            device!(ii - 1)
+            crecv = collect(recvbuf[ii])
+            @test all(crecv .≈ answer)
+        end
     end
 end