spmd_types 0.2.0__tar.gz → 0.2.2__tar.gz

{spmd_types-0.2.0/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.0
+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.0 → spmd_types-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.0"
+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.0 → spmd_types-0.2.2}/spmd_types/__init__.py

@@ -1,14 +1,14 @@
+# 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
 
-import spmd_types._dtensor_checker as _dtensor_checker  # noqa: F401
-
-# These still require _checker (the type inference engine)
-from spmd_types._checker import (  # noqa: F401
-    _SpmdTypeBackwardCompatibleMode as SpmdTypeMode,
-    is_type_checking,
-    no_typecheck,
-    typecheck,
+from spmd_types._backward_hooks import (  # noqa: F401
+    register_local_backward_hook,
 )
 from spmd_types._collectives import (  # noqa: F401
     all_gather,
@@ -36,7 +36,13 @@ 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 current_mesh  # 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,
@@ -47,9 +53,15 @@ 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,
+    local,
+    local_map,
     mutate_type,
     register_autograd_function,
+    register_decomposition,
     register_local_autograd_function,
     trace,
 )
@@ -62,6 +74,7 @@ from spmd_types.types import (  # noqa: F401
     LocalSpmdType,
     normalize_axis,
     normalize_mesh,
+    normalize_partition_spec,
     P,
     Partial,
     PartitionSpec,
@@ -72,6 +85,7 @@ from spmd_types.types import (  # noqa: F401
     Replicate,
     S,
     Shard,
+    SpmdType,
     SpmdTypeError,
     TensorSharding,
     V,

spmd_types-0.2.2/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