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

Redistribute#

class nvmath.distributed.distribution.Redistribute(
operand,
/,
input_distribution: Distribution,
output_distribution: Distribution,
*,
options: RedistributeOptions | dict[str, Any] | None = None,
stream: AnyStream | None = None,
)[source]#

Create a stateful object that encapsulates the specified Redistribute operation and required resources. This object ensures the validity of resources during use and releases them when they are no longer needed to prevent misuse.

This object encompasses all functionalities of the function-form API redistribute(), which is a convenience wrapper around it. The stateful object also allows for the amortization of preparatory costs when the same Redistribute operation is to be performed on multiple operands with the same problem specification (see reset_operand() for more details).

Using the stateful object typically involves the following steps:

  1. Problem Specification: Initialize the object with a defined operation and options.

  2. Preparation: Use plan() to determine the best distributed algorithmic implementation for this specific Redistribute operation.

  3. Execution: Perform the redistribution with execute().

  4. Resource Management: Ensure all resources are released either by explicitly calling free() or by managing the stateful object within a context manager.

Detailed information on each step described above can be obtained by passing in a logging.Logger object to RedistributeOptions or by setting the appropriate options in the root logger object, which is used by default:

>>> import logging
>>> logging.basicConfig(
...     level=logging.INFO,
...     format="%(asctime)s %(levelname)-8s %(message)s",
...     datefmt="%m-%d %H:%M:%S",
... )

Changed in version 1.0: This class 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.

  • 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.

See also

plan(), reset_operand(), execute(), release_operand()

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.

Create a Redistribute object encapsulating the problem specification above:

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

Options can be provided above to control the behavior of the operation using the options argument (see RedistributeOptions).

Next, plan the Redistribute:

>>> r.plan()

Now execute the Redistribute, and obtain the result b as a CuPy ndarray. Redistribute always performs the distributed operation on the GPU.

>>> b = r.execute()

Finally, free the Redistribute object’s resources. To avoid this explicit call, it’s recommended to use the Redistribute object as a context manager as shown below, if possible.

>>> r.free()

Any symmetric memory that is owned by the user must be deleted explicitly (this is a collective call and must be called by all processes):

>>> nvmath.distributed.free_symmetric_memory(a, b)

Note that all Redistribute methods execute on the current stream by default. Alternatively, the stream argument can be used to run a method on a specified stream.

Let’s now look at the same problem with NumPy ndarrays on the CPU.

>>> import numpy as np
>>> a = np.random.rand(*shape)  # each process holds a different section

Create a Redistribute object encapsulating the problem specification described earlier and use it as a context manager.

>>> with nvmath.distributed.distribution.Redistribute(a, Slab.X, Slab.Y) as r:
...     r.plan()
...
...     # Execute the Redistribute to redistribute the ndarray.
...     b = r.execute()

All the resources used by the object are released at the end of the block.

Redistribute always executes on the GPU. In this case, because a resides in host memory, the NumPy array is temporarily copied to device memory (on the symmetric memory heap), re-distributed on the GPU, and the result is copied to host memory as a NumPy array.

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

Methods

execute(
stream: AnyStream | None = None,
release_workspace: bool = False,
sync_symmetric_memory: bool = True,
)[source]#

Execute the Redistribute operation.

Parameters:

  • 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.

  • release_workspace – A value of True specifies that the Redistribute object should release workspace memory back to the symmetric memory pool on function return, while a value of False specifies that the object should retain the memory. This option may be set to True if the application performs other operations that consume a lot of memory between successive calls to the (same or different) execute() API, but incurs an overhead due to obtaining and releasing workspace memory from and to the symmetric memory pool on every call. The default is False. NOTE: All processes must use the same value or the application can deadlock.

  • 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.

Returns:

The redistributed operand, which remains on the same device and utilizes the same package as the input operand. For GPU operands, the result will be in symmetric memory and the user is responsible for explicitly deallocating it (for example, using nvmath.distributed.free_symmetric_memory(tensor)).

free()[source]#

Free Redistribute resources.

It is recommended that the Redistribute object be used within a context, but if it’s not possible then this method must be called explicitly to ensure that the Redistribute resources (especially internal library objects) are properly cleaned up.

plan(
stream: AnyStream | None = None,
)[source]#

Plan the Redistribute.

Parameters:

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.

release_operand()[source]#

Added in version 0.9.0.

This method does two things:

  • Releases internal references to the user-provided operand, so that this instance no longer contributes to its reference count.

  • Frees any internal copies (mirrors) that were created when the user-provided operand resides in a different memory space than the execution (i.e., copies made during construction or within reset_operand() / reset_operand_unchecked()).

This functionality can be useful in memory-constrained scenarios, e.g. where multiple stateful objects need to coexist. Leveraging this functionality, the caller can reduce memory usage while retaining the planned state.

Parameters:

None

Returns:

None

Semantics:

  • Preserves the planned state of the stateful object.

  • Repeated calls are safe: if the operand is already released, the call is a no-op and an informational message is logged.

  • After calling this method, reset_operand() (or reset_operand_unchecked()) must be called to supply a new operand before calling execute() again. Failure to do so will result in a runtime error. Device-side copies will be re-allocated as needed.

  • For cross-space scenarios (e.g. CPU operand with GPU execution, or GPU operand with CPU execution): execution is guaranteed to be always blocking. It is therefore always safe to call this method after calling execute() without additional synchronization.

  • When the operand is in the same memory space as the execution (e.g. GPU operand with GPU execution): in such case, this method drops this instance’s internal reference to the user-provided operand. If the reference count of the operand reaches zero, its memory may be freed, so particular attention should be paid. The caller is responsible to ensure that if such deallocation happens, it is ordered after pending computation (e.g. by retaining a reference until the computation is complete, or by synchronizing the stream). Failure to do so is analogous to use-after-free.

See Overview, Stateful APIs: Design and Usage Patterns for operand lifecycle and usage patterns, and Stream Semantics for stream ordering rules.

reset_operand(
operand,
*,
stream: AnyStream | None = None,
)[source]#

Reset the operand held by this Redistribute instance. This method is used to provide a new operand for execution.

Parameters:

  • operand –

    A tensor (ndarray-like object) compatible with the previous one. The new operand is considered compatible if all the following properties match with the previous one:

    • The operand distribution.

    • The package that the new operand belongs to.

    • The dtype of the new operand.

    • The shape and strides of the new operand.

    • The memory space of the new operand (CPU or GPU).

    • The device that new operand belongs to if it is on GPU.

  • 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.

See also

release_operand()

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
>>> nranks = process_group.nranks

Create a 3-D complex128 ndarray on GPU symmetric memory, initially partitioned on the X axis (the global shape is (128, 128, 128)):

>>> shape = 128 // nranks, 128, 128
>>> dtype = cp.complex128
>>> a = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=dtype)
>>> a[:] = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)

Create a Redistribute object as a context manager

>>> from nvmath.distributed.distribution import Redistribute
>>> with Redistribute(a, in_distribution, out_distribution) as r:
...     # Plan the Redistribute
...     r.plan()
...
...     # Execute the Redistribute to get the first result.
...     result1 = r.execute()
...
...     # Reset the operand to a new CuPy ndarray.
...     b = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=dtype)
...     b[:] = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)
...     r.reset_operand(b)
...
...     # Execute to get the new result corresponding to the updated operand.
...     result2 = r.execute()

With reset_operand(), minimal overhead is achieved as problem specification and planning are only performed once.

For the particular example above, explicitly calling reset_operand() is equivalent to updating the operand in-place, i.e, replacing r.reset_operand(b) with a[:]=b. Note that updating the operand in-place should be adopted with caution as it can only yield the expected result and incur no additional copies under the additional constraints below:

  • The operand’s distribution is the same.

For more details, please refer to inplace update example.

reset_operand_unchecked(
operand,
*,
stream: AnyStream | None = None,
)[source]#

This method is a performance-optimized alternative to reset_operand() that eliminates validation and logging overhead, making it ideal for performance-critical loops where operand compatibility is guaranteed by the caller.

This method accepts the same parameters as reset_operand().

Semantics:

The semantics are the same as in reset_operand(), except that this method does not perform any validation (e.g. package match, data type match, shape/strides match, etc.) or logging.

When to Use:

  • Performance-critical loops with repeated executions on different operand

  • After verifying correctness with reset_operand() during development

  • When operand compatibility is guaranteed by construction or invariant

See also

reset_operand(), release_operand()