spmd_types 0.2.0__tar.gz → 0.2.1__tar.gz

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

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

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.0"
+version = "0.2.1"
 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.1}/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,
@@ -48,8 +54,13 @@ from spmd_types._type_attr import (  # noqa: F401
 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,
 )
@@ -62,6 +73,7 @@ from spmd_types.types import (  # noqa: F401
     LocalSpmdType,
     normalize_axis,
     normalize_mesh,
+    normalize_partition_spec,
     P,
     Partial,
     PartitionSpec,

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