Upgrade spmd_types to 0.2.4

.ci/docker/requirements.txt

@@ -7,4 +7,4 @@ tokenizers >= 0.15.0
 safetensors
 einops
 pillow
-spmd_types==0.2.3
+spmd_types==0.2.4

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "wandb",
     "einops",
     "pillow",
-    "spmd_types==0.2.3",
+    "spmd_types==0.2.4",
 ]
 dynamic = ["version"]
 

tests/unit_tests/test_parallel_dims.py

@@ -443,8 +443,8 @@ def test_spmd_redistribute_per_axis_allgather(self):
             result = spmd_redistribute_per_axis(
                 x,
                 mesh,
-                src.per_axis_spmd_types(),
-                dst.per_axis_spmd_types(),
+                src,
+                dst,
             )
 
         self.assertEqual(comm_mode.get_total_counts(), 1)

torchtitan/components/loss.py

@@ -62,7 +62,6 @@ def cross_entropy_loss(
     )
 
 
-@spmd.register_autograd_function
 class _LossParallelCrossEntropy(torch.autograd.Function):
     """
     Vocab-parallel cross-entropy on local ``[T, V_local]`` logits.
@@ -80,30 +79,22 @@ class _LossParallelCrossEntropy(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         logits: torch.Tensor,
         labels: torch.Tensor,
         tp_group: dist.ProcessGroup,
-        global_vocab_size: int,
-        reduction: str = "sum",
-    ) -> torch.Tensor:
+    ) -> None:
         """
         SPMD type: logits S(-1)@TP, labels I@TP -> loss I@TP.
         Non-TP axes are passed through from logits to the output.
         """
         spmd.assert_type(logits, {tp_group: spmd.S(logits.dim() - 1)})
         spmd.assert_type(labels, {tp_group: spmd.I})
-        result = _LossParallelCrossEntropy.apply(
-            logits,
-            labels,
-            tp_group,
-            global_vocab_size,
-            reduction,
-        )
         output_type = dict(spmd.get_local_type(logits))
-        output_type[tp_group] = spmd.I
+        output_type[spmd.MeshAxis.of(tp_group)] = spmd.I
         spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     # pyrefly: ignore [bad-override]
@@ -261,13 +252,16 @@ def __call__(
         """Return the scaled loss and any metrics computed by the loss."""
         del kwargs
         loss = self.fn(pred, labels)
+        # loss: V->P, annotate global_valid_tokens
+        if get_spmd_backend() == "spmd_types" and current_spmd_mesh() is not None:
+            spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P})
+            if global_valid_tokens is not None:
+                spmd.assert_type(
+                    global_valid_tokens,
+                    {"dp": spmd.R, "cp": spmd.R, "tp": spmd.I},
+                )
         if global_valid_tokens is not None:
-            # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
-            is_type_checking = spmd.is_type_checking()
-            with spmd.no_typecheck():
-                loss = loss / global_valid_tokens
-                if is_type_checking:
-                    spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I})
+            loss = loss / global_valid_tokens
         return loss, {}
 
 
@@ -293,13 +287,16 @@ def __call__(
     ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         del kwargs
         loss = self.fn(pred, labels, global_vocab_size=self.global_vocab_size)
+        # loss: V->P, annotate global_valid_tokens
+        if get_spmd_backend() == "spmd_types" and current_spmd_mesh() is not None:
+            spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P})
+            if global_valid_tokens is not None:
+                spmd.assert_type(
+                    global_valid_tokens,
+                    {"dp": spmd.R, "cp": spmd.R, "tp": spmd.I},
+                )
         if global_valid_tokens is not None:
-            # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
-            is_type_checking = spmd.is_type_checking()
-            with spmd.no_typecheck():
-                loss = loss / global_valid_tokens
-                if is_type_checking:
-                    spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I})
+            loss = loss / global_valid_tokens
         return loss, {}
 
 
@@ -350,15 +347,13 @@ def compute_logprobs(
     elif get_spmd_backend() == "spmd_types" and spmd_mesh_size("tp") > 1:
         # spmd_types returns a plain local vocab shard. Labels are global token
         # ids, so cross_entropy needs full-vocab logits.
-        mesh = current_spmd_mesh()
-        assert mesh is not None
         # dst=I, not R: the vocab all-gather's grad is the replicated upstream
         # grad sliced back to this rank's vocab shard (I's backward), not an
         # all-reduce (R's backward). The latter over-counts by tp_degree and
         # diverges from the DTensor path above, whose redistribute grad slices.
         logits = spmd.redistribute(
             logits,
-            mesh.get_group("tp"),
+            "tp",
             src=spmd.S(-1),
             dst=spmd.I,
         )
@@ -644,15 +639,11 @@ def _chunk(t):
 
             total_loss = hidden_states.new_zeros((), dtype=torch.float32)
             if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
-                # TODO(pianpwk): would be nice if mutate_type accepted multiple axes.
-                for axis_name, dst in {
-                    "dp": spmd.P,
-                    "cp": spmd.P,
-                    "tp": spmd.I,
-                }.items():
-                    total_loss = spmd.mutate_type(
-                        total_loss, axis_name, src=spmd.R, dst=dst
-                    )
+                total_loss = spmd.mutate_type(
+                    total_loss,
+                    src=spmd.R,
+                    dst={"dp": spmd.P, "cp": spmd.P, "tp": spmd.I},
+                )
             metrics: dict[str, torch.Tensor] = {}
 
             # Disable FSDP reshard on lm_head to keep weight unsharded across

torchtitan/components/quantization/nvfp4.py

@@ -119,9 +119,7 @@ def build(self, **kwargs):
                 instance = Linear.Config.build(self, **kwargs)
                 if instance._sharding_config is not None:
                     sc = instance._sharding_config
-                    weight_tp = (
-                        sc.state_shardings["weight"].per_axis_spmd_types().get(TP)
-                    )
+                    weight_tp = sc.state_shardings["weight"].local_type.get(TP)
                     rowwise = isinstance(weight_tp, spmd.Shard) and weight_tp.dim == 1
                     if rowwise:
                         in_layout = dense_activation_placement(

torchtitan/distributed/linear.py

@@ -44,7 +44,6 @@ def ensure_symm_mem_ops():
     return symm_mem
 
 
-@spmd.register_autograd_function
 class AllGatherLinear(torch.autograd.Function):
     """All-gather the sequence shard, then apply a column-parallel linear.
 
@@ -80,32 +79,32 @@ class AllGatherLinear(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         x_shard_m: torch.Tensor,
         w_shard_n: torch.Tensor,
         bias_shard_n: torch.Tensor | None,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> torch.Tensor:
+    ) -> None:
         """SPMD type: x S(0)@TP, w S(0)@TP -> y S(1)@TP.
 
         The gather consumes the row shard, so the result is full on rows; the
         weight's output-feature shard survives the GEMM. Non-TP axes pass through
         from x.
         """
-        spmd.assert_type(x_shard_m, {group: spmd.S(0)})
+        spmd.assert_type(x_shard_m, {group_name: spmd.S(0)})
         # S(0), not S(1), even though this is the column-parallel direction: torch
         # stores the weight as [N, K] while the mental model of the GEMM is [K, N],
         # so sharding the output features N is dim 0 of what is actually stored.
-        spmd.assert_type(w_shard_n, {group: spmd.S(0)})
+        spmd.assert_type(w_shard_n, {group_name: spmd.S(0)})
         if bias_shard_n is not None:
-            spmd.assert_type(bias_shard_n, {group: spmd.S(0)})
-        result = AllGatherLinear.apply(
-            x_shard_m, w_shard_n, bias_shard_n, group, group_name
+            spmd.assert_type(bias_shard_n, {group_name: spmd.S(0)})
+        spmd.assert_local_type_like(
+            result,
+            x_shard_m,
+            {group_name: spmd.S(1)},  # pyrefly: ignore [bad-argument-type]
         )
-        output_type = {**spmd.get_local_type(x_shard_m), group: spmd.S(1)}
-        spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     def forward(  # pyrefly: ignore[bad-override]
@@ -183,7 +182,6 @@ def backward(ctx, grad_y_shard_n: torch.Tensor):  # pyrefly: ignore[bad-override
         return grad_x_shard_m, grad_w_shard_n, grad_bias, None, None
 
 
-@spmd.register_autograd_function
 class AllGatherLinearMulti(torch.autograd.Function):
     """One all-gather feeding a pair of column-parallel linears on the same input.
 
@@ -282,29 +280,28 @@ def backward(  # pyrefly: ignore[bad-override]
         )
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        results: tuple[torch.Tensor, torch.Tensor],
+        *,
         x_shard_m: torch.Tensor,
         wa_shard_n: torch.Tensor,
         wb_shard_n: torch.Tensor,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> None:
         """SPMD type: x S(0)@TP, both w S(0)@TP -> both y S(1)@TP."""
-        spmd.assert_type(x_shard_m, {group: spmd.S(0)})
+        spmd.assert_type(x_shard_m, {group_name: spmd.S(0)})
         # S(0) for the column-parallel direction; see AllGatherLinear for why the
         # stored [N, K] layout inverts the dim you would expect.
-        spmd.assert_type(wa_shard_n, {group: spmd.S(0)})
-        spmd.assert_type(wb_shard_n, {group: spmd.S(0)})
-        results = AllGatherLinearMulti.apply(
-            x_shard_m, wa_shard_n, wb_shard_n, group, group_name
-        )
-        output_type = {**spmd.get_local_type(x_shard_m), group: spmd.S(1)}
+        spmd.assert_type(wa_shard_n, {group_name: spmd.S(0)})
+        spmd.assert_type(wb_shard_n, {group_name: spmd.S(0)})
         for result in results:
-            spmd.assert_type(result, output_type)
-        return results
+            spmd.assert_local_type_like(
+                result,
+                x_shard_m,
+                {group_name: spmd.S(1)},  # pyrefly: ignore [bad-argument-type]
+            )
 
 
-@spmd.register_autograd_function
 class LinearReduceScatter(torch.autograd.Function):
     """Apply a row-parallel linear, then reduce-scatter over the sequence.
 
@@ -329,32 +326,32 @@ class LinearReduceScatter(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         x_shard_k: torch.Tensor,
         w_shard_k: torch.Tensor,
         bias: torch.Tensor | None,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> torch.Tensor:
+    ) -> None:
         """SPMD type: x S(1)@TP, w S(1)@TP, bias R@TP -> y S(0)@TP.
 
         The local matmul is a partial sum over the sharded K; the reduce-scatter
         completes it and shards rows instead. Non-TP axes pass through from x.
         """
-        spmd.assert_type(x_shard_k, {group: spmd.S(1)})
+        spmd.assert_type(x_shard_k, {group_name: spmd.S(1)})
         # S(1), the mirror of AllGatherLinear's S(0): torch stores the weight as
         # [N, K] while the mental model of the GEMM is [K, N], so sharding the
         # input features K -- the row-parallel direction -- is dim 1 of what is
         # actually stored.
-        spmd.assert_type(w_shard_k, {group: spmd.S(1)})
+        spmd.assert_type(w_shard_k, {group_name: spmd.S(1)})
         if bias is not None:
-            spmd.assert_type(bias, {group: spmd.R})
-        result = LinearReduceScatter.apply(
-            x_shard_k, w_shard_k, bias, group, group_name
+            spmd.assert_type(bias, {group_name: spmd.R})
+        spmd.assert_local_type_like(
+            result,
+            x_shard_k,
+            {group_name: spmd.S(0)},  # pyrefly: ignore [bad-argument-type]
         )
-        output_type = {**spmd.get_local_type(x_shard_k), group: spmd.S(0)}
-        spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     def forward(  # pyrefly: ignore[bad-override]