allocate_symmetric_memory#

nvmath.distributed.allocate_symmetric_memory(
shape: int | Sequence[int],
package: ModuleType,
*,
dtype: DTypeLike | torch.dtype | None = None,
axis_order: Literal['C', 'F'] | Sequence[int] | None = None,
make_symmetric: bool = False,
logger: Logger | None = None,
)[source]#

Return uninitialized tensor of given shape and type, allocated from the symmetric heap, on the device on which nvmath.distributed was initialized. The tensor type is determined by the provided package (e.g. cupy, torch). This is a collective operation and must be called by all processes. Note that the buffer size must be the same on all processes, or you can use make_symmetric=True to pad all buffers to the same max size.

Parameters:

  • shape – Shape of the tensor to allocate.

  • package – Python package determining the tensor type (e.g. cupy, torch).

  • dtype – Tensor dtype.

  • axis_order – Axis order.

  • make_symmetric – If buffer sizes do not match across processes, will allocate the maximum size on every process to ensure the allocation is symmetric.

  • logger (logging.Logger) – Python Logger object. The root logger will be used if a logger object is not provided.