physics_mcp/source/physicsnemo/distributed/shard_tensor.py

# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Iterable
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
from warnings import warn

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh, _mesh_resources

from physicsnemo.distributed import DistributedManager
from physicsnemo.utils.profiling import annotate, profile
from physicsnemo.utils.version_check import check_module_requirements

# Prevent importing this module if the minimum version of pytorch is not met.
check_module_requirements("physicsnemo.distributed.shard_tensor")

from torch.distributed.tensor import DTensor  # noqa: E402
from torch.distributed.tensor._dtensor_spec import (  # noqa: E402
    TensorMeta,
)
from torch.distributed.tensor.placement_types import (  # noqa: E402
    Placement,
    Replicate,
    Shard,
)

from physicsnemo.distributed._shard_redistribute import (  # noqa: E402
    ShardRedistribute,
)
from physicsnemo.distributed._shard_tensor_spec import (  # noqa: E402
    ShardTensorSpec,
    _infer_shard_tensor_spec_from_local_chunks,
    _stride_from_contiguous_shape_C_style,
)

aten = torch.ops.aten


class _ToTorchTensor(torch.autograd.Function):
    """Autograd function to convert a ShardTensor to a regular PyTorch tensor.

    This class handles the conversion from ShardTensor to torch.Tensor in both forward
    and backward passes, maintaining proper gradient flow.  Slices the ShardTensor
    to the local component only on the current rank.
    """

    @staticmethod
    def forward(
        ctx: torch.autograd.function.FunctionCtx,
        input: "ShardTensor",
        grad_placements: Optional[Sequence[Placement]] = None,
    ) -> torch.Tensor:
        """Convert ShardTensor to torch.Tensor in forward pass.

        Args:
            ctx: Autograd context for saving tensors/variables for backward
            input: ShardTensor to convert
            grad_placements: Optional sequence of placements to use for gradients

        Returns:
            torch.Tensor: Local tensor representation of the ShardTensor
        """
        ctx.shard_tensor_spec = input._spec
        ctx.grad_placements = grad_placements
        local_tensor = input._local_tensor

        # JUST LIKE DTENSOR:
        # We need to return a fresh Tensor object there as autograd metadata
        # will be inplaced into it. So we don't want to pollute the Tensor
        # object stored in the _local_tensor of this ShardTensor.
        return local_tensor.view_as(local_tensor)

    @staticmethod
    def backward(
        ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor
    ) -> Tuple["ShardTensor", None]:
        """Convert gradient torch.Tensor back to ShardTensor in backward pass.

        Args:
            ctx: Autograd context containing saved tensors/variables from forward
            grad_output: Gradient tensor to convert back to ShardTensor

        Returns:
            Tuple containing:
            - ShardTensor gradient
            - None for grad_placements gradient (not needed)
        """
        shard_tensor_spec = ctx.shard_tensor_spec
        mesh = shard_tensor_spec.mesh
        if ctx.grad_placements is not None:
            if ctx.grad_placements != shard_tensor_spec.placements:
                grad_placements = ctx.grad_placements
                grad_sharding_shapes = "infer"
            else:
                # If the placements are the same as the input placements,
                # we reuse the sharding sizes from the input placements.
                grad_placements = ctx.grad_placements
                grad_sharding_shapes = shard_tensor_spec._sharding_shapes
        else:
            grad_placements = shard_tensor_spec.placements
            grad_sharding_shapes = shard_tensor_spec._sharding_shapes
        if grad_sharding_shapes is None:
            grad_sharding_shapes = "infer"
        # Generate a spec based on grad outputs and the expected placements:
        grad_tensor_spec = _infer_shard_tensor_spec_from_local_chunks(
            grad_output, mesh, grad_placements, grad_sharding_shapes
        )

        return (
            ShardTensor(
                grad_output, grad_tensor_spec, requires_grad=grad_output.requires_grad
            ),
            None,
        )


class _FromTorchTensor(torch.autograd.Function):
    """Autograd function for converting a torch.Tensor to a ShardTensor.

    This class handles the forward and backward passes for converting between
    torch.Tensor and ShardTensor types, maintaining gradient information.

    Global shape information is inferred using collective communication on
    the specified device mesh.
    """

    @staticmethod
    def forward(
        ctx: torch.autograd.function.FunctionCtx,
        local_input: torch.Tensor,
        device_mesh: DeviceMesh,
        placements: Tuple[Placement, ...],
        sharding_shapes: Union[str, Dict[int, List[Tuple[int, ...]]]] = "chunk",
    ) -> "ShardTensor":
        """Convert a local torch.Tensor to a ShardTensor in forward pass.

        Args:
            ctx: Autograd context for saving tensors/variables for backward
            local_input: Local tensor to convert to ShardTensor
            device_mesh: Device mesh specifying process groups
            placements: Tuple of placement rules for sharding
            sharding_shapes: Controls how shard tensor spec is generated:
                - "chunk": Use torch.chunk shapes to infer shapes from global shape (no communication)
                - "infer": Use collective communication to infer shapes from mesh neighbors.
                - Manual dict mapping mesh dim to list of shard shapes: Use provided shapes.  Must pass on each rank!

        Returns:
            ShardTensor constructed from the local input tensor
        """
        ctx.previous_placement = placements
        ctx.previous_mesh = device_mesh

        # This function is simpler than the corresponding DTensor implementation on the surface
        # because under the hood, we have some logic here to infer the sharding shapes.
        shard_tensor_spec = _infer_shard_tensor_spec_from_local_chunks(
            local_input, device_mesh, placements, sharding_shapes
        )

        shard_tensor = ShardTensor(
            local_input,
            shard_tensor_spec,
            requires_grad=local_input.requires_grad,
        )

        return shard_tensor

    @staticmethod
    def backward(
        ctx: torch.autograd.function.FunctionCtx,
        grad_output: "ShardTensor",
    ) -> Tuple[torch.Tensor, None, None]:
        """Convert gradient ShardTensor back to torch.Tensor in backward pass.

        Args:
            ctx: Autograd context containing saved tensors/variables from forward
            grad_output: Gradient ShardTensor to convert back to torch.Tensor

        Returns:
            Tuple containing:
            - Local tensor gradient
            - None for device_mesh gradient (not needed)
            - None for placements gradient (not needed)

        Raises:
            RuntimeError: If gradient tensor has different placement than original
        """
        previous_placement = ctx.previous_placement
        if grad_output.placements != previous_placement:
            # Automatically redistribute to the previous placement as long as it's not a partial.
            if not any(p.is_partial() for p in previous_placement):
                grad_output = grad_output.redistribute(
                    grad_output._spec.mesh, previous_placement
                )
            else:
                raise RuntimeError(
                    "Resharding gradients with partial placements not implemented"
                )

        return grad_output.to_local(), None, None, None


class ShardTensor(DTensor):
    """
    A class similar to pytorch's native DTensor but with more
    flexibility for uneven data sharding.

    Leverages very similar API to DTensor (identical, where possible)
    but deliberately tweaking routines to avoid implicit assumptions
    about tensor sharding.

    The key differences from DTensor are:
    - Supports uneven sharding where different ranks can have different local tensor sizes
    - Tracks and propagates shard size information across operations
    - Handles redistribution of unevenly sharded tensors
    - Provides custom collective operations optimized for uneven sharding

    Like DTensor, operations are dispatched through PyTorch's dispatcher system.
    Most operations work by:
    1. Converting inputs to local tensors
    2. Performing the operation locally
    3. Constructing a new ShardTensor with appropriate sharding spec
    4. Handling any needed communication between ranks

    The class provides methods for:
    - Converting to/from local tensors
    - Redistributing between different sharding schemes
    - Performing collective operations like all_gather and reduce_scatter
    - Basic tensor operations that maintain sharding information
    """

    _local_tensor: torch.Tensor
    _spec: ShardTensorSpec
    __slots__ = ["_local_tensor", "_spec"]

    # For torch.ops.aten operators (low-level dispatch)
    _dispatch_registry: Dict[torch._ops.OpOverload, Callable] = {}

    # For Python-level functions (torch.mean, tensor.mean, etc.)
    _function_registry: Dict[Callable, Callable] = {}

    # For custom functions registered with PyTorch,
    # it is sometimes necessary to match by name.
    # For instance, if you declare an op with
    #
    # @torch.library.custom_op(
    #    "module::function_name", mutates_args=()
    # )
    # def function_external_to_torch(
    #
    # Then, you likely want to register the handler with
    #
    # ShardTensor.register_named_function_handler("module.function_name.default", handler)
    _named_function_registry: dict[str, Callable] = {}

    # Upon construction of any ShardTensor objects, this will be set to true.
    # Wrappers are triggered dynamically, so the wrapping will be pass-through
    # exclusively until true.
    _enable_shard_patches: bool = False

    @classmethod
    def patches_enabled(cls) -> bool:
        """
        Whether to enable patches for this class.

        Default is False, but can be changed by the user.
        """
        return cls._enable_shard_patches

    @classmethod
    def register_dispatch_handler(
        cls, op: torch._ops.OpOverload, handler: Callable
    ) -> None:
        """Register a handler for a specific PyTorch operator in the dispatch system."""
        cls._dispatch_registry[op] = handler

    @classmethod
    def register_function_handler(cls, func: Callable, handler: Callable) -> None:
        """Register a handler for a Python-level function or method."""
        cls._function_registry[func] = handler

    @classmethod
    def register_named_function_handler(cls, func_name: str, handler: Callable) -> None:
        """Register a named function that has been named via torch.library.custom_op"""
        cls._named_function_registry[func_name] = handler

    @staticmethod
    def __new__(
        cls,
        local_tensor: torch.Tensor,
        spec: ShardTensorSpec,
        *,
        requires_grad: bool,
    ) -> "ShardTensor":
        """
        Construct a new Shard Tensor from a local tensor, device mesh, and placement.

        Note that unlike DTensor, ShardTensor will automatically collect the Shard size
        information from all participating devices. This is to enable uneven and
        dynamic sharding.

        Heavily derived from torch DTensor

        Args:
            local_tensor: Local tensor to use as the data
            spec: ShardTensorSpec defining the sharding scheme
            requires_grad: Whether the tensor requires gradients

        Returns:
            A new ShardTensor instance
        """
        if local_tensor.requires_grad and not requires_grad:
            warn(
                "To construct a new ShardTensor from torch.Tensor, "
                "it's recommended to use local_tensor.detach() and "
                "make requires_grad consistent."
            )

        if spec.tensor_meta is None:
            raise ValueError("TensorMeta should not be None!")

        # Check the sharding information is known:
        ret = torch.Tensor._make_wrapper_subclass(
            cls,
            spec.tensor_meta.shape,
            strides=spec.tensor_meta.stride,
            dtype=local_tensor.dtype,
            device=local_tensor.device,
            layout=local_tensor.layout,
            requires_grad=requires_grad,
        )

        ret._spec = spec
        ret._local_tensor = local_tensor

        cls._enable_shard_patches = True

        return ret

    def __repr__(self) -> str:
        return f"ShardTensor(local_tensor={self._local_tensor}, device_mesh={self._spec.mesh}, placements={self._spec.placements})"

    @classmethod
    def from_dtensor(cls, dtensor: DTensor) -> "ShardTensor":
        """
        Convert a DTensor to a ShardTensor.  We assume the DTensor is properly constructed.

        Args:
            dtensor: DTensor to convert

        Returns:
            Equivalent ShardTensor
        """
        # Always ensure sharding is turned on:
        cls._enable_shard_patches = True

        # DTensor is locked to sharding a tensor according to chunk format.
        # We can use that to infer sharding sizes with no communication.

        # Create the spec by inferring the sharding sizes from the DTensor:
        spec = _infer_shard_tensor_spec_from_local_chunks(
            dtensor._local_tensor,
            dtensor._spec.mesh,
            dtensor._spec.placements,
            sharding_shapes="chunk",
            global_shape=dtensor.shape,
        )

        return ShardTensor.__new__(
            cls,
            local_tensor=dtensor._local_tensor,
            spec=spec,
            requires_grad=dtensor.requires_grad,
        )

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs={}):
        with annotate(f"__torch_function___{func.__name__}"):
            # Check for overrides:
            if func in cls._function_registry and cls._enable_shard_patches:
                res = cls._function_registry[func](func, types, args, kwargs)
                return res
            elif (
                str(func) in cls._named_function_registry and cls._enable_shard_patches
            ):
                res = cls._named_function_registry[str(func)](func, types, args, kwargs)
                return res
            # Fall back to the default behavior:
            return super().__torch_function__(func, types, args, kwargs)

    @classmethod
    @torch._disable_dynamo
    @profile
    def __torch_dispatch__(
        cls,
        func: torch._ops.OpOverload,
        types: Tuple[type, ...],
        args: Tuple[object, ...] = (),
        kwargs: Optional[Dict[str, object]] = None,
    ) -> Union["ShardTensor", Iterable["ShardTensor"], object]:
        with annotate(f"__torch_dispatch___{func.__name__}"):
            # Leverage DTensor Dispatch as much as possible, but, enable
            # the ability to operate on this output in the future:
            if func in cls._dispatch_registry:
                res = cls._dispatch_registry[func](*args, **kwargs)
                return res

            # We assume that if we reach this point, the operator has not been
            # intercepted by a wrapper or in the registry.  So the DTensor
            # default behavior is likely to be correct.

            if func == aten.view.default:
                # For view, we need input tensors to be contiguous:
                for arg in args:
                    if isinstance(arg, ShardTensor) or isinstance(arg, DTensor):
                        if not arg._local_tensor.is_contiguous():
                            arg._local_tensor = arg._local_tensor.contiguous()

            dispatch_res = DTensor._op_dispatcher.dispatch(func, args, kwargs or {})

            # Return a shard tensor instead of a dtensor.
            def _convert_dtensor_with_input_check(dtensor, input_args):
                """
                This function searches the input for ShardTensors that match output shapes.
                It prevents collectives, since we can copy the sharding shapes for irregular shards.
                The idea here is that, if the global shape is unchanged, and
                the placements are unchanged, the sharding shapes should be unchanged.

                If no matches are found, it falls back to inference based on DTensor.

                This is only used when we already went back through the DTensor dispatch.
                """
                # Check if this matches any input ShardTensor
                for arg in input_args:
                    if (
                        isinstance(arg, ShardTensor)
                        and dtensor._spec.tensor_meta == arg._spec.tensor_meta
                        and dtensor._spec.placements == arg._spec.placements
                    ):
                        return ShardTensor.__new__(
                            ShardTensor,
                            local_tensor=dtensor._local_tensor,
                            spec=arg._spec,
                            requires_grad=dtensor.requires_grad,
                        )

                # Fall back to default conversion
                return ShardTensor.from_dtensor(dtensor)

            if isinstance(dispatch_res, DTensor):
                return _convert_dtensor_with_input_check(dispatch_res, args)

            if isinstance(dispatch_res, Iterable):
                return type(dispatch_res)(
                    _convert_dtensor_with_input_check(d, args)
                    if isinstance(d, DTensor)
                    else d
                    for d in dispatch_res
                )

            return dispatch_res

    @staticmethod
    def from_local(
        local_tensor: torch.Tensor,
        device_mesh: Optional[DeviceMesh] = None,
        placements: Optional[Sequence[Placement]] = None,
        sharding_shapes: Union[str, Dict[int, List[Tuple[int, ...]]]] = "infer",
    ) -> "ShardTensor":
        """
        Generate a new ShardTensor from local torch tensors. Uses
        device mesh and placements to infer global tensor properties.

        No restriction is made on forcing tensors to have equal shapes
        locally. Instead, the requirement is that tensor shapes could
        be concatenated into a single tensor according to the placements.

        Args:
            local_tensor: Local chunk of tensor. All participating tensors must be
                of the same rank and concatable across the mesh dimensions
            device_mesh: Target Device Mesh, if not specified will use the current mesh
            placements: Target placements, must have same number of elements as device_mesh.ndim
            sharding_shapes: Controls how shard tensor spec is generated:
                - "chunk": Use torch.chunk shapes to infer shapes from global shape (no communication)
                - "infer": Use collective communication to infer shapes from mesh neighbors.
                - Manual dict mapping mesh dim to list of shard shapes: Use provided shapes.  Must pass on each rank!
        Returns:
            A new ShardTensor instance
        """

        # this turns on shard patches globally for this process.
        ShardTensor._enable_shard_patches = True

        # This implementation follows the pytorch DTensor Implementation Closely.
        device_mesh = device_mesh or _mesh_resources.get_current_mesh()
        device_type = device_mesh.device_type

        # convert the local tensor to desired device base on device mesh's device_type
        if device_type != local_tensor.device.type and not local_tensor.is_meta:
            local_tensor = local_tensor.to(device_type)

        # set default placements to replicated if not specified
        if placements is None:
            placements = [Replicate() for _ in range(device_mesh.ndim)]
        else:
            placements = list(placements)
            for idx, placement in enumerate(placements):
                # normalize shard dim to be positive
                if placement.is_shard():
                    placement = cast(Shard, placement)
                    if placement.dim < 0:
                        placements[idx] = Shard(placement.dim + local_tensor.ndim)

        # `from_local` is differentiable, and the gradient of the dist tensor this function
        # created should flow back the gradients to the local_tensor, so we call an autograd
        # function to construct the dist tensor instead.
        return _FromTorchTensor.apply(  # pyre-ignore[16]: autograd func
            local_tensor,
            device_mesh,
            tuple(placements),
            sharding_shapes,
        )

    def offsets(self, mesh_dim: Optional[int] = None) -> List[int]:
        """
        Get offsets of shards along a mesh dimension.

        Args:
            mesh_dim: Mesh dimension to get offsets for. If None, returns all offsets.

        Returns:
            List of offsets for shards along specified dimension
        """
        return self._spec.offsets(mesh_dim)

    def redistribute(
        self,
        device_mesh: Optional[DeviceMesh] = None,
        placements: Optional[Sequence[Placement]] = None,
        *,
        async_op: bool = False,
    ) -> "ShardTensor":
        """
        Redistribute tensor across device mesh with new placement scheme.
        Like DTensor redistribute but uses custom layer for shard redistribution.

        Args:
            device_mesh: Target device mesh. Uses current if None.
            placements: Target placement scheme. Required.
            async_op: Whether to run asynchronously

        Returns:
            Redistributed ShardTensor

        Raises:
            RuntimeError: If placements not specified or invalid
        """

        # if device_mesh is not specified, use the current device_mesh
        device_mesh = device_mesh or self.device_mesh
        # raise error if new placements not specified
        if placements is None:
            raise RuntimeError("placements is needed for redistribute!")

        placements = list(placements)
        for i, placement in enumerate(placements):
            if placement.is_partial():
                raise RuntimeError(
                    "Can not redistribute to Partial, redistributing to Partial is for internal use only!"
                )
            elif isinstance(placement, Shard) and placement.dim < 0:
                # normalize shard dim to be positive
                placements[i] = Shard(placement.dim + self.ndim)
        placements = tuple(placements)

        return ShardRedistribute.apply(self, device_mesh, placements, async_op)

    def to_local(
        self, *, grad_placements: Optional[Sequence[Placement]] = None
    ) -> torch.Tensor:
        """
        Get local tensor from this ShardTensor.

        Args:
            grad_placements: Future layout of gradients. Optional.

        Returns:
            Local torch.Tensor. Shape may vary between ranks for sharded tensors.
        """

        if not torch.is_grad_enabled():
            return self._local_tensor

        if grad_placements is not None and not isinstance(grad_placements, tuple):
            grad_placements = tuple(grad_placements)

        return _ToTorchTensor.apply(self, grad_placements)

    def full_tensor(
        self, *, grad_placements: Optional[Sequence[Placement]] = None
    ) -> torch.Tensor:
        """
        Need to re-implement here to ensure a ShardTensor is used as the output
        of redistribute.
        """

        redist_res = self.redistribute(
            placements=[Replicate()] * self.device_mesh.ndim, async_op=False
        )
        return _ToTorchTensor.apply(redist_res, grad_placements)

    def backward(self, *args, **kwargs):
        """
        Backward pass for ShardTensor.

        This method is used to perform the backward pass for a ShardTensor.
        It handles the redistribution of the tensor to the desired placements
        and then calls the backward pass on the local tensor.
        """

        # Before calling backward, we need to resolve any partial placements.
        new_placements = []
        # grad_placements = []
        needs_redistribute = False
        for i, placement in enumerate(self._spec.placements):
            if placement.is_partial():
                new_placements.append(Replicate())
                # grad_placements.append(Shard(i))
                needs_redistribute = True
            else:
                new_placements.append(placement)
                # grad_placements.append(placement)

        if needs_redistribute:
            self = self.redistribute(placements=new_placements)

        return self.to_local().backward(*args, **kwargs)


def scatter_tensor(
    tensor: torch.Tensor,
    global_src: int,
    mesh: DeviceMesh,
    placements: Tuple[Placement, ...],
    global_shape: Optional[torch.Size] = None,
    dtype: Optional[torch.dtype] = None,
    requires_grad: bool = False,
) -> "ShardTensor":
    """
    Take a tensor from source rank and distribute it across devices on the mesh according to placements.

    This function takes a tensor that exists on a single source rank and distributes it across
    a device mesh according to the specified placement scheme. For multi-dimensional meshes,
    it performs a flattened scatter operation before constructing the sharded tensor.

    Args:
        tensor: The tensor to distribute, must exist on source rank
        global_src: Global rank ID of the source process
        mesh: Device mesh defining the process topology
        placements: Tuple of placement specifications defining how to distribute the tensor

    Returns:
        ShardTensor: The distributed tensor with specified placements

    Raises:
        ValueError: If global_src is not an integer or not in the mesh
    """
    dm = DistributedManager()

    if not isinstance(global_src, int):
        raise ValueError("Global source must be an integer rank")
    if global_src not in mesh.mesh:
        raise ValueError("Please specify a tensor source in this mesh")

    is_src = dm.rank == global_src

    # For multi-dimensional meshes, we use a flattened process group
    mesh_group = dm.get_mesh_group(mesh)

    # Broadcast tensor metadata from source
    if global_shape is None or dtype is None:
        if dm.rank == global_src:
            meta = [TensorMeta(tensor.shape, tensor.stride(), tensor.dtype)]
        else:
            meta = [None]

        dist.broadcast_object_list(meta, src=global_src, group=mesh_group)

        local_meta = meta[0]
    else:
        stride = _stride_from_contiguous_shape_C_style(global_shape)
        local_meta = TensorMeta(global_shape, stride, dtype)

    # This needs to be optimized, but I want to get the whole pipeline optimized first.
    # This only gets done when scatter_tensor is called and it should be relatively small
    # in full applications.

    # What isn't optimmized?  Broadcasting the full tensor when placement is likely
    # Shard on at least one mesh dimension.  It would be more efficient to iteratively
    # scatter along Shard dimensions.  BUT, the focus is on performance of full applications
    # and this is a once-per-iteration cost.

    # Broadcast the tensor to all ranks
    if tensor is None and not is_src:
        # Tensor is allowed to be none if not on the root rank
        tensor = torch.empty(local_meta.shape, dtype=local_meta.dtype, device=dm.device)

    dist.broadcast(tensor, src=global_src, group=mesh_group)

    # Create a fully-replicated spec:
    spec = ShardTensorSpec(
        mesh=mesh,
        placements=[Replicate() for _ in range(mesh.ndim)],
        tensor_meta=local_meta,
        _sharding_shapes={},
    )

    # Make a "fully-replicated" tensor on all ranks:
    st = ShardTensor.__new__(
        ShardTensor,
        local_tensor=tensor,
        spec=spec,
        requires_grad=requires_grad,
    )

    # Redistribute the tensor to the desired placements:
    st = st.redistribute(mesh, placements, async_op=False)
    # This is an unoptimal step but is functional:
    if requires_grad:
        st = st.detach()
        st.requires_grad = True

    return st