[spmd_types] qwen3.5 enablement

Enable Qwen3.5 input annotations, local multimodal regions, sharding configs, and torch_native GDN coverage for the spmd_types backend.

Status: FSDP=2 matched default backend bitwise over 10 debug steps. TP=2 and FSDP=2+TP=2 run, but are not bitwise against default backend; grad_norm diverges from step 1. TP=2 typechecking passes for qwen35_debugmodel and qwen35_debugmodel_gdn_torch_native.

ghstack-source-id: 9115a56
Pull Request resolved: #3895

Author: pianpwk

torchtitan/distributed/spmd_types.py

@@ -168,6 +168,9 @@ def annotate_input_spmd_types(
 
     Hardcodes the standard decoder convention: inputs and positions are
     ``S(0)@DP, S(1)@CP, R@TP``; labels are ``S(0)@DP, S(1)@CP, I@TP``.
+    Qwen3.5 multimodal payload tensors are irregular across DP/CP shards and
+    invariant across TP ranks; multimodal compute should handle them in local
+    SPMD regions.
     Analogous to ``full_dtensor.parallelize_inputs()`` but for the
     ``spmd_types`` path.
     """
@@ -183,15 +186,27 @@ def annotate_input_spmd_types(
         MeshAxisName.CP: spmd.S(1),
         MeshAxisName.TP: spmd.I,
     }
+    multimodal_type = {
+        MeshAxisName.DP: spmd.V,
+        MeshAxisName.CP: spmd.V,
+        MeshAxisName.TP: spmd.I,
+    }
 
     mesh = parallel_dims.spmd_dense_mesh()
     with set_current_spmd_mesh(mesh):
         spmd.assert_type(inputs, token_type)
         spmd.assert_type(labels, label_type)
-        if "positions" in extra_kwargs and isinstance(
-            extra_kwargs["positions"], torch.Tensor
+        for name in ("positions", "mrope_positions"):
+            if name in extra_kwargs and isinstance(extra_kwargs[name], torch.Tensor):
+                spmd.assert_type(extra_kwargs[name], token_type)
+        for name in (
+            "pixel_values",
+            "pixel_values_videos",
+            "grid_thw",
+            "grid_thw_videos",
         ):
-            spmd.assert_type(extra_kwargs["positions"], token_type)
+            if name in extra_kwargs and isinstance(extra_kwargs[name], torch.Tensor):
+                spmd.assert_type(extra_kwargs[name], multimodal_type)
     return inputs, labels, extra_kwargs
 
 

torchtitan/models/qwen3_5/config_registry.py

@@ -69,6 +69,14 @@ def qwen35_debugmodel() -> Trainer.Config:
     )
 
 
+def qwen35_debugmodel_gdn_torch_native() -> Trainer.Config:
+    config = qwen35_debugmodel()
+    for layer in config.model_spec.model.layers:
+        if layer.delta_net is not None:
+            layer.delta_net.kernel.backend = "torch_native"
+    return config
+
+
 def qwen35_debugmodel_moe() -> Trainer.Config:
     model_spec = model_registry("debugmodel_moe", moe_comm_backend="standard")
     return Trainer.Config(

torchtitan/models/qwen3_5/model.py

@@ -8,6 +8,7 @@
 from dataclasses import dataclass
 from typing import Literal
 
+import spmd_types as spmd
 import torch
 import torch.nn.functional as F
 
@@ -19,6 +20,8 @@
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.experimental import local_map
 
+from torchtitan.distributed.spmd_types import current_spmd_mesh
+from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.models.common import Conv1d, FeedForward, Linear
 from torchtitan.models.common.attention import AttentionMasksType, BaseAttention
 from torchtitan.models.common.decoder import Decoder
@@ -64,7 +67,7 @@ def _torch_native_gated_delta(
     k = _l2norm(k.float(), dim=-1)
     v, g, beta = v.float(), g.float(), beta.float()
 
-    output = torch.zeros(B, L, H, D_v, dtype=torch.float32, device=q.device)
+    output = torch.zeros_like(v, dtype=torch.float32)
     state = torch.zeros(B, H, D_k, D_v, dtype=torch.float32, device=q.device)
 
     for t in range(L):
@@ -196,29 +199,30 @@ def forward(
         if self.backend == "torch_native":
             return _torch_native_gated_delta(q, k, v, g, beta)
 
-        if self.backend == "fla_chunked":
-            result = _fla_chunk_gated_delta_rule(
-                q,
-                k,
-                v,
-                g,
-                beta,
-                use_qk_l2norm_in_kernel=True,
-            )
-        elif self.backend == "fla_fused_recurrent":
-            result = _fla_fused_recurrent_gated_delta_rule(
-                q,
-                k,
-                v,
-                g,
-                beta=beta,
-                use_qk_l2norm_in_kernel=True,
-            )
-        else:
-            raise ValueError(
-                f"Unknown fla_backend '{self.backend}'. "
-                "Valid: 'fla_chunked', 'fla_fused_recurrent', 'torch_native'."
-            )
+        with spmd.no_typecheck():
+            if self.backend == "fla_chunked":
+                result = _fla_chunk_gated_delta_rule(
+                    q,
+                    k,
+                    v,
+                    g,
+                    beta,
+                    use_qk_l2norm_in_kernel=True,
+                )
+            elif self.backend == "fla_fused_recurrent":
+                result = _fla_fused_recurrent_gated_delta_rule(
+                    q,
+                    k,
+                    v,
+                    g,
+                    beta=beta,
+                    use_qk_l2norm_in_kernel=True,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown fla_backend '{self.backend}'. "
+                    "Valid: 'fla_chunked', 'fla_fused_recurrent', 'torch_native'."
+                )
 
         # FLA kernels return (output, final_state); we only need output
         return result[0]
@@ -281,6 +285,20 @@ def __init__(self, config: Config):
 
     def _causal_conv(self, x: torch.Tensor, conv: nn.Module) -> torch.Tensor:
         x = F.pad(x.transpose(1, 2), [self.conv_kernel_size - 1, 0])
+
+        def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
+            # groups == local out-channels for depthwise channel-sharded conv.
+            # pyrefly: ignore [no-matching-overload]
+            return F.conv1d(
+                x_local,
+                w_local,
+                None,
+                conv.stride,
+                conv.padding,
+                conv.dilation,
+                w_local.size(0),
+            )
+
         if isinstance(x, DTensor):
             # TODO: Remove once the DTensor Conv1d dispatch fix for sharded
             # groups lands in a released torch. local_map runs the conv on
@@ -290,19 +308,6 @@ def _causal_conv(self, x: torch.Tensor, conv: nn.Module) -> torch.Tensor:
             w = conv.weight
             w_plc = w.placements  # pyrefly: ignore [missing-attribute]
 
-            def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
-                # groups == local out-channels (depthwise, channel-sharded)
-                # pyrefly: ignore [no-matching-overload]
-                return F.conv1d(
-                    x_local,
-                    w_local,
-                    None,
-                    conv.stride,
-                    conv.padding,
-                    conv.dilation,
-                    w_local.size(0),
-                )
-
             conv_dt = local_map(
                 _conv,
                 out_placements=(x_plc,),
@@ -311,6 +316,15 @@ def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
                 device_mesh=x.device_mesh,
             )
             x = conv_dt(x, w)  # pyrefly: ignore
+        elif get_spmd_backend() == "spmd_types":
+            conv_spmd = spmd.local_map(
+                in_types=(
+                    {"dp": spmd.S(0), "cp": spmd.S(2), "tp": spmd.S(1)},
+                    {"dp": spmd.R, "cp": spmd.R, "tp": spmd.S(0)},
+                ),
+                out_types={"dp": spmd.S(0), "cp": spmd.S(2), "tp": spmd.S(1)},
+            )(_conv)
+            x = conv_spmd(x, conv.weight)
         else:
             x = conv(x)
         return F.silu(x).transpose(1, 2)
@@ -329,9 +343,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         xa = self.in_proj_a(x)
         xb = self.in_proj_b(x)
 
-        xq = xq.view(bs, seqlen, -1, self.key_head_dim)
-        xk = xk.view(bs, seqlen, -1, self.key_head_dim)
-        xv = xv.view(bs, seqlen, -1, self.value_head_dim)
+        def local_head_split(t: torch.Tensor, head_dim: int) -> torch.Tensor:
+            # Drop into local region, we can't propagate S(2) -> head unflatten.
+            # TODO(pianpwk): this should be doable once spmd_types tracks sharding evenness.
+            with spmd.local():
+                t = t.view(bs, seqlen, -1, head_dim)
+                if get_spmd_backend() == "spmd_types":
+                    spmd.assert_type(
+                        t, spmd.V, spmd.PartitionSpec("dp", "cp", "tp", None)
+                    )
+            return t
+
+        xq = local_head_split(xq, self.key_head_dim)
+        xk = local_head_split(xk, self.key_head_dim)
+        xv = local_head_split(xv, self.value_head_dim)
 
         # Gating signals, shape (bs, seqlen, n_value_heads):
         #   g:    decay rate per head, always negative
@@ -341,7 +366,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         output = self.kernel(xq, xk, xv, g, beta)
 
-        xz = xz.view(bs, seqlen, -1, self.value_head_dim)
+        xz = local_head_split(xz, self.value_head_dim)
         output = self.norm(output, xz)
 
         output = output.reshape(bs, seqlen, -1)
@@ -407,11 +432,22 @@ def forward(
     ) -> torch.Tensor:
         bs, seqlen, _ = x.shape
 
+        def local_head_split(t: torch.Tensor, head_dim: int) -> torch.Tensor:
+            # Drop into local region, we can't propagate S(2) -> head unflatten.
+            # TODO(pianpwk): this should be doable once spmd_types tracks sharding evenness.
+            with spmd.local():
+                t = t.view(bs, seqlen, -1, head_dim)
+                if get_spmd_backend() == "spmd_types":
+                    spmd.assert_type(
+                        t, spmd.V, spmd.PartitionSpec("dp", "cp", "tp", None)
+                    )
+            return t
+
         # wq is 2x wider: produces query + gate
-        xq_gate = self.wq(x).view(bs, seqlen, -1, self.head_dim * 2)
+        xq_gate = local_head_split(self.wq(x), self.head_dim * 2)
         xq, gate = xq_gate.chunk(2, dim=-1)
-        xk = self.wk(x).view(bs, seqlen, -1, self.head_dim)
-        xv = self.wv(x).view(bs, seqlen, -1, self.head_dim)
+        xk = local_head_split(self.wk(x), self.head_dim)
+        xv = local_head_split(self.wv(x), self.head_dim)
 
         # QK norm (before RoPE)
         xq = self.q_norm(xq)
@@ -675,12 +711,23 @@ def _scatter_vision_embeds(
         Returns:
             Updated embeddings
         """
+        # Vision compute is TP-invariant; I->R convert to mix with text embeddings.
+        if spmd.is_type_checking():
+            merged_embeds = spmd.convert(
+                merged_embeds,
+                current_spmd_mesh().get_group("tp"),  # pyrefly: ignore [missing-attribute]
+                src=spmd.I,
+                dst=spmd.R,
+                backward_options={"op_dtype": merged_embeds.dtype},
+            )
+
         for item_idx, sample_idx, vision_start, n_tokens in vision_positions:
             inputs_embeds[
                 sample_idx, vision_start : vision_start + n_tokens, :
             ] = merged_embeds[item_idx, :n_tokens, :]
         return inputs_embeds
 
+    @spmd.local_map(out_types={"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.R})
     def _prepare_multimodal_embeds(
         self,
         tokens: torch.Tensor,

torchtitan/models/qwen3_5/parallelize.py

@@ -14,7 +14,7 @@
 import torch
 import torch.nn as nn
 from torch.distributed._composable.fsdp import fully_shard
-from torch.distributed.fsdp import MixedPrecisionPolicy
+from torch.distributed.fsdp import DataParallelMeshDims, MixedPrecisionPolicy
 
 from torchtitan.config import (
     CompileConfig,
@@ -30,6 +30,11 @@
     apply_fsdp_to_decoder,
     get_fsdp_reshard_after_forward_policy,
 )
+from torchtitan.distributed.full_dtensor import (
+    resolve_fsdp_mesh,
+    resolve_sparse_fsdp_mesh,
+    validate_config,
+)
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
 
 
@@ -40,6 +45,7 @@ def _apply_fsdp_to_vision_encoder(
     reduce_dtype: torch.dtype,
     reshard_after_forward_policy: str = "default",
     pp_enabled: bool = False,
+    dp_mesh_dims: DataParallelMeshDims | None = None,
 ):
     """FSDP the vision encoder as a single unit.
 
@@ -56,6 +62,7 @@ def _apply_fsdp_to_vision_encoder(
         mesh=dp_mesh,
         mp_policy=mp_policy,
         reshard_after_forward=reshard_after_forward,
+        dp_mesh_dims=dp_mesh_dims,
     )
 
 
@@ -90,10 +97,14 @@ def parallelize_qwen3_5(
             "and multimodal CP needs vision scatter before CP sharding."
         )
 
-    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
-        if parallelism.enable_async_tensor_parallel and not model_compile_enabled:
-            raise RuntimeError("Async TP requires torch.compile")
+    if parallelism.enable_async_tensor_parallel and not model_compile_enabled:
+        raise RuntimeError("Async TP requires torch.compile")
 
+    if parallelism.spmd_backend == "spmd_types":
+        validate_config(parallel_dims, model)
+        # pyrefly: ignore [not-callable]
+        model.parallelize(parallel_dims)
+    elif parallel_dims.tp_enabled or parallel_dims.ep_enabled:
         # pyrefly: ignore [not-callable]
         model.parallelize(parallel_dims)
 
@@ -112,10 +123,24 @@ def parallelize_qwen3_5(
             # pyrefly: ignore [bad-argument-type]
             apply_compile(model.vision_encoder, compile_config)
 
-    dp_mesh_names = (
-        ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
-    )
-    dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
+    if parallelism.spmd_backend == "spmd_types":
+        dp_mesh, dp_mesh_dims = resolve_fsdp_mesh(parallel_dims)
+        edp_mesh, edp_mesh_dims = resolve_sparse_fsdp_mesh(parallel_dims)
+    else:
+        dp_mesh_names = (
+            ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
+        )
+        dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
+        dp_mesh_dims = None
+        edp_mesh = None
+        edp_mesh_dims = None
+        if parallel_dims.ep_enabled:
+            edp_mesh_names = (
+                ["dp_replicate", "efsdp"]
+                if parallel_dims.dp_replicate_enabled
+                else ["efsdp"]
+            )
+            edp_mesh = parallel_dims.get_optional_mesh(edp_mesh_names)
 
     if model.vision_encoder is not None:
         _apply_fsdp_to_vision_encoder(
@@ -125,17 +150,9 @@ def parallelize_qwen3_5(
             reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce],
             reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward,
             pp_enabled=parallel_dims.pp_enabled,
+            dp_mesh_dims=dp_mesh_dims,
         )
 
-    edp_mesh = None
-    if parallel_dims.ep_enabled:
-        edp_mesh_names = (
-            ["dp_replicate", "efsdp"]
-            if parallel_dims.dp_replicate_enabled
-            else ["efsdp"]
-        )
-        edp_mesh = parallel_dims.get_optional_mesh(edp_mesh_names)
-
     apply_fsdp_to_decoder(
         model,  # pyrefly: ignore [bad-argument-type]
         dp_mesh,
@@ -146,6 +163,9 @@ def parallelize_qwen3_5(
         reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward,
         ep_degree=parallel_dims.ep,
         edp_mesh=edp_mesh,
+        dp_mesh_dims=dp_mesh_dims,
+        edp_mesh_dims=edp_mesh_dims,
+        enable_symm_mem=parallelism.enable_fsdp_symm_mem,
     )
 
     return model

torchtitan/models/qwen3_5/rope.py

@@ -77,7 +77,8 @@ def _compute_mrope_cache(self, position_ids: torch.Tensor) -> torch.Tensor:
             if isinstance(position_ids, DTensor)
             else position_ids
         )
-        pos = pos.to(device=rope_cache.device)
+        if pos.device != rope_cache.device:
+            pos = pos.to(device=rope_cache.device)
 
         _maybe_check_max_pos(pos, max_valid_pos=rope_cache.shape[0] - 1)
         head_dim = rope_cache.shape[-1] // 2