Fix missing imports in FluxMPIExt (#2589)

* missing import of cpu_device

* test AMD GPU availability with functional(AMDGPUDevice)

* add missing import for NCCL

* test reduction for MPI and NCCL

Author: Alexander-Barth

ext/FluxMPIExt/FluxMPIExt.jl

@@ -3,7 +3,7 @@ module FluxMPIExt
 using Flux: MPIBackend, NCCLBackend, DistributedUtils,
             MPI_CUDA_AWARE, MPI_ROCM_AWARE
 using MPI: MPI
-using MLDataDevices: AbstractDevice, CUDADevice, AMDGPUDevice, functional, set_device!
+using MLDataDevices: AbstractDevice, CUDADevice, AMDGPUDevice, functional, set_device!, cpu_device
 
 
 function DistributedUtils.__initialize(
@@ -24,7 +24,7 @@ function DistributedUtils.__initialize(
         error(lazy"CUDA devices are not functional and `force_cuda` is set to `true`. This is caused by backend: $(caller).")
     end
 
-    if amdgpu_devices !== missing && AMDGPU.functional()
+    if amdgpu_devices !== missing && functional(AMDGPUDevice)
         if amdgpu_devices === nothing
             set_device!(AMDGPUDevice, nothing, local_rank + 1)
         else

test/ext_distributed/reduce_distributedtest.jl

@@ -0,0 +1,28 @@
+using Flux
+using Test
+
+backend_string = ARGS[1]
+
+if backend_string == "mpi"
+    import MPI
+    backend_type = MPIBackend
+elseif backend_string == "nccl"
+    import MPI, NCCL, CUDA
+    backend_type = NCCLBackend
+else
+    error("unsupported backend: $backend_string")
+end
+
+DistributedUtils.initialize(backend_type)
+backend = DistributedUtils.get_distributed_backend(backend_type)
+
+rank = DistributedUtils.local_rank(backend)
+total_workers = DistributedUtils.total_workers(backend)
+
+sendrecvbuf = fill(rank+1,4)
+
+DistributedUtils.reduce!(backend, sendrecvbuf, +)
+
+if rank == 0
+    @test all(sendrecvbuf .== sum(1:total_workers))
+end

test/ext_distributed/runtests.jl

@@ -1,5 +1,8 @@
 # Distributed Tests
 using MPI, Pkg, Test
+if get(ENV, "FLUX_TEST_DISTRIBUTED_NCCL", "false") == "true"
+    import CUDA
+end
 
 nprocs_str = get(ENV, "JULIA_MPI_TEST_NPROCS", "")
 nprocs = nprocs_str == "" ? clamp(Sys.CPU_THREADS, 2, 4) : parse(Int, nprocs_str)