helpers.NCCL_communicator module¶

class NCCLComm(comm)[source]¶

Bases: object

Wraps an MPI communicator and performs some calls to NCCL functions instead.

Allreduce(sendbuf, recvbuf, op=mpi4py.MPI.SUM)[source]¶

Reduce(sendbuf, recvbuf, op=mpi4py.MPI.SUM, root=0)[source]¶

static get_count(data)[source]¶: As NCCL doesn’t support complex numbers, we have to act as if we’re sending two real numbers if using complex.

static get_dtype(data)[source]¶: As NCCL doesn’t support complex numbers, we have to act as if we’re sending two real numbers if using complex.

get_op(MPI_op)[source]¶