More informative error messages for failing to load devices

LuxDL · Apr 7, 2024 · 3a4f1d9 · 3a4f1d9
1 parent e0089c4
commit 3a4f1d9
Show file tree

Hide file tree

Showing 8 changed files with 17 additions and 7 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -17,7 +17,6 @@ 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:

diff --git a/.gitignore b/.gitignore
@@ -31,3 +31,4 @@ docs/src/tutorials/advanced
 *.log
 
 bench/benchmark_results.json
+*.cov
diff --git a/ext/LuxMPIExt.jl b/ext/LuxMPIExt.jl
@@ -7,26 +7,31 @@ using LuxDeviceUtils: AbstractLuxDevice, LuxCUDADevice, LuxAMDGPUDevice, cpu_dev
 using MPI: MPI
 
 function DistributedUtils.__initialize(
- ::Type{MPIBackend}; cuda_devices=nothing, amdgpu_devices=nothing)
+ ::Type{MPIBackend}; cuda_devices=nothing, amdgpu_devices=nothing,
+ force_cuda::Bool=false, caller::String="", force_amdgpu::Bool=false) # Undocumented internal kwarg 
  !MPI.Initialized() && MPI.Init()
  DistributedUtils.MPI_Initialized[] = true
 
  local_rank = MPI.Comm_rank(MPI.COMM_WORLD)
 
  if cuda_devices !== missing && __is_functional(LuxCUDADevice)
  if cuda_devices === nothing
- set_device!(LuxCUDADevice, nothing, local_rank)
+ set_device!(LuxCUDADevice, nothing, local_rank + 1)
  else
- set_device!(LuxCUDADevice, cuda_devices[local_rank])
+ set_device!(LuxCUDADevice, cuda_devices[local_rank + 1])
  end
+ elseif force_cuda
+ error(lazy"CUDA devices are not functional (or `LuxCUDA.jl` not loaded) and `force_cuda` is set to `true`. This is caused by backend: $(caller).")
  end
 
  if amdgpu_devices !== missing && __is_functional(LuxAMDGPUDevice)
  if amdgpu_devices === nothing
- set_device!(LuxAMDGPUDevice, nothing, local_rank)
+ set_device!(LuxAMDGPUDevice, nothing, local_rank + 1)
  else
- set_device!(LuxAMDGPUDevice, amdgpu_devices[local_rank])
+ set_device!(LuxAMDGPUDevice, amdgpu_devices[local_rank + 1])
  end
+ elseif force_amdgpu
+ error(lazy"AMDGPU devices are not functional (or `LuxAMDGPU.jl` not loaded) and `force_amdgpu` is set to `true`. This is caused by backend: $(caller).")
  end
 
  return

diff --git a/ext/LuxMPINCCLExt.jl b/ext/LuxMPINCCLExt.jl
@@ -9,7 +9,8 @@ using Setfield: @set!
 function DistributedUtils.__initialize(
  ::Type{NCCLBackend}; cuda_devices=nothing, amdgpu_devices=missing)
  @assert amdgpu_devices===missing "`AMDGPU` is not supported by `NCCL`."
- DistributedUtils.__initialize(MPIBackend; cuda_devices, amdgpu_devices)
+ DistributedUtils.__initialize(
+ MPIBackend; cuda_devices, force_cuda=true, caller="NCCLBackend", amdgpu_devices)
  DistributedUtils.NCCL_Initialized[] = true
  return
 end

diff --git a/test/distributed/common_distributedtest.jl b/test/distributed/common_distributedtest.jl
@@ -1,4 +1,5 @@
 using Lux, MPI, NCCL, Test
+using LuxAMDGPU, LuxCUDA
 
 const input_args = length(ARGS) == 2 ? ARGS : ("CPU", "mpi")
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend

diff --git a/test/distributed/data_distributedtest.jl b/test/distributed/data_distributedtest.jl
@@ -1,4 +1,5 @@
 using Lux, MLUtils, MPI, NCCL, Random, Test
+using LuxAMDGPU, LuxCUDA
 
 const input_args = length(ARGS) == 2 ? ARGS : ("CPU", "mpi")
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend

diff --git a/test/distributed/optimizer_distributedtest.jl b/test/distributed/optimizer_distributedtest.jl
@@ -1,4 +1,5 @@
 using Lux, MPI, NCCL, Optimisers, Random, Test
+using LuxAMDGPU, LuxCUDA
 
 const input_args = length(ARGS) == 2 ? ARGS : ("CPU", "mpi")
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend

diff --git a/test/distributed/synchronize_distributedtest.jl b/test/distributed/synchronize_distributedtest.jl
@@ -1,4 +1,5 @@
 using ComponentArrays, Lux, MPI, NCCL, Optimisers, Random, Test
+using LuxAMDGPU, LuxCUDA
 
 const input_args = length(ARGS) == 2 ? ARGS : ("CPU", "mpi")
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend