ProcessGroup#

class nvmath.distributed.ProcessGroup[source]#

A ProcessGroup represents a set of processes collectively running nvmath.distributed operations. A ProcessGroup implements a small set of communication collectives used by nvmath-python to initialize the distributed runtime and set up distributed math operations. It is not used for compute (each math library uses its chosen communication backend(s), like NCCL or NVSHMEM).

Attributes

nranks#

Number of processes in the group.

rank#

Index of this process in the group.

Methods

abstractmethod allreduce_buffer(
array: ndarray,
*,
op: ReductionOp,
) → None[source]#

Allreduce an array.

Parameters:

  • array – Input and output of the collective. The function operates in-place.

  • op – One of the values from ReduceOp enum. Specifies an operation for element-wise reductions.

abstractmethod allreduce_object(
obj: T,
*,
op: Callable[[T, T], T],
) → T[source]#

Reduces all Python objects contributed by members of the group. The result is a single reduced object which is returned on every process.

Parameters:

  • obj – object contributed by this process.

  • op – A Python function that takes two objects and returns a single (reduced) object.

abstractmethod broadcast_buffer(
array: ndarray,
*,
root: int = 0,
) → None[source]#

Broadcast an array from one process to every process.

Parameters:

  • array – input (on root) and output of the collective.

  • root – rank of sending process.