spmd-types 0.2.1__tar.gz

spmd_types-0.2.1/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+(c) Meta Platforms, Inc. and affiliates.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list
+of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this
+list of conditions and the following disclaimer in the documentation
+and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may
+be used to endorse or promote products derived from this software without specific
+prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
+EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
+SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
+TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+DAMAGE.

spmd_types-0.2.1/PKG-INFO

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: spmd_types
+Version: 0.2.1
+Summary: A type system for distributed (SPMD) tensor computations in PyTorch
+Author: Meta Platforms, Inc.
+License: BSD 3-Clause License
+        
+        (c) Meta Platforms, Inc. and affiliates.
+        
+        Redistribution and use in source and binary forms, with or without modification,
+        are permitted provided that the following conditions are met:
+        
+        1. Redistributions of source code must retain the above copyright notice, this list
+        of conditions and the following disclaimer.
+        
+        2. Redistributions in binary form must reproduce the above copyright notice, this
+        list of conditions and the following disclaimer in the documentation
+        and/or other materials provided with the distribution.
+        
+        3. Neither the name of the copyright holder nor the names of its contributors may
+        be used to endorse or promote products derived from this software without specific
+        prior written permission.
+        
+        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
+        EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+        OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
+        SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+        INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
+        TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+        BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+        CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+        ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+        DAMAGE.
+        
+Project-URL: Homepage, https://github.com/meta-pytorch/spmd_types
+Project-URL: Repository, https://github.com/meta-pytorch/spmd_types
+Project-URL: Issues, https://github.com/meta-pytorch/spmd_types/issues
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.10.0
+Dynamic: license-file
+
+# spmd_types
+
+A type system for distributed (SPMD) tensor computations in PyTorch.
+
+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.
+
+## Installation
+
+```bash
+pip install spmd_types
+```
+
+## Quick start
+
+```python
+import torch
+import torch.distributed as dist
+import spmd_types as spmd
+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.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 [MEGATRON_QUICKSTART.md](MEGATRON_QUICKSTART.md) for a guide on porting
+Megatron-derived training frameworks to use spmd_types.
+
+See [DESIGN.md](DESIGN.md) for the full type system specification, including
+type inference rules, collective signatures, and forward-backward pairs.
+
+## License
+
+BSD 3-Clause License. See [LICENSE](LICENSE) for details.

spmd_types-0.2.1/README.md

@@ -0,0 +1,51 @@
+# spmd_types
+
+A type system for distributed (SPMD) tensor computations in PyTorch.
+
+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.
+
+## Installation
+
+```bash
+pip install spmd_types
+```
+
+## Quick start
+
+```python
+import torch
+import torch.distributed as dist
+import spmd_types as spmd
+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.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 [MEGATRON_QUICKSTART.md](MEGATRON_QUICKSTART.md) for a guide on porting
+Megatron-derived training frameworks to use spmd_types.
+
+See [DESIGN.md](DESIGN.md) for the full type system specification, including
+type inference rules, collective signatures, and forward-backward pairs.
+
+## License
+
+BSD 3-Clause License. See [LICENSE](LICENSE) for details.

spmd_types-0.2.1/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spmd_types"
+version = "0.2.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 = [
+    "torch>=2.10.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/meta-pytorch/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.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["*_test.py"]

spmd_types-0.2.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spmd_types-0.2.1/spmd_types/__init__.py

@@ -0,0 +1,109 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# spmd_types package
+from __future__ import annotations
+
+from spmd_types._backward_hooks import (  # noqa: F401
+    register_local_backward_hook,
+)
+from spmd_types._collectives import (  # noqa: F401
+    all_gather,
+    all_reduce,
+    all_to_all,
+    redistribute,
+    reduce_scatter,
+    unshard,
+)
+from spmd_types._dist import set_dist  # noqa: F401
+from spmd_types._dtensor import (  # noqa: F401
+    dtensor_placement_to_spmd_type,
+    spmd_redistribute,
+    spmd_type_to_dtensor_placement,
+)
+from spmd_types._local import (  # noqa: F401
+    convert,
+    invariant_to_replicate,
+    reinterpret,
+    shard,
+)
+from spmd_types._mesh import set_current_mesh  # noqa: F401
+from spmd_types._mesh_axis import MeshAxis  # noqa: F401
+
+# reinterpret_mesh lives in its own module
+from spmd_types._reinterpret_mesh import reinterpret_mesh  # noqa: F401
+from spmd_types._scalar import Scalar  # noqa: F401
+from spmd_types._state import (  # noqa: F401
+    current_mesh,
+    current_mesh_all_names,
+    current_mesh_names,
+    is_type_checking,
+    no_typecheck,
+)
+from spmd_types._traceback import traceback_filtering  # noqa: F401
+from spmd_types._type_attr import (  # noqa: F401
+    get_axis_local_type,
+    get_local_type,
+    maybe_get_axis_local_type,
+)
+
+# Collectives and operations -- runtime API (no _checker dependency)
+from spmd_types.runtime import (  # noqa: F401
+    assert_local_type,
+    assert_type,
+    assert_type_like,
+    Infer,
+    local,
+    local_map,
+    mutate_type,
+    register_autograd_function,
+    register_decomposition,
+    register_local_autograd_function,
+    trace,
+)
+
+# Types
+from spmd_types.types import (  # noqa: F401
+    DimSharding,
+    I,
+    Invariant,
+    LocalSpmdType,
+    normalize_axis,
+    normalize_mesh,
+    normalize_partition_spec,
+    P,
+    Partial,
+    PartitionSpec,
+    PerMeshAxisLocalSpmdType,
+    PerMeshAxisSpmdType,
+    PerMeshAxisSpmdTypes,
+    R,
+    Replicate,
+    S,
+    Shard,
+    SpmdTypeError,
+    TensorSharding,
+    V,
+    Varying,
+)
+
+
+class _TypeCheckingSentinel:
+    """Singleton whose bool value reflects whether type checking is active.
+
+    ``bool(TYPE_CHECKING)`` returns True when a ``typecheck()`` context is
+    active on the current thread, False otherwise.  This avoids the
+    sys.modules replacement trick which breaks torch.compile / Dynamo.
+    """
+
+    def __bool__(self) -> bool:
+        return is_type_checking()
+
+    def __repr__(self) -> str:
+        return f"TYPE_CHECKING({is_type_checking()})"
+
+
+TYPE_CHECKING = _TypeCheckingSentinel()

spmd_types-0.2.1/spmd_types/_backward_hooks.py

@@ -0,0 +1,133 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Per-hook SPMD type-propagation registry for nn.Module backward hooks.
+
+BackwardHookFunction (nn.Module.register_full_backward_hook) has a pure
+pass-through forward. Each user hook callable must be explicitly registered
+via ``register_local_backward_hook`` to declare that it does not alter SPMD
+types; unregistered hooks raise SpmdTypeError when type checking is active.
+
+TODO: a per-hook rule API (for hooks that do collectives on grads, e.g.
+all-reduce flipping I -> R) can be added later when a real use case
+motivates the exact shape.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.utils.hooks as _torch_hooks
+from spmd_types._state import is_type_checking
+from spmd_types.runtime import (
+    _set_partition_spec,
+    assert_type_like,
+    get_partition_spec,
+)
+from spmd_types.types import SpmdTypeError
+from torch.nn.modules._functions import BackwardHookFunction
+
+_LOCAL_BACKWARD_HOOKS: set[Callable] = set()
+
+
+def register_local_backward_hook(fn: Callable) -> Callable:
+    """Declare that ``fn`` does not alter SPMD types when it runs in backward.
+
+    Analogous to :func:`register_local_autograd_function`: the hook's backward
+    is treated as local (no collectives, no type-changing effects on grads).
+    Side-agnostic: covers both ``register_full_backward_hook`` (post) and
+    ``register_full_backward_pre_hook`` (pre) use.  Usable as a decorator.
+    """
+    _LOCAL_BACKWARD_HOOKS.add(fn)
+    return fn
+
+
+def _validate(hooks):
+    if not is_type_checking():
+        return
+    for fn in hooks:
+        if fn in _LOCAL_BACKWARD_HOOKS:
+            continue
+        raise SpmdTypeError(
+            f"Backward hook {getattr(fn, '__qualname__', fn)!r} attached to "
+            f"nn.Module is not registered for SPMD type propagation.  Call "
+            f"register_local_backward_hook(fn) if the hook does not alter "
+            f"gradient SPMD types."
+        )
+
+
+def _apply_types(hooks, inputs, outputs):
+    """Copy SPMD annotations across BackwardHookFunction wrappers.
+
+    BackwardHookFunction.apply creates fresh tensor objects for module
+    backward hooks. Registered local hooks are semantically pass-through, so
+    their wrapper tensors should keep the input SPMD metadata.
+    """
+    if not isinstance(inputs, tuple):
+        inputs = (inputs,)
+    if not isinstance(outputs, tuple):
+        outputs = (outputs,)
+
+    for fn in hooks:
+        if fn in _LOCAL_BACKWARD_HOOKS:
+            for inp, out in zip(inputs, outputs):
+                if isinstance(inp, torch.Tensor) and isinstance(out, torch.Tensor):
+                    spec = get_partition_spec(inp)
+                    assert_type_like(out, inp)
+                    if spec is not None:
+                        _set_partition_spec(out, spec)
+
+
+_orig_setup_input_hook = None
+_orig_setup_output_hook = None
+
+
+def _patched_setup_input_hook(self, input):
+    _validate(self.user_hooks)
+    result = _orig_setup_input_hook(self, input)
+    _apply_types(self.user_hooks, input, result)
+    _apply_types(self.user_pre_hooks, input, result)
+    return result
+
+
+def _patched_setup_output_hook(self, output):
+    _validate(self.user_pre_hooks)
+    result = _orig_setup_output_hook(self, output)
+    _apply_types(self.user_pre_hooks, output, result)
+    _apply_types(self.user_hooks, output, result)
+    return result
+
+
+def install() -> None:
+    """Install the BackwardHook monkey-patch.  Idempotent."""
+    global _orig_setup_input_hook, _orig_setup_output_hook
+
+    if _orig_setup_input_hook is not None:
+        return
+
+    from spmd_types._checker import register_autograd_function
+
+    BackwardHookFunction.typecheck_forward = staticmethod(BackwardHookFunction.apply)
+    register_autograd_function(BackwardHookFunction)
+
+    _orig_setup_input_hook = _torch_hooks.BackwardHook.setup_input_hook
+    _orig_setup_output_hook = _torch_hooks.BackwardHook.setup_output_hook
+    _torch_hooks.BackwardHook.setup_input_hook = _patched_setup_input_hook
+    _torch_hooks.BackwardHook.setup_output_hook = _patched_setup_output_hook
+
+
+def uninstall() -> None:
+    """Remove the BackwardHook monkey-patch.  Idempotent."""
+    global _orig_setup_input_hook, _orig_setup_output_hook
+
+    if _orig_setup_input_hook is None:
+        return
+
+    _torch_hooks.BackwardHook.setup_input_hook = _orig_setup_input_hook
+    _torch_hooks.BackwardHook.setup_output_hook = _orig_setup_output_hook
+    _orig_setup_input_hook = None
+    _orig_setup_output_hook = None