ShardTensor: fixes for torch.compile with domain parallelism (#1986)

CHANGELOG.md

@@ -394,6 +394,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Padé networks (pytorch/pytorch#193194).
 - Preserve unreduced ShardTensor placement dimensions in `sum` and `mean`
   when `keepdim=True`, including independently sharded simulation and spatial axes.
+- `torch.compile` with ShardTensor: neighborhood attention (`na1d`/`na2d`/
+  `na3d`) keeps its halo exchange when compiled, `select` handles a negative
+  `dim`, dynamic shapes no longer fail in the spec hash, a plain DTensor
+  operand next to a ShardTensor uses the DTensor fallback, and a plain
+  parameter used with a ShardTensor gets a reduced gradient as in eager.
 - Multi-node jobs launched with `srun` no longer hang in
   `DistributedManager.initialize()` when Slurm starts rank 0 on a different node
   from the one that ran `srun`. The rendezvous now uses the job step's first node.

physicsnemo/domain_parallel/_shard_tensor_spec.py

@@ -34,6 +34,21 @@
 from physicsnemo.distributed.utils import compute_split_shapes
 
 
+def _is_symbolic(spec: "ShardTensorSpec") -> bool:
+    r"""Whether the spec's global shape or stride holds a SymInt."""
+    meta = spec.tensor_meta
+    return meta is not None and any(
+        type(d) is not int for d in (*meta.shape, *meta.stride)
+    )
+
+
+def _same_dims(a, b) -> bool:
+    r"""Equal dims without SymInt arithmetic: ints by value, SymInts by identity."""
+    return len(a) == len(b) and all(
+        x is y or (type(x) is int and type(y) is int and x == y) for x, y in zip(a, b)
+    )
+
+
 @dataclass(kw_only=True)
 class ShardTensorSpec(DTensorSpec):
     r"""A distributed tensor specification that tracks sharding information.
@@ -98,8 +113,12 @@ def _hash_impl(self) -> int:
             hash_items.append(self.tensor_meta.dtype)
         if self._sharding_shapes is not None:
             hash_items.append(tuple(sorted(self._sharding_shapes.items())))
-        hash_tuple = tuple(hash_items)
-        return hash(hash_tuple)
+        try:
+            return hash(tuple(hash_items))
+        except TypeError:
+            # Symbolic (SymInt) shapes under torch.compile are unhashable;
+            # collisions are allowed and __eq__ still discriminates.
+            return hash((self.mesh, self.placements))
 
     def __setattr__(self, attr: str, value) -> None:
         r"""Invalidate the cached hash when the sharding shapes change.
@@ -218,6 +237,34 @@ def __eq__(self, other: object) -> bool:
         """
         if not isinstance(other, ShardTensorSpec):
             return False
+        if _is_symbolic(self) or _is_symbolic(other):
+            # Symbolic shapes hash by (mesh, placements) only, so collisions
+            # with specs from other traces reach here: never compare SymInts
+            # (that records guards or fails across traces). Same objects or
+            # equal ints only; anything else is a cache miss, not an error.
+            if self is other:
+                return True
+            a, b = self.tensor_meta, other.tensor_meta
+            if (
+                self.mesh != other.mesh
+                or self.placements != other.placements
+                or (a is None) != (b is None)
+            ):
+                return False
+            if a is not None and not (
+                a.dtype == b.dtype
+                and _same_dims(a.shape, b.shape)
+                and _same_dims(a.stride, b.stride)
+            ):
+                return False
+            sa, sb = self._sharding_shapes, other._sharding_shapes
+            if sa is None or sb is None:
+                return sa is sb
+            return sa.keys() == sb.keys() and all(
+                len(sa[k]) == len(sb[k])
+                and all(_same_dims(x, y) for x, y in zip(sa[k], sb[k]))
+                for k in sa
+            )
         if not super().__eq__(other):
             return False
         if self._sharding_shapes != other._sharding_shapes:
@@ -731,6 +778,8 @@ def validate_aligned_sharding(specs, op_name: str) -> None:
             continue
         cached = getattr(spec, "_sharding_shapes", None) or {}
         shape = tuple(spec.tensor_meta.shape)
+        if any(type(d) is not int for d in shape):
+            continue  # symbolic shapes: boundaries can't be compared here
         key = (spec.mesh, len(shape), placements, tuple(shape[d] for _, d in sharded))
         known = reference.setdefault(key, {})
         for mesh_dim, tensor_dim in sharded:

physicsnemo/domain_parallel/shard_tensor.py

@@ -27,7 +27,8 @@
 
 import torch
 import torch.distributed as dist
-from torch._subclasses.fake_tensor import is_fake
+from torch._subclasses.fake_tensor import FakeTensor, is_fake
+from torch._subclasses.functional_tensor import FunctionalTensor
 from torch.distributed.device_mesh import DeviceMesh, _mesh_resources
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._dtensor_spec import (
@@ -797,6 +798,38 @@ def _is_tracing(args: object, kwargs: object = None) -> bool:
     return False
 
 
+# Plain tensors as they appear inside a traced ``__torch_function__``: a plain
+# tensor or parameter is a FakeTensor (FunctionalTensor during AOT's trace).
+# Exact types only: other subclasses (torchao, AsyncCollectiveTensor, ...) keep
+# their own type under tracing and must not be wrapped in a DTensor.
+_PLAIN_TRACED_TYPES = (
+    torch.Tensor,
+    torch.nn.Parameter,
+    FakeTensor,
+    FunctionalTensor,
+)
+
+
+def _needs_fallback_under_tracing(args: object, kwargs: object = None) -> bool:
+    r"""Whether a traced op must take the DTensor fallback instead of dispatch.
+
+    True for a plain ``DTensor`` operand (dispatch would reject the
+    ShardTensor) or a grad-requiring plain tensor to promote: the fallback's
+    differentiable promotion reduces its ``Partial`` gradient to a plain one,
+    as in eager; promotion at dispatch level records no backward.
+    """
+    promote = (
+        ShardTensor._promotion_mode is not TensorPromotionMode.DISABLED
+        and torch.is_grad_enabled()
+    )
+    for leaf in torch.utils._pytree.tree_leaves((args, kwargs)):
+        if isinstance(leaf, DTensor) and not isinstance(leaf, ShardTensor):
+            return True
+        if promote and type(leaf) in _PLAIN_TRACED_TYPES and leaf.requires_grad:
+            return True
+    return False
+
+
 class ShardTensor(torch.Tensor):
     r"""A distributed tensor class with support for uneven data sharding.
 
@@ -1629,10 +1662,14 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         elif str(func) in cls._named_function_registry and cls._enable_shard_patches:
             args, kwargs = _promote_plain_handler_args(args, kwargs)
             res = cls._named_function_registry[str(func)](func, types, args, kwargs)
-        elif _is_tracing(args, kwargs):
+        elif _is_tracing(args, kwargs) and not _needs_fallback_under_tracing(
+            args, kwargs
+        ):
             # Under torch.compile tracing, route unpatched ops straight to
             # ``__torch_dispatch__`` (like DTensor) so the aten-level path -- which AOT
             # handles correctly -- owns the op and keeps the graph differentiable.
+            # Plain DTensor operands and grad-requiring plain tensors use the
+            # fallback instead (see ``_needs_fallback_under_tracing``).
             with torch._C.DisableTorchFunctionSubclass():
                 return func(*args, **kwargs)
         else:

physicsnemo/domain_parallel/shard_utils/index_ops.py

@@ -884,6 +884,8 @@ def sharded_select_helper(tensor: ShardTensor, dim: int, index: int) -> ShardTen
     # If it's along an unsharded dimension, there is nearly nothing to do.
 
     input_spec = tensor._spec
+    # Dispatch passes ``dim`` as written (e.g. -1 from stack's backward).
+    dim = dim % len(input_spec.shape)
 
     input_placements = input_spec.placements
 
@@ -985,6 +987,7 @@ def sharded_select_backward_helper(
     # If it's along an unsharded dimension, there is nearly nothing to do.
 
     input_placements = grad_output._spec.placements
+    dim = dim % len(input_sizes)
 
     output_stride = _stride_from_contiguous_shape_C_style(input_sizes)
 

physicsnemo/domain_parallel/shard_utils/natten_patches.py

@@ -18,7 +18,6 @@
 
 from typing import Any, Callable
 
-import torch
 from torch.distributed.tensor.placement_types import Shard
 
 from physicsnemo.core.version_check import OptionalImport
@@ -33,14 +32,11 @@
     MissingShardPatch,
     UndeterminedShardingError,
 )
-from physicsnemo.nn.functional.natten import na1d, na2d, na3d
+from physicsnemo.nn.functional import natten as _na_ops  # noqa: F401 (defines the ops)
 
 _natten = OptionalImport("natten")
-_raw_func_map = {
-    na1d: lambda: _natten.functional.na1d,
-    na2d: lambda: _natten.functional.na2d,
-    na3d: lambda: _natten.functional.na3d,
-}
+# Argument order of the ``physicsnemo::na{1,2,3}d`` ops.
+_OP_ARGS = ("q", "k", "v", "kernel_size", "dilation", "is_causal", "scale", "backend")
 
 __all__ = ["na1d_wrapper", "na2d_wrapper", "na3d_wrapper"]
 
@@ -226,61 +222,63 @@ def _natten_wrapper(
     types: tuple[Any, ...],
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
-) -> torch.Tensor | ShardTensor:
-    r"""Shared wrapper for natten functions to support sharded tensors.
+) -> ShardTensor:
+    r"""ShardTensor handler for the ``physicsnemo::na{1,2,3}d`` ops.
 
-    Registered with :meth:`ShardTensor.register_function_handler` so that calls
-    to :func:`~physicsnemo.nn.functional.natten.na1d`,
-    :func:`~physicsnemo.nn.functional.natten.na2d`, or
-    :func:`~physicsnemo.nn.functional.natten.na3d` automatically route through
-    this handler when any argument is a :class:`ShardTensor`.
+    :func:`~physicsnemo.nn.functional.natten.na1d`,
+    :func:`~physicsnemo.nn.functional.natten.na2d` and
+    :func:`~physicsnemo.nn.functional.natten.na3d` call these ops for tensor
+    subclass operands. Registered by op name, so ``__torch_function__`` routes
+    the op here in eager and under ``torch.compile``; the halo exchange and
+    the local NATTEN call carry autograd.
 
     Parameters
     ----------
     func : Callable
-        The wrapped natten function (passed by ``__torch_function__``).
+        The op overload, e.g. ``torch.ops.physicsnemo.na2d.default``.
     types : tuple[Any, ...]
         The types of the inputs (unused).
     args : tuple[Any, ...]
-        Positional arguments containing query, key, value tensors and kernel_size.
+        Op arguments, in the order of ``_OP_ARGS``.
     kwargs : dict[str, Any]
-        Keyword arguments including ``dilation``.
+        Op arguments passed by keyword.
 
     Returns
     -------
-    Union[torch.Tensor, ShardTensor]
-        Result tensor as either ``torch.Tensor`` or ShardTensor depending on input types.
+    ShardTensor
+        Neighborhood attention output, sharded like ``q``.
 
     Raises
     ------
     UndeterminedShardingError
-        If input tensor types are mismatched.
+        If q, k and v are not all ShardTensors.
     """
-    q, k, v, kernel_size = args[0], args[1], args[2], args[3]
-
-    dilation = kwargs.get("dilation", 1)
-    natten_kwargs = {_k: _v for _k, _v in kwargs.items() if _k != "dilation"}
-
-    if all(type(_t) is torch.Tensor for _t in (q, k, v)):
-        return func(
-            q, k, v, kernel_size=kernel_size, dilation=dilation, **natten_kwargs
-        )
-    elif all(isinstance(_t, ShardTensor) for _t in (q, k, v)):
-        raw_func = _raw_func_map[func]()
-        return _partial_natten(
-            q, k, v, kernel_size, dilation, base_func=raw_func, **natten_kwargs
-        )
-    else:
+    bound = {**dict(zip(_OP_ARGS, args)), **kwargs}
+    q, k, v = bound["q"], bound["k"], bound["v"]
+    if not all(isinstance(_t, ShardTensor) for _t in (q, k, v)):
         raise UndeterminedShardingError(
-            "q, k, and v must all be the same types (torch.Tensor or ShardTensor)"
+            "q, k, and v must all be ShardTensors for sharded neighborhood attention"
         )
+    raw_func = getattr(_natten.functional, func.__name__.split(".")[0])
+    return _partial_natten(
+        q,
+        k,
+        v,
+        bound["kernel_size"],
+        bound["dilation"],
+        base_func=raw_func,
+        is_causal=bound["is_causal"],
+        scale=bound["scale"],
+        backend=bound["backend"],
+    )
 
 
 # Public aliases for explicit registration
 na1d_wrapper = _natten_wrapper
 na2d_wrapper = _natten_wrapper
 na3d_wrapper = _natten_wrapper
 
-ShardTensor.register_function_handler(na1d, na1d_wrapper)
-ShardTensor.register_function_handler(na2d, na2d_wrapper)
-ShardTensor.register_function_handler(na3d, na3d_wrapper)
+for _name in ("na1d", "na2d", "na3d"):
+    ShardTensor.register_named_function_handler(
+        f"physicsnemo.{_name}.default", _natten_wrapper
+    )