[spmd_types] Kimi 2.7 enablement

tests/integration_tests/models.py

@@ -13,9 +13,7 @@
 def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefinitions:
     """Use ``backend`` for every variant, or return an unsupported test unchanged."""
     if backend == "spmd_types" and any(
-        "--module kimi_k2_7" in arg or "--module muse_glimmer" in arg
-        for variant in t.override_args
-        for arg in variant
+        "--module muse_glimmer" in arg for variant in t.override_args for arg in variant
     ):
         return t
 

torchtitan/models/common/attention.py

@@ -156,13 +156,15 @@ def forward(
         max_q = attention_masks.max_q
         max_k = attention_masks.max_k
 
-        B, L, _, H = q_BLNH.shape
+        B, L, _, qk_head_dim = q_BLNH.shape
+        value_head_dim = v_BLNH.shape[-1]
         T = B * L
 
-        # varlen attention expects (T, N, H)
-        q_TNH = q_BLNH.reshape(T, -1, H)
-        k_TNH = k_BLNH.reshape(T, -1, H)
-        v_TNH = v_BLNH.reshape(T, -1, H)
+        # varlen attention expects (T, N, H). The value head dimension may
+        # differ from the query/key head dimension, as in MLA.
+        q_TNH = q_BLNH.reshape(T, -1, qk_head_dim)
+        k_TNH = k_BLNH.reshape(T, -1, qk_head_dim)
+        v_TNH = v_BLNH.reshape(T, -1, value_head_dim)
 
         # Some operators can upcast under AMP, but varlen attention currently only
         # supports bf16/fp16 inputs. If this changes, or fp16 training support
@@ -219,7 +221,7 @@ def forward(
                 # rejected during q_TNH reshape propagation.
                 q_ps = get_partition_spec(q_TNH)
                 spmd.assert_type(result, q_local, q_ps)
-            out_BLNH = result.view(B, L, -1, H).to(q_BLNH.dtype)
+            out_BLNH = result.view(B, L, -1, value_head_dim).to(q_BLNH.dtype)
             return out_BLNH
 
         out_TNH, lse_NT = result
@@ -231,7 +233,7 @@ def forward(
             lse_ps = None if q_ps is None else spmd.PartitionSpec(q_ps[1], q_ps[0])
             spmd.assert_type(lse_NT, q_local, lse_ps)
 
-        out_BLNH = out_TNH.view(B, L, -1, H).to(q_BLNH.dtype)
+        out_BLNH = out_TNH.view(B, L, -1, value_head_dim).to(q_BLNH.dtype)
         # FA varlen returns the LSE as (N, T); reorder to (B, L, N) so
         # out_transform can broadcast per (token, head).
         lse_BLN = lse_NT.transpose(0, 1).reshape(B, L, -1)

torchtitan/models/kimi_k2_7/__init__.py

@@ -17,6 +17,7 @@
     Embedding,
     Linear,
     RMSNorm,
+    ScaledBiasRowwiseLinear,
     TransformerBlock,
 )
 from torchtitan.models.common.nn_modules import LayerNorm
@@ -104,6 +105,17 @@ def _vl_linear(in_features: int, out_features: int) -> Linear.Config:
     )
 
 
+def _scaled_bias_rowwise_linear(
+    in_features: int, out_features: int
+) -> ScaledBiasRowwiseLinear.Config:
+    return ScaledBiasRowwiseLinear.Config(
+        in_features=in_features,
+        out_features=out_features,
+        bias=True,
+        param_init=_LINEAR_INIT,
+    )
+
+
 def _vl_layernorm(dim: int, eps: float = 1e-5) -> LayerNorm.Config:
     return LayerNorm.Config(normalized_shape=dim, eps=eps)
 
@@ -138,11 +150,11 @@ def _vision_encoder_config(
             wq=_vl_linear(dim, dim),
             wk=_vl_linear(dim, dim),
             wv=_vl_linear(dim, dim),
-            proj=_vl_linear(dim, dim),
+            proj=_scaled_bias_rowwise_linear(dim, dim),
         ),
         mlp=VisionMLP.Config(
             fc1=_vl_linear(dim, ffn_dim),
-            fc2=_vl_linear(ffn_dim, dim),
+            fc2=_scaled_bias_rowwise_linear(ffn_dim, dim),
         ),
     )
 
@@ -168,7 +180,7 @@ def _vision_encoder_config(
             merged_dim=merged_dim,
             pre_norm=_vl_layernorm(dim),
             linear_1=_vl_linear(merged_dim, merged_dim),
-            linear_2=_vl_linear(merged_dim, text_hidden_size),
+            linear_2=_scaled_bias_rowwise_linear(merged_dim, text_hidden_size),
         ),
     )
 

torchtitan/models/kimi_k2_7/model.py

@@ -8,10 +8,14 @@
 https://github.com/sgl-project/sglang/blob/e0c0c0a45cb1bda90392bfa2bba4184f5b0638a0/python/sglang/srt/models/kimi_k25.py
 """
 
+import contextlib
 from dataclasses import dataclass
 
+import spmd_types as spmd
 import torch
 
+from torchtitan.distributed.spmd_types import spmd_mesh_size
+from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.models.common.attention import AttentionMasksType
 from torchtitan.models.common.decoder import Decoder
 from torchtitan.models.common.multimodal import (
@@ -20,7 +24,10 @@
 )
 from torchtitan.models.deepseek_v3.model import DeepSeekV3Model
 
-from .sharding import set_kimi_k2_5_sharding_config
+from .sharding import (
+    annotate_multimodal_input_spmd_types,
+    set_kimi_k2_5_sharding_config,
+)
 from .vision_encoder import KimiK25VisionEncoder
 
 
@@ -76,6 +83,12 @@ def __init__(self, config: Config):
             config.vision_encoder.build() if config.vision_encoder is not None else None
         )
 
+    def multimodal_context(self) -> contextlib.AbstractContextManager[None]:
+        """Use local DP typechecking while preparing multimodal inputs."""
+        if get_spmd_backend() == "spmd_types" and spmd_mesh_size("dp") > 1:
+            return spmd.set_current_mesh(local_axes=("dp",))
+        return contextlib.nullcontext()
+
     def _prepare_multimodal_embeds(
         self,
         tokens: torch.Tensor,
@@ -166,17 +179,31 @@ def forward(  # pyrefly: ignore [bad-override]
         Returns:
             (batch, seq_len, vocab_size) logits.
         """
-        if self.tok_embeddings is not None:
-            x = self._prepare_multimodal_embeds(
-                tokens,
-                pixel_values=pixel_values,
-                grid_thw=grid_thw,
-                pixel_values_videos=pixel_values_videos,
-                grid_thw_videos=grid_thw_videos,
-                special_tokens=special_tokens,  # pyrefly: ignore [bad-argument-type]
-            )
-        else:
-            x = tokens
+        with self.multimodal_context():
+            if get_spmd_backend() == "spmd_types":
+                annotate_multimodal_input_spmd_types(
+                    pixel_values=pixel_values,
+                    grid_thw=grid_thw,
+                    pixel_values_videos=pixel_values_videos,
+                    grid_thw_videos=grid_thw_videos,
+                )
+
+            if self.tok_embeddings is not None:
+                x = self._prepare_multimodal_embeds(
+                    tokens,
+                    pixel_values=pixel_values,
+                    grid_thw=grid_thw,
+                    pixel_values_videos=pixel_values_videos,
+                    grid_thw_videos=grid_thw_videos,
+                    special_tokens=special_tokens,  # pyrefly: ignore [bad-argument-type]
+                )
+            else:
+                x = tokens
+
+        if get_spmd_backend() == "spmd_types":
+            # The scatter restores a token-aligned tensor, so text-model DP
+            # resumes as global batch sharding after the multimodal region.
+            spmd.assert_type(x, {"dp": spmd.S(0), "tp": spmd.R})
 
         for layer in self.layers.values():
             x = layer(x, attention_masks, positions)

torchtitan/models/kimi_k2_7/parallelize.py

@@ -32,6 +32,7 @@
 from torchtitan.distributed.full_dtensor import (
     resolve_fsdp_mesh,
     resolve_sparse_fsdp_mesh,
+    validate_config,
 )
 
 
@@ -67,11 +68,10 @@ def parallelize_kimi_k2_5(
         compile_config.enable and "model" in compile_config.components
     )
 
-    if (
-        parallelism.spmd_backend == "spmd_types"
-        or parallel_dims.tp_enabled
-        or parallel_dims.ep_enabled
-    ):
+    if parallelism.spmd_backend == "spmd_types":
+        validate_config(parallel_dims, model)
+        model.parallelize(parallel_dims)  # pyrefly: ignore [not-callable]
+    elif parallel_dims.tp_enabled or parallel_dims.ep_enabled:
         model.parallelize(parallel_dims)  # pyrefly: ignore [not-callable]
 
     if ac_config is not None:

torchtitan/models/kimi_k2_7/sharding.py

@@ -12,39 +12,62 @@
 - Decoder (MLA + MoE): reuses ``set_deepseek_v3_sharding_config``. Multimodal
   configs keep the token embedding ``Replicate`` for the vision scatter and
   resume SP at layer 0 (see ``_shard_decoder_after_embedding_scatter``).
-- Vision encoder: activations flow ``Replicate`` (no SP -- the patch sequence is
+- Vision encoder: activations flow ``Invariant`` (no SP -- the patch sequence is
   short, so sequence-sharding would add gather/scatter around the block-diagonal
   attention for little memory gain). Only the linear layers are Colwise/Rowwise
-  sharded for memory; norms and position embeddings stay ``Replicate``.
+  sharded for memory; norms and position embeddings stay ``Invariant``.
 """
 
 from typing import TYPE_CHECKING
 
 import spmd_types as spmd
+import torch
 
+from torchtitan.distributed.parallel_dims import MeshAxisName
 from torchtitan.models.common.decoder_sharding import (
-    colwise_config,
     dense_activation_placement,
     dense_param_placement,
-    rowwise_config,
-    set_gqa_inner_attention_local_map,
+    dense_sequence_parallel_placement,
+)
+from torchtitan.models.common.vision_encoder_sharding import (
+    invariant_norm_config,
+    set_vision_transformer_block_sharding_config,
+    vision_colwise_config,
+    vision_invariant_linear_config,
+    vision_scaled_bias_rowwise_config,
 )
 from torchtitan.models.deepseek_v3.sharding import set_deepseek_v3_sharding_config
-from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig
+from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig, SpmdLayout
+
+DP = MeshAxisName.DP
+TP = MeshAxisName.TP
 
 if TYPE_CHECKING:
     from torchtitan.models.kimi_k2_7.model import KimiK25Model
 
-_REPLICATE_PARAM = dense_param_placement(tp=spmd.R)
 _REPLICATE_ACT = dense_activation_placement(tp=spmd.R)
 
-_REPLICATE_NORM = ShardingConfig(
-    state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM},
-    in_src_shardings={"input": _REPLICATE_ACT},
-    in_dst_shardings={"input": _REPLICATE_ACT},
-    out_src_shardings=_REPLICATE_ACT,
-    out_dst_shardings=_REPLICATE_ACT,
-)
+
+def annotate_multimodal_input_spmd_types(
+    *,
+    pixel_values: torch.Tensor | None,
+    grid_thw: torch.Tensor | None,
+    pixel_values_videos: torch.Tensor | None,
+    grid_thw_videos: torch.Tensor | None,
+) -> None:
+    """Annotate Kimi K2.5 multimodal inputs with their local SPMD types."""
+    multimodal_type = {
+        MeshAxisName.DP: spmd.V,
+        MeshAxisName.TP: spmd.I,
+    }
+    for tensor in (
+        pixel_values,
+        grid_thw,
+        pixel_values_videos,
+        grid_thw_videos,
+    ):
+        if tensor is not None:
+            spmd.assert_type(tensor, multimodal_type)
 
 
 def set_kimi_k2_5_sharding_config(
@@ -65,15 +88,7 @@ def set_kimi_k2_5_sharding_config(
 
 
 def _shard_decoder_after_embedding_scatter(config: "KimiK25Model.Config") -> None:
-    """Keep ``tok_embeddings`` ``Replicate`` and resume SP at layer 0's output.
-
-    The vision scatter writes features at arbitrary sequence positions, so it
-    needs the full (``Replicate``) embedding -- a ``Shard(1)`` one cannot be
-    indexed by sequence position locally. Layer 0 then takes a ``Replicate``
-    input and its rowwise ``wo`` reduce-scatters back to ``Shard(1)``, so the
-    residual is sequence-parallel from layer 0's output and layers ``1..N-1``
-    are unchanged full SP.
-    """
+    """Keep the full embedding through vision scatter, then resume SP at layer 0."""
     config.tok_embeddings.sharding_config = ShardingConfig(
         state_shardings={"weight": dense_param_placement(tp=spmd.S(0))},
         in_src_shardings={"input": _REPLICATE_ACT},
@@ -84,62 +99,46 @@ def _shard_decoder_after_embedding_scatter(config: "KimiK25Model.Config") -> Non
     )
 
     layer0 = config.layers[0]
-    layer0.attention_norm.sharding_config = ShardingConfig(
-        state_shardings={"weight": _REPLICATE_PARAM},
-        in_src_shardings={"input": _REPLICATE_ACT},
-        out_src_shardings=_REPLICATE_ACT,
-    )
-    layer0.attention.sharding_config = ShardingConfig(
+    layer0.sharding_config = ShardingConfig(
         in_src_shardings={"x": _REPLICATE_ACT},
-        in_dst_shardings={"x": _REPLICATE_ACT},
+        in_dst_shardings={"x": dense_sequence_parallel_placement()},
+        out_src_shardings=dense_sequence_parallel_placement(),
     )
 
 
 def _set_vision_encoder_sharding(ve_cfg) -> None:
-    """Replicate-activation TP plan for the MoonViT3d vision encoder.
+    """Invariant-activation TP plan for the MoonViT3d vision encoder.
 
     Linear layers are Colwise/Rowwise sharded for memory; norms and the
-    learnable position table are Replicate. ``patch_embed`` wraps the plain
-    ``pixel_values`` input as ``DTensor(Replicate)`` so the rest of the encoder
-    runs in DTensor space.
+    learnable position table stay Invariant. ``patch_embed`` wraps the plain
+    ``pixel_values`` input as a TP-invariant tensor so the rest of the encoder
+    runs in distributed tensor space.
     """
-    # The encoder's own ``pos_embed`` table is Replicate (F.interpolate runs on it).
+    # The encoder's own ``pos_embed`` table is invariant across TP ranks.
     ve_cfg.sharding_config = ShardingConfig(
-        state_shardings={"pos_embed": _REPLICATE_PARAM},
+        state_shardings={
+            "pos_embed": SpmdLayout({DP: spmd.R, TP: spmd.I}),
+        },
+        out_src_shardings=SpmdLayout({DP: spmd.V, TP: spmd.I}),
+        out_dst_shardings=SpmdLayout({DP: spmd.V, TP: spmd.R}),
     )
-
-    # patch_embed (Linear): receives plain pixel_values -> wrap as Replicate.
-    ve_cfg.patch_embed_proj.sharding_config = ShardingConfig(
-        state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM},
-        in_src_shardings={"input": _REPLICATE_ACT},
-        in_dst_shardings={"input": _REPLICATE_ACT},
-        out_src_shardings=_REPLICATE_ACT,
-        out_dst_shardings=_REPLICATE_ACT,
+    ve_cfg.rotary_pos_emb.sharding_config = ShardingConfig(
+        state_shardings={
+            "inv_freq": SpmdLayout({DP: spmd.R, TP: spmd.I}),
+        },
+        out_src_shardings=SpmdLayout({DP: spmd.R, TP: spmd.I}),
     )
 
-    # Transformer block sub-modules (shared VisionTransformerBlock: norm1/norm2).
-    block = ve_cfg.block
-    block.norm1.sharding_config = _REPLICATE_NORM
-    block.norm2.sharding_config = _REPLICATE_NORM
+    ve_cfg.patch_embed_proj.sharding_config = vision_invariant_linear_config()
 
-    # The stacked 2D rope_cache enters the attention as a plain (Replicate)
-    # tensor input so it is DTensor-wrapped before meeting head-sharded q/k.
-    block.attn.sharding_config = ShardingConfig(
-        in_src_shardings={"rope_cache": _REPLICATE_ACT},
-        in_dst_shardings={"rope_cache": _REPLICATE_ACT},
+    set_vision_transformer_block_sharding_config(
+        ve_cfg.block,
+        rope_cache_dp=spmd.V,
     )
-    block.attn.wq.sharding_config = colwise_config()
-    block.attn.wk.sharding_config = colwise_config()
-    block.attn.wv.sharding_config = colwise_config()
-    block.attn.proj.sharding_config = rowwise_config(output_sp=False)
-    set_gqa_inner_attention_local_map(block.attn.inner_attention)
-
-    block.mlp.fc1.sharding_config = colwise_config()
-    block.mlp.fc2.sharding_config = rowwise_config(output_sp=False)
 
     # Final norm + projector.
-    ve_cfg.final_norm.sharding_config = _REPLICATE_NORM
+    ve_cfg.final_norm.sharding_config = invariant_norm_config()
     proj = ve_cfg.projector
-    proj.pre_norm.sharding_config = _REPLICATE_NORM
-    proj.linear_1.sharding_config = colwise_config()
-    proj.linear_2.sharding_config = rowwise_config(output_sp=False)
+    proj.pre_norm.sharding_config = invariant_norm_config()
+    proj.linear_1.sharding_config = vision_colwise_config()
+    proj.linear_2.sharding_config = vision_scaled_bias_rowwise_config()