spmd-types 0.2.1__py3-none-any.whl

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/_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