spmd_types 0.2.4__tar.gz → 0.2.5__tar.gz

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

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

{spmd_types-0.2.4 → spmd_types-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.4"
+version = "0.2.5"
 description = "A type system for distributed (SPMD) tensor computations in PyTorch"
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -48,3 +48,10 @@ 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" = ["../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.2.5}/spmd_types/_collectives.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 from typing import List, Optional
 
 import torch
+import torch.distributed._functional_collectives as funcol
 from spmd_types import _dist
 from spmd_types._dtype_utils import (
     _apply_op_dtype,
@@ -67,8 +68,9 @@ class _AllReduce(torch.autograd.Function):
             ctx.mark_dirty(x)
             result = x
         else:
-            result = x.clone()
-            _dist.dist.all_reduce(result, op=_dist.dist.ReduceOp.SUM, group=pg)
+            result = funcol.all_reduce(x, "sum", pg)
+            if isinstance(result, funcol.AsyncCollectiveTensor):
+                result = result.wait()
         if inplace:
             assert dtype_options.out_dtype == result.dtype
             return result
@@ -300,16 +302,9 @@ class _AllGatherShard(torch.autograd.Function):
         x = _apply_op_dtype(x, dtype_options.op_dtype)
         pg = axis
         x = x.contiguous()
-        world_size = _dist.dist.get_world_size(pg)
-        out_shape = list(x.shape)
-        out_shape[0] *= world_size
-        out = torch.empty(out_shape, device=x.device, dtype=x.dtype)
-        _dist.dist.all_gather_into_tensor(out, x, group=pg)
-        if gather_dim == 0:
-            result = out
-        else:
-            chunks = out.chunk(world_size, dim=0)
-            result = torch.cat(chunks, dim=gather_dim)
+        result = funcol.all_gather_tensor(x, gather_dim, pg)
+        if isinstance(result, funcol.AsyncCollectiveTensor):
+            result = result.wait()
         return _apply_out_dtype(result, dtype_options.out_dtype)
 
     @staticmethod
@@ -635,22 +630,9 @@ class _ReduceScatterShard(torch.autograd.Function):
         ctx.backward_options = dtype_options.backward_options
         x = _apply_op_dtype(x, dtype_options.op_dtype)
         pg = axis
-        world_size = _dist.dist.get_world_size(pg)
-        # reduce_scatter_tensor always scatters along dim 0, so we
-        # movedim before/after when scatter_dim != 0.
-        needs_permute = scatter_dim != 0
-        if needs_permute:
-            x = x.movedim(scatter_dim, 0).contiguous()
-
-        output_shape = list(x.shape)
-        output_shape[0] //= world_size
-        result = x.new_empty(output_shape)
-        _dist.dist.reduce_scatter_tensor(
-            result, x, op=_dist.dist.ReduceOp.SUM, group=pg
-        )
-
-        if needs_permute:
-            result = result.movedim(0, scatter_dim)
+        result = funcol.reduce_scatter_tensor(x, "sum", scatter_dim, pg)
+        if isinstance(result, funcol.AsyncCollectiveTensor):
+            result = result.wait()
         return _apply_out_dtype(result, dtype_options.out_dtype)
 
     @staticmethod

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

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