nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/distributed/_internal/tensor_ifc_torch.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use distributed Torch tensor objects.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["TorchDistributedTensor"]
+
+import math
+
+try:
+    import torch
+except ImportError:
+    torch = None  # type: ignore
+
+from collections.abc import Sequence
+
+import nvmath.distributed
+from nvmath.distributed._internal.nvshmem import nvshmem_empty_dlpack
+from nvmath.internal.package_ifc import StreamHolder
+from nvmath.internal.tensor_ifc_torch import TorchTensor
+from nvmath.internal.utils import device_ctx
+
+from .tensor_ifc import DistributedTensor
+
+
+# Most methods aren't redefined, because they simply act on the local array
+class TorchDistributedTensor(TorchTensor, DistributedTensor):
+    """
+    TensorHolder for distributed torch tensors.
+    """
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @classmethod
+    def empty(
+        cls, shape, device_id="cpu", *, dtype="float32", strides=None, stream_holder: StreamHolder | None = None, **context
+    ) -> TorchDistributedTensor:
+        """
+        Create an empty tensor of the specified shape and data type on the specified device
+        (None, 'cpu', or device id).
+
+        Note, that the strides, if specified, should correspond to a dense
+        (possibly permuted) tensor and MUST NOT overlap.
+        Otherwise, the behaviour is not defined.
+        """
+        symmetric_memory = context.pop("symmetric_memory", False)
+        make_symmetric = context.pop("make_symmetric", False)
+        skip_symmetric_check = context.pop("skip_symmetric_check", False)
+
+        if device_id == "cpu":
+            if symmetric_memory or make_symmetric or skip_symmetric_check:
+                raise ValueError("symmetric memory options cannot be used when allocating tensor on CPU")
+            return super().empty(shape, device_id, dtype=dtype, strides=strides, **context)
+
+        ctx = nvmath.distributed.get_context()
+        assert ctx is not None, "nvmath.distributed has not been initialized"
+
+        logger = context.get("logger")
+
+        if not symmetric_memory:
+            if make_symmetric or skip_symmetric_check:
+                raise ValueError("Use of symmetric memory option with symmetric_memory=False")
+            return super().empty(
+                shape, device_id=device_id, dtype=dtype, strides=strides, stream_holder=stream_holder, **context
+            )
+
+        dtype = TorchTensor.name_to_dtype[dtype]
+
+        with device_ctx(device_id):
+            size = math.prod(shape, start=dtype.itemsize)
+            # TODO: ideally strides should be set in DLPack, but cuda.core doesn't support
+            # ndarray yet and instead returns a flat buffer.
+            dlpack_buf = nvshmem_empty_dlpack(
+                size,
+                device_id,
+                ctx.process_group,
+                make_symmetric=make_symmetric,
+                skip_symmetric_check=skip_symmetric_check,
+                logger=logger,
+            )
+            tensor = torch.from_dlpack(dlpack_buf)
+            # Buffer may be padded if make_symmetric=True.
+            tensor = tensor[:size]
+            if strides is None:
+                tensor = tensor.view(dtype).view(shape)
+            else:
+                tensor = torch.as_strided(tensor.view(dtype), shape, strides)
+
+        return cls(tensor)
+
+    def to(self, device_id, stream_holder, symmetric_memory: bool = False) -> TorchDistributedTensor:
+        """
+        In addition to the base class semantics:
+          - Source or target device must be the one used to initialize NVSHMEM on this
+            process. This implies that copy from one CUDA device to another is not allowed.
+          - Memory layout is preserved.
+          - Strides must be non-overlapping.
+        """
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        blocking = (device_id == "cpu") or (self.device_id == "cpu")
+
+        # To CPU or same device
+        if device_id == "cpu" or self.device_id == device_id:
+            with stream_holder.ctx:
+                tensor = self.tensor.to(device=device_id, non_blocking=not blocking)
+            result = TorchDistributedTensor(tensor)
+            assert result.is_symmetric_memory == symmetric_memory
+            return result
+
+        # Currently we don't allow copy from one device to another for distributed
+        # tensors (see comment in CupyDistributedTensor.to()).
+        if self.device_id != "cpu" and self.device_id != device_id:
+            raise ValueError("Cannot copy distributed tensor to a different CUDA device")
+
+        # CPU to GPU
+        with stream_holder.ctx:
+            tensor_device = TorchDistributedTensor.empty(
+                self.shape,
+                device_id=device_id,
+                dtype=self.dtype,
+                strides=self.strides,
+                stream_holder=stream_holder,
+                make_symmetric=symmetric_memory,
+                symmetric_memory=symmetric_memory,
+            )
+            tensor_device.tensor.copy_(self.tensor, non_blocking=not blocking)
+            assert tensor_device.is_symmetric_memory == symmetric_memory
+            return tensor_device
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None) -> TorchDistributedTensor:
+        if copy and self.device_id != "cpu":
+            raise NotImplementedError("reshape with copy=True is not supported for TorchDistributedTensor on GPU")
+        return super().reshape(shape, copy=copy)

nvmath/distributed/_internal/tensor_wrapper.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Entry point to using tensors from different libraries seamlessly.
+"""
+
+__all__ = ["wrap_operand", "copy_"]
+
+import warnings
+from collections.abc import Sequence
+
+from nvmath.internal.tensor_ifc import Tensor, TensorHolder
+from nvmath.internal.tensor_wrapper import (
+    infer_tensor_package as base_infer_tensor_package,
+)
+from nvmath.internal.tensor_wrapper import (
+    maybe_register_package as base_maybe_register_package,
+)
+
+from .tensor_ifc import DistributedTensor
+from .tensor_ifc_numpy import CudaDistributedTensor, NumpyDistributedTensor
+
+_TENSOR_TYPES: dict[str, type[DistributedTensor]] = {"numpy": NumpyDistributedTensor, "cuda": CudaDistributedTensor}
+
+
+def infer_tensor_package(tensor):
+    """
+    Infer the package that defines this tensor.
+    """
+    package = base_infer_tensor_package(tensor)
+    # Use call_base=False because base_infer_tensor_package already
+    # called base_maybe_register_package
+    maybe_register_package(package, call_base=False)
+    return package
+
+
+def maybe_register_package(package, call_base=True):
+    if call_base:
+        base_maybe_register_package(package)
+    if package == "torch":
+        from .tensor_ifc_torch import TorchDistributedTensor
+
+        _TENSOR_TYPES[package] = TorchDistributedTensor
+    elif package == "cupy":
+        from .tensor_ifc_cupy import CupyDistributedTensor, HostDistributedTensor
+
+        _TENSOR_TYPES["cupy"] = CupyDistributedTensor
+        _TENSOR_TYPES["cupy_host"] = HostDistributedTensor
+    elif package != "numpy":
+        raise AssertionError(f"Internal error: unrecognized package {package}")
+
+
+def wrap_operand(native_operand: Tensor) -> DistributedTensor[Tensor]:
+    """
+    Wrap one "native" operand so that package-agnostic API can be used.
+    """
+    if isinstance(native_operand, TensorHolder):
+        msg = (
+            "wrap_operand() is being called unnecessarily because the input is already a TensorHolder. "
+            "Only public facing APIs should call wrap_operand(). "
+            "Internal APIs should assume the operands are TensorHolder already."
+            "Trying to wrap a TensorHolder will become an error in the future."
+        )
+        warnings.warn(msg, DeprecationWarning)
+        assert isinstance(native_operand, DistributedTensor)
+        return native_operand
+    wrapped_operand = _TENSOR_TYPES[infer_tensor_package(native_operand)](native_operand)
+    return wrapped_operand
+
+
+def copy_(src: Sequence[TensorHolder], dest: Sequence[TensorHolder], stream_holder):
+    """
+    Copy the wrapped operands in dest to the corresponding wrapped operands in src.
+    """
+    for s, d in zip(src, dest, strict=True):
+        d.copy_(s, stream_holder=stream_holder)

nvmath/distributed/_utils.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+__all__ = ["allocate_symmetric_memory", "free_symmetric_memory"]
+
+from collections.abc import Iterable, Sequence
+from logging import Logger
+from types import ModuleType
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    import torch
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+import nvmath.distributed
+from nvmath.distributed._internal import tensor_wrapper
+from nvmath.internal.utils import device_ctx
+
+# Supported packages for tensors backed by symmetric memory.
+_SUPPORTED_PACKAGES = ("cupy", "torch")
+
+
+# TODO: do a general refactor of this function in nvmath
+def calculate_strides(shape: Sequence[int], axis_order: Iterable[int]):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    strides = [0] * len(shape)
+
+    stride = 1
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides
+
+
+def 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,
+):
+    """Return uninitialized tensor of given shape and type, allocated from the NVSHMEM
+    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.
+
+    Args:
+        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.
+    """
+
+    if package.__name__ not in _SUPPORTED_PACKAGES:
+        raise ValueError(f"The package must be one of {_SUPPORTED_PACKAGES}. Got {package}.")
+
+    distributed_ctx = nvmath.distributed.get_context()
+    if distributed_ctx is None:
+        raise RuntimeError(
+            "nvmath.distributed has not been initialized. Refer to "
+            "https://docs.nvidia.com/cuda/nvmath-python/latest/distributed-apis/runtime.html"
+            " for more information."
+        )
+    if not distributed_ctx.nvshmem_available:
+        raise RuntimeError("nvmath.distributed wasn't initialized with NVSHMEM backend")
+
+    device_id = distributed_ctx.device_id
+
+    if isinstance(shape, int):
+        shape = (shape,)
+
+    strides = None
+    if axis_order is not None:
+        if axis_order == "C":
+            strides = calculate_strides(shape, reversed(range(len(shape))))
+        elif axis_order == "F":
+            strides = calculate_strides(shape, range(len(shape)))
+        else:
+            strides = calculate_strides(shape, axis_order)
+
+    if package.__name__ == "cupy":
+        from ._internal.tensor_ifc_cupy import CupyDistributedTensor
+
+        if dtype is None:
+            # This mimics numpy and cupy
+            dtype = np.float64
+
+        dtype = np.dtype(dtype).name  # type: ignore
+        return CupyDistributedTensor.empty(
+            shape,
+            dtype=dtype,
+            device_id=device_id,
+            strides=strides,
+            symmetric_memory=True,
+            make_symmetric=make_symmetric,
+            skip_symmetric_check=skip_symmetric_check,
+            logger=logger,
+        ).tensor
+    elif package.__name__ == "torch":
+        from ._internal.tensor_ifc_torch import TorchDistributedTensor
+
+        if dtype is None:
+            import torch
+
+            dtype = torch.get_default_dtype()
+
+        dtype = str(dtype).split(".")[1]
+        return TorchDistributedTensor.empty(
+            shape,
+            dtype=dtype,
+            device_id=device_id,
+            strides=strides,
+            symmetric_memory=True,
+            make_symmetric=make_symmetric,
+            skip_symmetric_check=skip_symmetric_check,
+            logger=logger,
+        ).tensor
+
+
+def free_symmetric_memory(*tensors) -> None:
+    """Frees tensors' data buffer where the buffer was allocated on the NVSHMEM
+    symmetric heap. Note that this is only meant to be called on tensors returned
+    by ``allocate_symmetric_memory()``.
+
+    **This is a collective operation and must be called by all processes, with tensors
+    in the same order**."""
+
+    device_id = tensor_wrapper.wrap_operand(tensors[0]).device_id
+    if device_id == "cpu":
+        raise TypeError("free_symmetric_memory called on CPU array/tensor")
+
+    with device_ctx(device_id):
+        for tensor in tensors:
+            wrapped_tensor = tensor_wrapper.wrap_operand(tensor)
+            if wrapped_tensor.device_id == "cpu":
+                raise TypeError("free_symmetric_memory called on CPU array/tensor")
+
+            assert wrapped_tensor.device_id == device_id, "Internal error: symmetric memory tensors are not on the same device"
+
+            wrapped_tensor.free_symmetric()