spmd_types 0.3.1__tar.gz → 0.3.2__tar.gz

{spmd_types-0.3.1/spmd_types.egg-info → spmd_types-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.3.1
+Version: 0.3.2
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License

spmd_types-0.3.2/pyproject.toml

@@ -0,0 +1,110 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spmd_types"
+version = "0.3.2"
+description = "A type system for distributed (SPMD) tensor computations in PyTorch"
+readme = "README.md"
+license = {file = "LICENSE"}
+requires-python = ">=3.10"
+authors = [
+    {name = "Meta Platforms, Inc."},
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+]
+
+[project.optional-dependencies]
+docs = [
+    "myst-parser==2.0.0",
+    "pytorch-sphinx-theme2==0.4.13",
+    "sphinx==7.2.6",
+]
+
+[project.urls]
+Homepage = "https://github.com/meta-pytorch/spmd_types"
+Documentation = "https://meta-pytorch.org/spmd_types/"
+Repository = "https://github.com/meta-pytorch/spmd_types"
+Issues = "https://github.com/meta-pytorch/spmd_types/issues"
+
+[tool.setuptools.packages.find]
+include = ["spmd_types*"]
+
+[tool.setuptools.package-data]
+spmd_types = ["py.typed"]
+
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
+
+[tool.ruff.format]
+exclude = ["*.md"]
+quote-style = "double"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["*_test.py"]
+
+# ops imports spmd_types from most of its interfaces and farm from its
+# gradient code, so an unscoped import graph selects their suites, GPU plans
+# included, for every change here. Import-graph selection stays inside this
+# tree; the tests below opt back in.
+[tool.test-selector.graph-dependents]
+"**" = ["**"]
+
+[tool.test-selector.extra-deps]
+
+# The farm tests that import spmd_types directly.
+"**" = [
+    "../farm/farm/model/tests/test_input_side_fsdp_order.py",
+    "../farm/farm/optimize/tests/test_grad_clipping_shared_tensor.py",
+    "../farm/farm/optimize/tests/test_replication_metadata.py",
+]
+
+# BEGIN GENERATED buck-manual-deps -- do not edit by hand
+# Regenerate: export_buck_manual_deps.py
+"spmd_types/__init__.py" = [
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_artifacts.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_artifacts_v22.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_encoder_r6_i11_v22.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_conditioner_calibration_v21.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_ctts_pipeline_r6_i11_v22.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_recipe_r5_v21.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r5.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r6_i11_v22.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r6_i11_v22_gb300.py",
+    "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_style_r5_v21.py",
+    "../ops/interfaces/attn/triton_attention_fp32.py",
+    "../ops/interfaces/gemm/gemm.py",
+    "../rl/projects/rvg/data/data_pipeline/kimi_labeling.py",
+]
+# END GENERATED buck-manual-deps
+
+[dependency-groups]
+# Build backend of this package, locked: in the mslsrc/tbd_training export the
+# members build against the venv (the root's `no-build-isolation-package`), so
+# the backend has to be in a default group for `uv sync` to install it first.
+build = ["setuptools", "wheel"]
+# What the tests import beyond the package itself; the export runs them in
+# this project's uv env.
+# @tbd-training-enable: dev = ["expecttest", "pytest", "torch"]
+
+[tool.uv]
+default-groups = ["build"] # @tbd-training-disable
+# @tbd-training-enable: default-groups = ["build", "dev"]

{spmd_types-0.3.1 → spmd_types-0.3.2}/spmd_types/_collectives.py

@@ -68,11 +68,12 @@ class _AllReduce(torch.autograd.Function):
         ctx.backward_options = dtype_options.backward_options
         if inplace:
             assert dtype_options.op_dtype == x.dtype
+            # c10d reduces x's own storage, which must be contiguous.
+            # funcol.all_reduce copies strided inputs itself.
+            assert x.is_contiguous(), "inplace all_reduce input must be contiguous"
         else:
             x = _apply_op_dtype(x, dtype_options.op_dtype)
         pg = axis
-        # TODO: check if contiguous assertion is really necessary
-        assert x.is_contiguous(), "all_reduce input must be contiguous"
         # TODO: check if world_size == 1 short-circuit is really necessary
         if _dist.dist.get_world_size(pg) == 1:
             if inplace:
@@ -241,7 +242,7 @@ def all_reduce(
             input_dtype=x.dtype,
             requires_grad=x.requires_grad,
             async_op=async_op,
-            input_is_leaf=x.requires_grad and x.grad_fn is None,
+            input_is_leaf=x.requires_grad and x.is_leaf,
             inplace=inplace,
         )
         return _AllReduce.apply(x, pg, dst, inplace, dtype_options)
@@ -603,7 +604,7 @@ def all_gather(
             input_dtype=x.dtype,
             requires_grad=x.requires_grad,
             async_op=async_op,
-            input_is_leaf=x.requires_grad and x.grad_fn is None,
+            input_is_leaf=x.requires_grad and x.is_leaf,
         )
         if split_sizes is not None:
             return _AllGatherUneven.apply(
@@ -640,7 +641,10 @@ class _ReduceScatterStack(torch.autograd.Function):
         # NCCL reduce_scatter_tensor always scatters along dim 0, so move
         # scatter_dim into position 0 when it isn't already there.
         if scatter_dim != 0:
-            x = x.movedim(scatter_dim, 0).contiguous()
+            x = x.movedim(scatter_dim, 0)
+        # NCCL reduce_scatter_tensor reads the input as contiguous without
+        # checking; incoming gradients (e.g. from sum()) can be expanded views.
+        x = x.contiguous()
         result = x.new_empty([1] + list(x.shape[1:]))
         _dist.dist.reduce_scatter_tensor(
             result, x, op=_dist.dist.ReduceOp.SUM, group=pg
@@ -882,7 +886,7 @@ def reduce_scatter(
         input_dtype=x.dtype,
         requires_grad=x.requires_grad,
         async_op=async_op,
-        input_is_leaf=x.requires_grad and x.grad_fn is None,
+        input_is_leaf=x.requires_grad and x.is_leaf,
     )
 
     if dst is V:
@@ -915,7 +919,9 @@ class _AllToAllStack(torch.autograd.Function):
         ctx.backward_options = dtype_options.backward_options
         x = _apply_op_dtype(x, dtype_options.op_dtype)
         pg = axis
-        input_chunks = list(torch.unbind(x, dim=split_dim))
+        # Process groups require contiguous chunks; unbind along a non-zero
+        # split_dim returns strided views.
+        input_chunks = [c.contiguous() for c in torch.unbind(x, dim=split_dim)]
         output_chunks = [
             torch.empty_like(input_chunks[0]) for _ in range(len(input_chunks))
         ]
@@ -1045,6 +1051,7 @@ def all_to_all(
     ::
 
         def all_to_all_spec(xs, src, dst):
+            n = mesh_axis_size
             match src, dst:
                 case V, V:
                     x = torch.stack(xs)
@@ -1052,8 +1059,13 @@ def all_to_all(
                     return x.unbind()
 
                 case S(i), S(j):
-                    x = torch.concat(xs, i)
-                    return x.chunk(mesh_axis_size, j)
+                    # Every rank splits its tensor into n chunks on dim j and
+                    # sends chunk r to rank r; rank r concatenates the chunks
+                    # it receives on dim i, in source rank order.
+                    return [
+                        torch.concat([x.chunk(n, j)[r] for x in xs], i)
+                        for r in range(n)
+                    ]
 
     The varying and shard versions of ``all_to_all`` are pretty different (even
     though under the hood they both have an all-to-all communication pattern),
@@ -1062,8 +1074,33 @@ def all_to_all(
     * ``all_to_all(V,V)`` transposes the logical mesh axis with the dim 0 local
       tensor axis.
 
-    * ``all_to_all(S(i),S(j))`` intuitively unshards the tensor on dim i, and
-      then reshards it on dim j (but skips actually doing the all-gather).
+    * ``all_to_all(S(i),S(j))`` with ``i != j`` intuitively unshards the tensor
+      on dim i, and then reshards it on dim j (but skips actually doing the
+      all-gather).  The global tensor is preserved:
+      ``torch.concat(ys, j) == torch.concat(xs, i)``.
+
+    * ``all_to_all(S(i),S(i))`` does the same exchange with split and concat on
+      the same dim.  The output is still ``S(i)``, but the global tensor is
+      block-permuted: view global dim i as ``n * n`` equal blocks indexed by
+      (source rank, destination rank); the exchange swaps the two indices.
+      With ``x = torch.concat(xs, i)`` and ``y = torch.concat(ys, i)``::
+
+          y = x.unflatten(i, (n, n, -1)).transpose(i, i + 1).flatten(i, i + 2)
+
+      For example, if two ranks hold ``[0, 1, 2, 3]`` and ``[4, 5, 6, 7]``,
+      they end up with ``[0, 1, 4, 5]`` and ``[2, 3, 6, 7]``: the global tensor
+      goes from ``[0, 1, 2, 3, 4, 5, 6, 7]`` to ``[0, 1, 4, 5, 2, 3, 6, 7]``.
+      This is a valid global SPMD operation.  ``S(i)`` says how the global
+      tensor is reconstructed from the local tensors; it does not promise that
+      a collective leaves the global value unchanged.  Each explicit collective
+      defines its own operation on the global tensor, which may depend on the
+      size of the mesh axis.  Use ``redistribute`` to change the type while
+      preserving the global value; ``redistribute(S(i),S(i))`` is a no-op.
+
+    With ``input_split_sizes``/``output_split_sizes`` (all-to-all-v), the
+    exchange is the same on dim 0 except that rank s sends its
+    ``input_split_sizes[r]`` rows to rank r, so in general it moves blocks of
+    the global tensor rather than preserving it.
 
     ::
 
@@ -1122,7 +1159,8 @@ def all_to_all(
         [C0, C1, C2]      [A2, B2, C2]
 
     ``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)``.
+    ``all_to_all(S(j),S(i)): S(j) -> S(i)``.  When ``i == j`` this is the same
+    block transpose, which is its own inverse.
     """
     pg = _pg_for_axis(axis)
     if has_torch_function_unary(x):
@@ -1154,7 +1192,7 @@ def all_to_all(
 
     if isinstance(src, Shard) and isinstance(dst, Shard):
         # S(i) -> S(j): physically split on new shard dim (j), concat on old (i).
-        # Conceptual spec: concat(xs, i) then chunk(ws, j).
+        # This preserves the global tensor only when i != j; see the docstring.
         if split_dim is not None and split_dim != dst.dim:
             raise ValueError(
                 f"all_to_all S({src.dim})->S({dst.dim}) got split_dim={split_dim} "
@@ -1200,7 +1238,7 @@ def all_to_all(
         input_dtype=x.dtype,
         requires_grad=x.requires_grad,
         async_op=async_op,
-        input_is_leaf=x.requires_grad and x.grad_fn is None,
+        input_is_leaf=x.requires_grad and x.is_leaf,
     )
 
     if output_split_sizes is not None or input_split_sizes is not None:
@@ -1250,7 +1288,11 @@ def redistribute(  # noqa: C901
         redistribute(P,R)       =   all_reduce(P,R)
         redistribute(P,I)       =   all_reduce(P,I)
         redistribute(P,S(i))    =   reduce_scatter(P,S(i))
-        redistribute(S(i),S(j)) =   all_to_all(S(i),S(j))
+        redistribute(S(i),S(j)) =   all_to_all(S(i),S(j))     (i != j)
+        redistribute(S(i),S(i)) =   x                         (no-op)
+
+    ``redistribute(S(i),S(i))`` does not route to ``all_to_all(S(i),S(i))``,
+    which block-permutes the global tensor (see ``all_to_all``).
 
     These only work if the mesh axis is the LAST to shard a particular tensor
     dimension.

{spmd_types-0.3.1 → spmd_types-0.3.2}/spmd_types/_local.py

@@ -62,7 +62,8 @@ def _build_rank_to_group_local(pg) -> dict[int, int]:
       - Groups are [0,1] and [2,3]
       - Returns {0: 0, 1: 1, 2: 0, 3: 1}
     """
-    ranks, group_offsets, _offset = _prepare_collective_groups(pg)
+    # Older torch also returns the group's base offset as a third element.
+    ranks, group_offsets = _prepare_collective_groups(pg)[:2]
     mapping = {}
     for group_offset in group_offsets:
         for local_rank, r in enumerate(ranks):
@@ -402,7 +403,7 @@ def reinterpret(  # noqa: C901
         input_dtype=x.dtype,
         requires_grad=x.requires_grad,
         async_op=async_op,
-        input_is_leaf=x.requires_grad and x.grad_fn is None,
+        input_is_leaf=x.requires_grad and x.is_leaf,
     )
 
     if src is R and dst is V:
@@ -1099,7 +1100,7 @@ def convert(  # noqa: C901
         input_dtype=x.dtype,
         requires_grad=x.requires_grad,
         async_op=async_op,
-        input_is_leaf=x.requires_grad and x.grad_fn is None,
+        input_is_leaf=x.requires_grad and x.is_leaf,
     )
 
     if src_base is R and dst_base is V:
@@ -1184,6 +1185,7 @@ def shard(
     )
 
 
+@api_boundary
 def invariant_to_replicate(
     x,
     axis: DeviceMeshAxis,
@@ -1245,6 +1247,6 @@ def invariant_to_replicate(
         input_dtype=x.dtype,
         requires_grad=x.requires_grad,
         async_op=async_op,
-        input_is_leaf=x.requires_grad and x.grad_fn is None,
+        input_is_leaf=x.requires_grad and x.is_leaf,
     )
     return _InvariantToReplicate.apply(x, axis, dtype_options)

{spmd_types-0.3.1 → spmd_types-0.3.2}/spmd_types/rules.py

@@ -851,7 +851,10 @@ def all_to_all(
     dst: PerMeshAxisSpmdType = V,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor | NdimWithSpmdType:
-    """Type of ``spmd.all_to_all(x, axis, src=, dst=)``: ``V|S(i) -> V|S(j)``."""
+    """Type of ``spmd.all_to_all(x, axis, src=, dst=)``: ``V|S(i) -> V|S(j)``.
+
+    ``i == j`` is accepted, as at runtime; see ``spmd.all_to_all`` for its
+    global semantics (a block permutation of the global tensor)."""
     return _transition(x, "all_to_all", axis, src, dst, [("V|S", "V|S")], out=out)
 
 

{spmd_types-0.3.1 → spmd_types-0.3.2}/spmd_types/types.py

@@ -70,7 +70,9 @@ class Shard:
     specifying how ranks relate to a global tensor, Shard(dim) additionally
     says "the global tensor can be reconstructed by concatenating the local
     tensors along dimension ``dim``."  This is analogous to DTensor's
-    ``Shard(dim)`` placement.
+    ``Shard(dim)`` placement.  It describes how to reconstruct the global
+    tensor, not that an operation from ``S(i)`` to ``S(i)`` preserves its value
+    (e.g., ``all_to_all(S(i), S(i))`` block-permutes the global tensor).
 
     This is not a true type (notice it doesn't inherit from PerMeshAxisLocalSpmdType)
     but it is accepted at any src/dst argument and, from a typing perspective,

{spmd_types-0.3.1 → spmd_types-0.3.2/spmd_types.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.3.1
+Version: 0.3.2
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License

spmd_types-0.3.1/pyproject.toml

@@ -1,69 +0,0 @@
-[build-system]
-requires = ["setuptools>=68.0", "wheel"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "spmd_types"
-version = "0.3.1"
-description = "A type system for distributed (SPMD) tensor computations in PyTorch"
-readme = "README.md"
-license = {file = "LICENSE"}
-requires-python = ">=3.10"
-authors = [
-    {name = "Meta Platforms, Inc."},
-]
-classifiers = [
-    "Development Status :: 4 - Beta",
-    "Intended Audience :: Developers",
-    "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: BSD License",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
-    "Programming Language :: Python :: 3.12",
-    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-]
-dependencies = [
-]
-
-[project.optional-dependencies]
-docs = [
-    "myst-parser==2.0.0",
-    "pytorch-sphinx-theme2==0.4.13",
-    "sphinx==7.2.6",
-]
-
-[project.urls]
-Homepage = "https://github.com/meta-pytorch/spmd_types"
-Documentation = "https://meta-pytorch.org/spmd_types/"
-Repository = "https://github.com/meta-pytorch/spmd_types"
-Issues = "https://github.com/meta-pytorch/spmd_types/issues"
-
-[tool.setuptools.packages.find]
-include = ["spmd_types*"]
-
-[tool.setuptools.package-data]
-spmd_types = ["py.typed"]
-
-
-[tool.ruff]
-line-length = 88
-target-version = "py310"
-
-[tool.ruff.lint]
-select = ["E4", "E7", "E9", "F"]
-
-[tool.ruff.format]
-exclude = ["*.md"]
-quote-style = "double"
-
-[tool.pytest.ini_options]
-testpaths = ["tests"]
-python_files = ["*_test.py"]
-
-[tool.test-selector.extra-deps]
-
-# BEGIN GENERATED buck-manual-deps -- do not edit by hand
-# Regenerate: export_buck_manual_deps.py
-"spmd_types/__init__.py" = ["../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_artifacts.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_artifacts_v22.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_audio_encoder_r6_i11_v22.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_conditioner_calibration_v21.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_ctts_pipeline_r6_i11_v22.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_recipe_r5_v21.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r5.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r6_i11_v22.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_sampling_r6_i11_v22_gb300.py", "../../llama4x/projects/experimental/rtai_avatar/pumpkin/sythetic_data/t2av/conversation_landscape_style_r5_v21.py", "../ops/interfaces/attn/triton_attention_fp32.py", "../ops/interfaces/gemm/gemm.py"]
-# END GENERATED buck-manual-deps