spmd_types 0.2.4__tar.gz → 0.3.0__tar.gz

{spmd_types-0.2.4/spmd_types.egg-info → spmd_types-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.4
+Version: 0.3.0
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License
@@ -33,6 +33,7 @@ License: BSD 3-Clause License
         DAMAGE.
         
 Project-URL: Homepage, https://github.com/meta-pytorch/spmd_types
+Project-URL: Documentation, https://meta-pytorch.org/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
@@ -47,6 +48,10 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Provides-Extra: docs
+Requires-Dist: myst-parser==2.0.0; extra == "docs"
+Requires-Dist: pytorch-sphinx-theme2==0.4.13; extra == "docs"
+Requires-Dist: sphinx==7.2.6; extra == "docs"
 Dynamic: license-file
 
 # `spmd_types`
@@ -107,10 +112,17 @@ dist.destroy_process_group()
 
 ## Documentation
 
+See the [SPMD Types documentation](https://meta-pytorch.org/spmd_types/) for
+the complete user and API 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 [Rules](docs/rules.md) for how to write an `spmd_typecheck` typing rule
+for your own custom autograd functions by composing `spmd_types.rules`, and
+how to test it numerically with `rulecheck`.
+
 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.

{spmd_types-0.2.4 → spmd_types-0.3.0}/README.md

@@ -56,10 +56,17 @@ dist.destroy_process_group()
 
 ## Documentation
 
+See the [SPMD Types documentation](https://meta-pytorch.org/spmd_types/) for
+the complete user and API 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 [Rules](docs/rules.md) for how to write an `spmd_typecheck` typing rule
+for your own custom autograd functions by composing `spmd_types.rules`, and
+how to test it numerically with `rulecheck`.
+
 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.

spmd_types-0.3.0/pyproject.toml

@@ -0,0 +1,69 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spmd_types"
+version = "0.3.0"
+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

{spmd_types-0.2.4 → spmd_types-0.3.0}/spmd_types/__init__.py

@@ -7,6 +7,7 @@
 # spmd_types package
 from __future__ import annotations
 
+from spmd_types import rules  # noqa: F401
 from spmd_types._collectives import (  # noqa: F401
     all_gather,
     all_reduce,
@@ -65,6 +66,7 @@ from spmd_types.runtime import (  # noqa: F401
     no_typecheck,
     register_autograd_function,
     register_decomposition,
+    reinterpret_no_grad,
     trace,
 )
 

{spmd_types-0.2.4 → spmd_types-0.3.0}/spmd_types/_backward_hooks.py

@@ -97,10 +97,22 @@ def install() -> None:
     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)
+    def _spmd_typecheck(outputs, *, args):
+        # BackwardHookFunction is the identity on its tensor arguments: each
+        # output carries its input's type and sharding (rules imports runtime,
+        # so import lazily here).
+        from spmd_types import rules
+
+        if len(outputs) != len(args):
+            raise SpmdTypeError(
+                "BackwardHookFunction returned a different number of values than "
+                f"it received ({len(outputs)} vs. {len(args)})"
+            )
+        for out, arg in zip(outputs, args):
+            if isinstance(out, torch.Tensor) and isinstance(arg, torch.Tensor):
+                rules.einsum("...->...", arg, linear_in=(0,), out=out)
+
+    BackwardHookFunction.spmd_typecheck = staticmethod(_spmd_typecheck)
 
     _orig_setup_input_hook = _torch_hooks.BackwardHook.setup_input_hook
     _orig_setup_output_hook = _torch_hooks.BackwardHook.setup_output_hook