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] = 'C',
make_symmetric: bool = False,
skip_symmetric_check: 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 in a form recognized by the package. If None, will use the package’s default dtype.

  • axis_order – Axis order. The default is ‘C’ (row-major ordering).

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

  • skip_symmetric_check – Skip checking that the allocation is symmetric (which requires inter-process communication). The default is False.

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