spmd_types 0.2.1__tar.gz → 0.2.2__tar.gz

{spmd_types-0.2.1/spmd_types.egg-info → spmd_types-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.1
+Version: 0.2.2
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License
@@ -50,14 +50,29 @@ License-File: LICENSE
 Requires-Dist: torch>=2.10.0
 Dynamic: license-file
 
-# spmd_types
+# `spmd_types`
 
-A type system for distributed (SPMD) tensor computations in PyTorch.
+A type system for distributed (SPMD) tensor computations in PyTorch.  This
+package provides two type systems:
 
-spmd_types tracks per-mesh-axis types on tensors -- Replicate (R), Invariant (I),
-Varying (V), Partial (P), and Shard (S) -- and enforces type-correct transitions
-through collective operations and local rewrites. It catches distributed
-programming errors at development time without requiring a GPU cluster.
+* **Local SPMD types**, which allow you to use Megatron-style differentiable
+  collectives in a safe way by tracking whether or not your backward gradients
+  are pending reduction or not.
+
+* **Global SPMD types**, a DTensor-like abstraction for writing code that has
+  the same semantics whether run on a single device or in a distributed
+  fashion, but with explicit communication operations so you are never
+  guessing when a redistribute occurs.
+
+In both cases, the SPMD types makes it possible for you to check that your
+code computes correct gradients (local SPMD) or gives equivalent results
+across different parallelizations (global SPMD), without having to actually
+run a full E2E distributed training run to check for loss matching.
+
+The goal of this package is to provide a flexible type system that can
+typecheck realistic training code.  We have used local SPMD types to typecheck
+a realistic pretraining codebase, and global SPMD types is actively under
+construction!
 
 ## Installation
 
@@ -71,6 +86,7 @@ pip install spmd_types
 import torch
 import torch.distributed as dist
 import spmd_types as spmd
+import spmd_types.checker
 from torch.distributed.device_mesh import init_device_mesh
 
 # Set up a fake process group (no GPUs needed)
@@ -79,12 +95,12 @@ mesh = init_device_mesh("cpu", (2, 4), mesh_dim_names=("dp", "tp"))
 dp = mesh.get_group("dp")
 tp = mesh.get_group("tp")
 
-with spmd.set_current_mesh(mesh), spmd.typecheck():
+with spmd.set_current_mesh(mesh), spmd.checker.typecheck():
     x = torch.randn(4)
-    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})  # R on dp, partial on tp
+    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})       # R on dp, partial on tp
     y = spmd.all_reduce(x, tp, src=spmd.P, dst=spmd.R)  # sum across tp ranks
-    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})   # now replicated everywhere
-    z = torch.mul(y, y)                              # type inference: R * R -> R
+    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})       # now replicated everywhere
+    z = torch.mul(y, y)                                 # type inference: R * R -> R
     spmd.assert_type(z, {dp: spmd.R, tp: spmd.R})
 
 dist.destroy_process_group()
@@ -92,12 +108,14 @@ dist.destroy_process_group()
 
 ## Documentation
 
-See [MEGATRON_QUICKSTART.md](MEGATRON_QUICKSTART.md) for a guide on porting
-Megatron-derived training frameworks to use spmd_types.
+See [Local SPMD types](docs/local_spmd_types.md) for a hands-on guide on
+porting Megatron-derived training frameworks, including the Megatron-to-spmd_types
+function mapping table and advice on Invariant vs Replicate.
 
-See [DESIGN.md](DESIGN.md) for the full type system specification, including
-type inference rules, collective signatures, and forward-backward pairs.
+See [Design](docs/design.md) for the full type system specification, including
+local vs global SPMD modes, collective signatures with diagrams, forward-backward
+pairs, expert mode, cross-mesh compatibility, and partition spec redistribute.
 
 ## License
 
-BSD 3-Clause License. See [LICENSE](LICENSE) for details.
+BSD 3-Clause License. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to contribute.

spmd_types-0.2.2/README.md

@@ -0,0 +1,69 @@
+# `spmd_types`
+
+A type system for distributed (SPMD) tensor computations in PyTorch.  This
+package provides two type systems:
+
+* **Local SPMD types**, which allow you to use Megatron-style differentiable
+  collectives in a safe way by tracking whether or not your backward gradients
+  are pending reduction or not.
+
+* **Global SPMD types**, a DTensor-like abstraction for writing code that has
+  the same semantics whether run on a single device or in a distributed
+  fashion, but with explicit communication operations so you are never
+  guessing when a redistribute occurs.
+
+In both cases, the SPMD types makes it possible for you to check that your
+code computes correct gradients (local SPMD) or gives equivalent results
+across different parallelizations (global SPMD), without having to actually
+run a full E2E distributed training run to check for loss matching.
+
+The goal of this package is to provide a flexible type system that can
+typecheck realistic training code.  We have used local SPMD types to typecheck
+a realistic pretraining codebase, and global SPMD types is actively under
+construction!
+
+## Installation
+
+```bash
+pip install spmd_types
+```
+
+## Quick start
+
+```python
+import torch
+import torch.distributed as dist
+import spmd_types as spmd
+import spmd_types.checker
+from torch.distributed.device_mesh import init_device_mesh
+
+# Set up a fake process group (no GPUs needed)
+dist.init_process_group(backend="fake", rank=0, world_size=8)
+mesh = init_device_mesh("cpu", (2, 4), mesh_dim_names=("dp", "tp"))
+dp = mesh.get_group("dp")
+tp = mesh.get_group("tp")
+
+with spmd.set_current_mesh(mesh), spmd.checker.typecheck():
+    x = torch.randn(4)
+    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})       # R on dp, partial on tp
+    y = spmd.all_reduce(x, tp, src=spmd.P, dst=spmd.R)  # sum across tp ranks
+    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})       # now replicated everywhere
+    z = torch.mul(y, y)                                 # type inference: R * R -> R
+    spmd.assert_type(z, {dp: spmd.R, tp: spmd.R})
+
+dist.destroy_process_group()
+```
+
+## Documentation
+
+See [Local SPMD types](docs/local_spmd_types.md) for a hands-on guide on
+porting Megatron-derived training frameworks, including the Megatron-to-spmd_types
+function mapping table and advice on Invariant vs Replicate.
+
+See [Design](docs/design.md) for the full type system specification, including
+local vs global SPMD modes, collective signatures with diagrams, forward-backward
+pairs, expert mode, cross-mesh compatibility, and partition spec redistribute.
+
+## License
+
+BSD 3-Clause License. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to contribute.

{spmd_types-0.2.1 → spmd_types-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.1"
+version = "0.2.2"
 description = "A type system for distributed (SPMD) tensor computations in PyTorch"
 readme = "README.md"
 license = {file = "LICENSE"}

{spmd_types-0.2.1 → spmd_types-0.2.2}/spmd_types/__init__.py

@@ -53,6 +53,7 @@ from spmd_types._type_attr import (  # noqa: F401
 # Collectives and operations -- runtime API (no _checker dependency)
 from spmd_types.runtime import (  # noqa: F401
     assert_local_type,
+    assert_local_type_like,
     assert_type,
     assert_type_like,
     Infer,
@@ -84,6 +85,7 @@ from spmd_types.types import (  # noqa: F401
     Replicate,
     S,
     Shard,
+    SpmdType,
     SpmdTypeError,
     TensorSharding,
     V,

{spmd_types-0.2.1 → spmd_types-0.2.2}/spmd_types/_checker/__init__.py

@@ -74,8 +74,6 @@ from spmd_types.runtime import (  # noqa: F401
     trace,
 )
 from spmd_types.types import (
-    _canonicalize_shard,
-    _check_orthogonality,
     DeviceMeshAxis,
     DTensorPropagationError,
     format_axis,
@@ -885,13 +883,15 @@ class _OpSpec:
         tensor_kwargs: Kwarg names that are tensor inputs.
         tensor_varargs_from: If set, all positional args from this index onward
             are tensor inputs (for ops with *args like einsum).
+        nonlinear_args: Positional arg indices whose operand enters the op
+            nonlinearly and therefore must not be P.
     """
 
     linearity: OpLinearity
     tensor_args: tuple[int, ...] = ()
     tensor_kwargs: tuple[str, ...] = ()
     tensor_varargs_from: int | None = None
-    fixed_args: tuple[int, ...] = ()
+    nonlinear_args: tuple[int, ...] = ()
 
 
 _OP_REGISTRY: dict[Callable, _OpSpec] = {
@@ -941,6 +941,8 @@ _OP_REGISTRY: dict[Callable, _OpSpec] = {
     torch.Tensor.mean: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.nansum: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.nanmean: _OpSpec(OpLinearity.LINEAR, (0,)),
+    torch.cumsum: _OpSpec(OpLinearity.LINEAR, (0,)),
+    torch.Tensor.cumsum: _OpSpec(OpLinearity.LINEAR, (0,)),
     # =================================================================
     # LINEAR -- concat / stack (tensor list at pos 0)
     # =================================================================
@@ -1028,31 +1030,68 @@ _OP_REGISTRY: dict[Callable, _OpSpec] = {
     torch.Tensor.repeat_interleave: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.Tensor.unfold: _OpSpec(OpLinearity.LINEAR, (0,)),
     # =================================================================
-    # LINEAR with fixed_args -- division (linear in numerator only)
+    # MULTILINEAR with nonlinear_args -- division (linear in numerator only)
+    # Division is multiplicative, not additive: a lone P numerator scaled by an
+    # R/scalar divisor stays P, while LINEAR would reject this as affine.
+    # nonlinear_args additionally forbids P in the divisor because 1/y is
+    # nonlinear in y.
     # =================================================================
-    torch.div: _OpSpec(OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)),
+    torch.div: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
+    ),
     torch.divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
     ),
     torch.true_divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
+    ),
+    torch.Tensor.div: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
-    torch.Tensor.div: _OpSpec(OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)),
     torch.Tensor.divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     torch.Tensor.true_divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
+    ),
+    torch.Tensor.div_: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
-    torch.Tensor.div_: _OpSpec(OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)),
     torch.Tensor.__truediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     torch.Tensor.__rtruediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(0,)
+        OpLinearity.MULTILINEAR, (0, 1), ("other",), nonlinear_args=(0,)
     ),
     torch.Tensor.__itruediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     # =================================================================
     # MULTILINEAR -- multiplication
@@ -1539,80 +1578,40 @@ def _set_result_type(result: object, output_type: LocalSpmdType) -> None:
             _apply(item)
 
 
-def _apply_fixed_args(  # noqa: C901
+def _validate_nonlinear_args(
     func: Callable,
     args: tuple,
-    kwargs: dict,
     spec: _OpSpec,
-    input_types_list: list[LocalSpmdType],
-) -> list[LocalSpmdType]:
-    """Filter fixed_args for LINEAR ops when Partial is present.
-
-    ``fixed_args`` lists positional arg indices that must be held fixed (not P)
-    for the op to be linear in the remaining args.  For example, ``div(a, b)``
-    is linear in ``a`` when ``b`` is fixed, so ``fixed_args=(1,)``.
-
-    When P is present among the non-fixed tensor args:
-    1. Validate that tensor args at fixed_args positions don't have P on any axis.
-    2. Exclude their types from the returned list so LINEAR inference sees only
-       the "free" args (and doesn't reject P + R mixing).
-
-    When no P is present in the free args, return the original list unchanged
-    so normal inference rules apply (e.g., R + V -> V).
+) -> None:
+    """Validate that operands entering the op nonlinearly are not Partial.
 
-    Args:
-        func: The torch function being called.
-        args: Positional arguments to the torch operation.
-        kwargs: Keyword arguments to the torch operation.
-        spec: The op specification (must have non-empty fixed_args).
-        input_types_list: The already-collected input types list.
+    Division is the current user: ``x / y`` is linear (scaling) in ``x`` only,
+    since ``1/y`` is nonlinear in ``y``. Nonlinear operands may be untyped
+    Python scalars or typed non-Partial values, but they cannot be P.
     """
-    # Identify which input_types_list entries came from fixed_args positions.
-    # We re-walk args to figure out which types are "fixed" vs "free".
-    fixed_positions = set(spec.fixed_args)
-    free_types: list[LocalSpmdType] = []
-    fixed_types: list[LocalSpmdType] = []
 
-    for i in spec.tensor_args:
-        if i < len(args):
-            val = args[i]
-            if isinstance(val, torch.Tensor):
-                if i in fixed_positions:
-                    fixed_types.append(get_local_type(val))
-                else:
-                    free_types.append(get_local_type(val))
-            elif isinstance(val, (list, tuple)):
-                for item in val:
-                    if isinstance(item, torch.Tensor):
-                        if i in fixed_positions:
-                            fixed_types.append(get_local_type(item))
-                        else:
-                            free_types.append(get_local_type(item))
-
-    # Check if P is present in any free arg on any axis
-    has_p_in_free = any(P in typ.values() for typ in free_types)
-    # Also check default V (missing from dict means V, not P) -- V != P, so fine.
-
-    if not has_p_in_free:
-        # No P in free args -- normal inference, include all types.
-        return input_types_list
-
-    # P is present in free args -- validate fixed args don't have P
-    for fixed_type in fixed_types:
-        for axis, typ in fixed_type.items():
-            if typ is P:
+    def _check_type(typ: LocalSpmdType) -> None:
+        for axis, axis_type in typ.items():
+            if axis_type is P:
                 raise SpmdTypeError(
-                    f"{func.__name__}: Partial type in fixed argument "
-                    f"(denominator/divisor) on axis {format_axis(axis)} is not allowed. "
-                    f"Division is only linear in the numerator."
+                    f"{func.__name__}: Partial type in nonlinear argument "
+                    f"on axis {format_axis(axis)} is not allowed. Reduce or "
+                    f"convert the Partial operand before using it in a "
+                    f"nonlinear position."
                 )
 
-    # Exclude fixed_types from input_types_list: return only free_types
-    # plus any scalar types that were appended by _collect_scalar_types.
-    # scalar_types count = len(input_types_list) - len(free_types) - len(fixed_types)
-    n_tensor_types = len(free_types) + len(fixed_types)
-    scalar_types = input_types_list[n_tensor_types:]
-    return free_types + scalar_types
+    def _check_value(val: object) -> None:
+        if isinstance(val, torch.Tensor):
+            _check_type(get_local_type(val))
+        elif isinstance(val, Scalar):
+            _check_type(get_local_type(val))
+        elif isinstance(val, (list, tuple)):
+            for item in val:
+                _check_value(item)
+
+    for i in spec.nonlinear_args:
+        if i < len(args):
+            _check_value(args[i])
 
 
 # Every function in this registry must accept (x: Tensor, axis, *, src=..., dst=...)
@@ -1680,7 +1679,7 @@ def _set_result_partition_spec(
             _set_partition_spec(item, spec)
 
 
-def _collect_shard_axes(
+def _collect_shard_axes(  # noqa: C901
     partition_specs: list[PartitionSpec | None],
     global_axes: set[MeshAxis],
 ) -> tuple[list[MeshAxis], dict[MeshAxis, set[MeshAxis]]]:
@@ -2086,11 +2085,16 @@ class _ShardPropagator:
         """
         mesh = self._get_mesh(axis.size())
 
+        from torch.utils._python_dispatch import _disable_current_modes
+
+        # Keep checker-internal fake DTensor propagation invisible to active
+        # TorchDispatchModes such as selective activation checkpointing's
+        # caching mode.
         # Disable LocalTensorMode during meta DTensor propagation.
         # DTensor.from_local calls mesh.get_coordinate() which, under
         # LocalTensorMode(N), tries to map all N simulated ranks into
         # the propagator's fake 1-D mesh.
-        with maybe_disable_local_tensor_mode():
+        with maybe_disable_local_tensor_mode(), _disable_current_modes():
             meta_args = torch.utils._pytree.tree_map(
                 lambda x: self._to_meta_dtensor(x, axis, mesh), args
             )
@@ -2902,10 +2906,8 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                     linearity = (
                         spec.linearity if spec is not None else OpLinearity.NONLINEAR
                     )
-                    if spec is not None and spec.fixed_args:
-                        input_types_list = _apply_fixed_args(
-                            func, args, kwargs, spec, input_types_list
-                        )
+                    if spec is not None and spec.nonlinear_args:
+                        _validate_nonlinear_args(func, original_args, spec)
                     # NB: output_type may be further updated because global
                     # shard propagation below may promote local type V to P.
                     output_type = infer_output_type(

{spmd_types-0.2.1 → spmd_types-0.2.2}/spmd_types/_collectives.py

@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import List, Optional, TYPE_CHECKING
+from typing import List, Optional
 
 import torch
 from spmd_types import _dist
@@ -18,13 +18,20 @@ from spmd_types._dtype_utils import (
     _process_dtype_options,
 )
 from spmd_types._local import convert, reinterpret
+from spmd_types._mesh import _pg_for_axis
 from spmd_types._traceback import api_boundary
-from spmd_types.types import _canonicalize_shard, I, P, PerMeshAxisSpmdType, R, Shard, V
+from spmd_types.types import (
+    _canonicalize_shard,
+    DeviceMeshAxis,
+    I,
+    P,
+    PerMeshAxisSpmdType,
+    R,
+    Shard,
+    V,
+)
 from torch.overrides import handle_torch_function, has_torch_function_unary
 
-if TYPE_CHECKING:
-    from torch.distributed import ProcessGroup
-
 # =============================================================================
 # all_reduce: P -> R | I
 # =============================================================================
@@ -93,7 +100,7 @@ class _AllReduce(torch.autograd.Function):
 @api_boundary
 def all_reduce(
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType = P,
     dst: PerMeshAxisSpmdType,
@@ -169,12 +176,13 @@ def all_reduce(
     common case), this role is filled by ``reduce_scatter(): P->V``
     instead.
     """
+    pg = _pg_for_axis(axis)
     if has_torch_function_unary(x):
         return handle_torch_function(
             all_reduce,
             (x,),
             x,
-            axis,
+            pg,
             src=src,
             dst=dst,
             inplace=inplace,
@@ -184,7 +192,7 @@ def all_reduce(
         )
     if src is not P:
         if src is V:
-            x = reinterpret(x, axis, src=V, dst=P)
+            x = reinterpret(x, pg, src=V, dst=P)
         elif src is R or src is I:
             raise ValueError(
                 f"all_reduce src must be P, got {src}. "
@@ -204,7 +212,7 @@ def all_reduce(
             requires_grad=x.requires_grad,
             inplace=inplace,
         )
-        return _AllReduce.apply(x, axis, dst, inplace, dtype_options)
+        return _AllReduce.apply(x, pg, dst, inplace, dtype_options)
     else:
         raise ValueError(f"all_reduce dst must be R or I, got {dst}")
 
@@ -387,7 +395,7 @@ class _AllGatherUneven(torch.autograd.Function):
 @api_boundary
 def all_gather(
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType = V,
     dst: PerMeshAxisSpmdType,
@@ -510,12 +518,13 @@ def all_gather(
     welcome to write your own custom collective and give it the same type as
     ``spmd.all_gather``.
     """
+    pg = _pg_for_axis(axis)
     if has_torch_function_unary(x):
         return handle_torch_function(
             all_gather,
             (x,),
             x,
-            axis,
+            pg,
             src=src,
             dst=dst,
             gather_dim=gather_dim,
@@ -560,16 +569,16 @@ def all_gather(
         if split_sizes is not None:
             return _AllGatherUneven.apply(
                 x,
-                axis,
+                pg,
                 dst,
                 gather_dim,
                 split_sizes,
                 dtype_options,
             )
         elif stack:
-            return _AllGatherStack.apply(x, axis, dst, gather_dim, dtype_options)
+            return _AllGatherStack.apply(x, pg, dst, gather_dim, dtype_options)
         else:
-            return _AllGatherShard.apply(x, axis, dst, gather_dim, dtype_options)
+            return _AllGatherShard.apply(x, pg, dst, gather_dim, dtype_options)
     else:
         raise ValueError(f"all_gather dst must be R or I, got {dst}")
 
@@ -698,7 +707,7 @@ class _ReduceScatterUneven(torch.autograd.Function):
 @api_boundary
 def reduce_scatter(
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType = P,
     dst: PerMeshAxisSpmdType = V,
@@ -774,12 +783,13 @@ def reduce_scatter(
     It is common to want to reduce-scatter on varying data; just
     ``reinterpret(V,P)`` the data as partial before calling ``reduce_scatter``.
     """
+    pg = _pg_for_axis(axis)
     if has_torch_function_unary(x):
         return handle_torch_function(
             reduce_scatter,
             (x,),
             x,
-            axis,
+            pg,
             src=src,
             dst=dst,
             scatter_dim=scatter_dim,
@@ -793,7 +803,7 @@ def reduce_scatter(
 
     if src is not P:
         if src is V:
-            x = reinterpret(x, axis, src=V, dst=P)
+            x = reinterpret(x, pg, src=V, dst=P)
         elif src is R or src is I:
             raise ValueError(
                 f"reduce_scatter src must be P, got {src}. "
@@ -832,17 +842,17 @@ def reduce_scatter(
     )
 
     if dst is V:
-        return _ReduceScatterStack.apply(x, axis, scatter_dim, dtype_options)
+        return _ReduceScatterStack.apply(x, pg, scatter_dim, dtype_options)
     elif split_sizes is not None:
         return _ReduceScatterUneven.apply(
             x,
-            axis,
+            pg,
             scatter_dim,
             split_sizes,
             dtype_options,
         )
     else:
-        return _ReduceScatterShard.apply(x, axis, scatter_dim, dtype_options)
+        return _ReduceScatterShard.apply(x, pg, scatter_dim, dtype_options)
 
 
 # =============================================================================
@@ -971,7 +981,7 @@ class _AllToAllShard(torch.autograd.Function):
 @api_boundary
 def all_to_all(
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType = V,
     dst: PerMeshAxisSpmdType = V,
@@ -1059,12 +1069,13 @@ def all_to_all(
     ``all_to_all(S(i),S(j)): S(i) -> S(j)``, the backwards is
     ``all_to_all(S(j),S(i)): S(j) -> S(i)``.
     """
+    pg = _pg_for_axis(axis)
     if has_torch_function_unary(x):
         return handle_torch_function(
             all_to_all,
             (x,),
             x,
-            axis,
+            pg,
             src=src,
             dst=dst,
             split_dim=split_dim,
@@ -1142,15 +1153,15 @@ def all_to_all(
             )
         return _AllToAllUneven.apply(
             x,
-            axis,
+            pg,
             output_split_sizes,
             input_split_sizes,
             dtype_options,
         )
     elif src is V and dst is V:
-        return _AllToAllStack.apply(x, axis, split_dim, concat_dim, dtype_options)
+        return _AllToAllStack.apply(x, pg, split_dim, concat_dim, dtype_options)
     else:
-        return _AllToAllShard.apply(x, axis, split_dim, concat_dim, dtype_options)
+        return _AllToAllShard.apply(x, pg, split_dim, concat_dim, dtype_options)
 
 
 # =============================================================================
@@ -1161,7 +1172,7 @@ def all_to_all(
 @api_boundary
 def redistribute(  # noqa: C901
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType,
     dst: PerMeshAxisSpmdType,
@@ -1383,7 +1394,7 @@ def redistribute(  # noqa: C901
 
 def unshard(
     x,
-    axis: ProcessGroup,
+    axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType,
     dst: PerMeshAxisSpmdType,