in src/nccl_op.cc [332:383]
void CreateCustomComm(int op_num, const std::set<int>& mpi_ranks, int comm_id, StreamExecutor* executor)
{
auto ranks_and_id = std::make_pair(mpi_ranks, comm_id);
if (op_to_comm.count(op_num))
{
return;
}
else if (ranks_to_comm.count(ranks_and_id))
{
op_to_comm[op_num] = ranks_to_comm[ranks_and_id];
return;
}
ncclUniqueId id;
if (mpi_rank == 0)
{
NCCL_CHECK(ncclGetUniqueId(&id), "ncclGetUniqueId");
}
MPI_CHECK(MPI_Bcast((void*)(&id), sizeof(ncclUniqueId), MPI_BYTE, 0, mpi_comm), "MPI_Bcast");
int custom_mpi_size = static_cast<int>(mpi_ranks.size());
int custom_mpi_rank = -1;
if (mpi_ranks.count(mpi_rank))
{
custom_mpi_rank = 0;
for (auto r : mpi_ranks)
{
if (mpi_rank == r)
{
break;
}
custom_mpi_rank++;
}
}
if (custom_mpi_rank != -1)
{
custom_comms.push_back(new NcclComm(custom_mpi_size, custom_mpi_rank, -1, -1, id, id, 0, executor));
}
else
{
custom_comms.push_back(new NcclComm(-1, -1, -1, -1, id, id, 0, executor));
}
int custom_comm_idx = static_cast<int>(custom_comms.size()) - 1;
ranks_to_comm[ranks_and_id] = custom_comm_idx;
op_to_comm[op_num] = custom_comm_idx;
}