The nvmath.distributed.distribution module is experimental and potentially subject to future changes.

redistribute#

nvmath.distributed.distribution.redistribute(
operand,
input_distribution: Distribution,
output_distribution: Distribution,
sync_symmetric_memory: bool = True,
options: RedistributeOptions | dict[str, Any] | None = None,
stream: AnyStream | None = None,
)[source]#

Perform a redistribute operation on the provided operand to change its distribution across processes.

Changed in version 1.0: This function was previously nvmath.distributed.reshape.reshape.

Parameters:

  • operand –

    A tensor (ndarray-like object). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.

    Important

    GPU operands must be on the symmetric heap (for example, allocated with nvmath.distributed.allocate_symmetric_memory()).

  • input_distribution – The distribution of the input operand across processes, determining which portion of the global array each process holds.

  • output_distribution – The desired distribution of the result across processes, where each process specifies which portion of the global array it will hold after reshaping.

  • sync_symmetric_memory – Indicates whether to issue a symmetric memory synchronization operation on the execute stream before the redistribute operation. Note that before Redistribute starts executing, it is required that the source operand be ready on all processes. A symmetric memory synchronization ensures completion and visibility by all processes of previously issued local stores to symmetric memory. Advanced users who choose to manage the synchronization on their own using the appropriate NVSHMEM API, or who know that GPUs are already synchronized on the source operand, can set this to False.

  • options – Specify options for Redistribute as a RedistributeOptions object. Alternatively, a dict containing the parameters for the RedistributeOptions constructor can also be provided. If not specified, the value will be set to the default-constructed RedistributeOptions object.

  • stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Returns:

A tensor that remains on the same device and belongs to the same package as the input operand, with shape according to output_box.

See also

Redistribute.

Examples

>>> import cupy as cp
>>> import nvmath.distributed

Get process group used to initialize nvmath.distributed (for information on initializing nvmath.distributed, you can refer to the documentation or to the Redistribute examples in nvmath/examples/distributed/redistribute):

>>> process_group = nvmath.distributed.get_context().process_group
>>> rank = process_group.rank

Let’s create a 3D floating-point ndarray on GPU, distributed across a certain number of processes, according to the Slab distribution with partitioning on the X axis.

>>> shape = Slab.X.shape(rank, (global_shape))

Redistribute uses the NVSHMEM PGAS model, which requires GPU operands to be on the symmetric heap:

>>> a = nvmath.distributed.allocate_symmetric_memory(shape, cp)
>>> a[:] = cp.random.rand(*shape)

With Redistribute, we will change how the ndarray is distributed, by specifying the input distribution and the desired output distribution. In this example, we redistribute from Slab.X to Slab.Y.

Perform the redistribute operation using redistribute():

>>> r = nvmath.distributed.distribution.redistribute(a, Slab.X, Slab.Y)

See RedistributeOptions for the complete list of available options.

The package current stream is used by default, but a stream can be explicitly provided to the Redistribute operation. This can be done if the Redistribute operand is computed on a different stream, for example:

>>> s = cp.cuda.Stream()
>>> with s:
...     a = nvmath.distributed.allocate_symmetric_memory(shape, cp)
...     a[:] = cp.random.rand(*shape)
>>> r = nvmath.distributed.distribution.redistribute(a, Slab.X, Slab.Y, stream=s)

The operation above runs on stream s and is ordered with respect to the input computation.

Create a NumPy ndarray on the CPU.

>>> import numpy as np
>>> b = np.random.rand(*shape)

Provide the NumPy ndarray to redistribute(), with the result also being a NumPy ndarray:

>>> r = nvmath.distributed.distribution.redistribute(b, Slab.X, Slab.Y)

Notes

  • This function is a convenience wrapper around Redistribute and is specifically meant for single use. The same computation can be performed with the stateful API.

Further examples can be found in the nvmath/examples/distributed/redistribute directory.

On this page