saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use local rank as GPU ID (#6719)

Browse source
This commit is contained in:
Wei-Sheng Chin 2021-02-17 22:42:54 +08:00 committed by GitHub
parent 9043df8b66
commit 9e67b88c83
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc
View file

@ -229,21 +229,12 @@ void NcclService::Initialize() {
// CPUs
// Other devices
// Make sure MPI is initialized.
int is_mpi_initialized = 0;
MPI_CHECK(MPI_Initialized(&is_mpi_initialized));
if (!is_mpi_initialized) {
int mpi_threads_provided = 0;
MPI_CHECK(MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &mpi_threads_provided));
}
int mpi_rank;
int mpi_size;
MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank));
MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &mpi_size));
const int mpi_rank = onnxruntime::training::MPIContext::GetInstance().GetWorldRank();
const int mpi_local_rank = onnxruntime::training::MPIContext::GetInstance().GetLocalRank();
const int mpi_size = onnxruntime::training::MPIContext::GetInstance().GetWorldSize();
// Set device this NCCL communicator runs on.
CUDA_CALL(cudaSetDevice(mpi_rank));
CUDA_CALL(cudaSetDevice(mpi_local_rank));
// Create communication stream.
CUDA_CALL(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));