[muse_glimmer] Enable CP for the multimodal model

tests/integration_tests/models.py

@@ -200,4 +200,11 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             test_name="muse_glimmer_mm_fsdp+tp+sp",
             ngpu=4,
         ),
+        OverrideDefinitions(
+            configs=[recipes.muse_glimmer_debugmodel_mm_tp2_cp2_pp2],
+            test_descr="Muse Glimmer multimodal TP+CP+PP+SP",
+            test_name="muse_glimmer_mm_tp+cp+pp+sp",
+            ngpu=8,
+            use_real_pg=True,
+        ),
     ]

torchtitan/models/common/multimodal.py

@@ -6,19 +6,21 @@
 
 """Model-agnostic vision<->text fusion for VLMs.
 
-The decoder embeds the full token sequence; the placeholder tokens
-get a throwaway text embedding that ``scatter_vision_embeds``
-overwrites with the vision encoder's per-item features at the positions
-``get_vision_positions`` locates.
+``get_vision_positions`` and ``scatter_vision_embeds`` support span-based
+fusion over a full token sequence. ``build_vision_bank_indices`` and
+``VisionScatter`` support the equivalent gather-based fusion after token
+sharding by carrying an absolute packed-bank row for every placeholder token.
 """
 
 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.protocols.module import Module
 
 
 def multimodal_context() -> contextlib.AbstractContextManager[None]:
@@ -96,6 +98,42 @@ def get_vision_positions(
     return positions
 
 
+def build_vision_bank_indices(
+    tokens_T: torch.Tensor,
+    *,
+    placeholder_id: int,
+) -> torch.Tensor:
+    """Map vision placeholder tokens to absolute packed-bank rows."""
+    vision_mask_T = tokens_T == placeholder_id
+    vision_bank_indices_T = torch.cumsum(vision_mask_T.to(torch.long), dim=0) - 1
+    return vision_bank_indices_T.masked_fill(~vision_mask_T, -1)
+
+
+class VisionScatter(Module):
+    """Sharding boundary for token-local packed-vision gathering."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        pass
+
+    def __init__(self, config: "VisionScatter.Config") -> None:
+        super().__init__()
+        del config
+
+    def forward(
+        self,
+        inputs_TD: torch.Tensor,
+        vision_bank_VD: torch.Tensor,
+        vision_bank_indices_T: torch.Tensor,
+    ) -> torch.Tensor:
+        if vision_bank_VD.shape[0] == 0:
+            return inputs_TD
+        vision_bank_VD = vision_bank_VD.to(inputs_TD.dtype)
+        is_vision_T1 = (vision_bank_indices_T >= 0).unsqueeze(-1)
+        gathered_TD = vision_bank_VD[vision_bank_indices_T.clamp(min=0)]
+        return torch.where(is_vision_T1, gathered_TD, inputs_TD)
+
+
 def scatter_vision_embeds(
     inputs_embeds: torch.Tensor,
     *,

torchtitan/models/common/vision_encoder_sharding.py

@@ -12,17 +12,29 @@
 from spmd_types import SpmdType
 
 from torchtitan.distributed.parallel_dims import MeshAxisName
-from torchtitan.models.common.decoder_sharding import set_gqa_inner_attention_local_map
 from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig
 
 if TYPE_CHECKING:
     from torchtitan.models.common.vision_encoder import VisionTransformerBlock
 
 
 DP = MeshAxisName.DP
+CP = MeshAxisName.CP
 TP = MeshAxisName.TP
 
 
+def _vision_state_placement(*, tp: spmd.PerMeshAxisSpmdType) -> SpmdType:
+    return SpmdType({DP: spmd.R, CP: spmd.R, TP: tp})
+
+
+def _vision_activation_placement(
+    *,
+    dp: spmd.PerMeshAxisSpmdType = spmd.V,
+    tp: spmd.PerMeshAxisSpmdType = spmd.I,
+) -> SpmdType:
+    return SpmdType({DP: dp, CP: spmd.R, TP: tp})
+
+
 def multimodal_input_sharding() -> dict[str, SpmdType]:
     """SPMD layouts for VLM vision inputs (folded into a model's input_sharding).
 
@@ -31,7 +43,7 @@ def multimodal_input_sharding() -> dict[str, SpmdType]:
     ``multimodal_context`` (a DP-local mesh) and the vision encoder runs per-rank.
     Shared by every VLM decoder (Qwen3.5, Kimi K2.5, Muse Glimmer).
     """
-    layout = SpmdType({DP: spmd.V, TP: spmd.I})
+    layout = _vision_activation_placement()
     return {
         "pixel_values": layout,
         "pixel_values_videos": layout,
@@ -44,35 +56,35 @@ def invariant_norm_config() -> ShardingConfig:
     """Norm whose state and activations are invariant across TP ranks."""
     return ShardingConfig(
         state_shardings={
-            "weight": SpmdType({DP: spmd.R, TP: spmd.I}),
-            "bias": SpmdType({DP: spmd.R, TP: spmd.I}),
+            "weight": _vision_state_placement(tp=spmd.I),
+            "bias": _vision_state_placement(tp=spmd.I),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(),
         },
-        out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.I}),
-        out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}),
+        out_src_shardings=_vision_activation_placement(),
+        out_dst_shardings=_vision_activation_placement(),
     )
 
 
 def vision_invariant_linear_config() -> ShardingConfig:
     """Unsharded linear whose state and activations are invariant at TP."""
     return ShardingConfig(
         state_shardings={
-            "weight": SpmdType({DP: spmd.R, TP: spmd.I}),
-            "bias": SpmdType({DP: spmd.R, TP: spmd.I}),
+            "weight": _vision_state_placement(tp=spmd.I),
+            "bias": _vision_state_placement(tp=spmd.I),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(),
         },
-        out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.I}),
-        out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}),
+        out_src_shardings=_vision_activation_placement(),
+        out_dst_shardings=_vision_activation_placement(),
     )
 
 
@@ -82,36 +94,37 @@ def vision_colwise_config(
     """Colwise vision linear with a TP-replicated local matmul input."""
     return ShardingConfig(
         state_shardings={
-            "weight": SpmdType({DP: spmd.R, TP: spmd.S(0)}),
-            "bias": SpmdType({DP: spmd.R, TP: spmd.S(0)}),
+            "weight": _vision_state_placement(tp=spmd.S(0)),
+            "bias": _vision_state_placement(tp=spmd.S(0)),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: input_tp}),
+            "input": _vision_activation_placement(tp=input_tp),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.R}),
+            "input": _vision_activation_placement(tp=spmd.R),
         },
-        out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.S(-1)}),
+        out_src_shardings=_vision_activation_placement(tp=spmd.S(-1)),
     )
 
 
 def vision_scaled_bias_rowwise_config() -> ShardingConfig:
     """Scaled-bias rowwise vision linear returning a TP-invariant activation."""
-    input_layout = SpmdType({DP: spmd.V, TP: spmd.S(1)})
+    input_layout = _vision_activation_placement(tp=spmd.S(1))
+    input_grad_layout = SpmdType({DP: spmd.V, CP: spmd.P, TP: spmd.S(1)})
     return ShardingConfig(
         state_shardings={
-            "weight": SpmdType({DP: spmd.R, TP: spmd.S(1)}),
-            "bias": SpmdType({DP: spmd.R, TP: spmd.R}),
+            "weight": _vision_state_placement(tp=spmd.S(1)),
+            "bias": _vision_state_placement(tp=spmd.R),
         },
         in_src_shardings={
             "input": input_layout,
         },
         in_dst_shardings={
             "input": input_layout,
         },
-        out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.P}),
-        out_dst_shardings=SpmdType({DP: spmd.V, TP: spmd.I}),
-        local_map=LocalMapConfig(in_grad_placements=(input_layout,)),
+        out_src_shardings=_vision_activation_placement(tp=spmd.P),
+        out_dst_shardings=_vision_activation_placement(),
+        local_map=LocalMapConfig(in_grad_placements=(input_grad_layout,)),
     )
 
 
@@ -126,19 +139,39 @@ def set_vision_transformer_block_sharding_config(
 
     block.attn.sharding_config = ShardingConfig(
         in_src_shardings={
-            "x": SpmdType({DP: spmd.V, TP: spmd.I}),
-            "rope_cache": SpmdType({DP: rope_cache_dp, TP: spmd.I}),
+            "x": _vision_activation_placement(),
+            "rope_cache": _vision_activation_placement(dp=rope_cache_dp),
         },
         in_dst_shardings={
-            "x": SpmdType({DP: spmd.V, TP: spmd.R}),
-            "rope_cache": SpmdType({DP: rope_cache_dp, TP: spmd.R}),
+            "x": _vision_activation_placement(tp=spmd.R),
+            "rope_cache": _vision_activation_placement(
+                dp=rope_cache_dp,
+                tp=spmd.R,
+            ),
         },
     )
     block.attn.wq.sharding_config = vision_colwise_config(input_tp=spmd.R)
     block.attn.wk.sharding_config = vision_colwise_config(input_tp=spmd.R)
     block.attn.wv.sharding_config = vision_colwise_config(input_tp=spmd.R)
     block.attn.proj.sharding_config = vision_scaled_bias_rowwise_config()
-    set_gqa_inner_attention_local_map(block.attn.inner_attention)
+    attention_layout = _vision_activation_placement(tp=spmd.S(1))
+    attention_grad_layout = SpmdType({DP: spmd.V, CP: spmd.P, TP: spmd.S(1)})
+    block.attn.inner_attention.sharding_config = ShardingConfig(
+        in_src_shardings={
+            "q_TNH": attention_layout,
+            "k_TNH": attention_layout,
+            "v_TNH": attention_layout,
+        },
+        in_dst_shardings={
+            "q_TNH": attention_layout,
+            "k_TNH": attention_layout,
+            "v_TNH": attention_layout,
+        },
+        out_src_shardings=attention_layout,
+        local_map=LocalMapConfig(
+            in_grad_placements=(attention_grad_layout,) * 3,
+        ),
+    )
 
     block.mlp.fc1.sharding_config = vision_colwise_config()
     block.mlp.fc2.sharding_config = vision_scaled_bias_rowwise_config()

torchtitan/models/muse_glimmer/README.md

@@ -46,12 +46,13 @@ configs are built in [`__init__.py`](./__init__.py).
 
 The `*_mm` flavors make `MuseGlimmerModel` own a vision encoder + adapter
 ([`vision_encoder.py`](./vision_encoder.py)) and run them inside `forward`,
-scattering the projected features into token embeddings at `vision_mask`
-positions.
+building packed vision features and gathering them into token embeddings using
+bank indices prepared before token sharding.
 
 ## Parallelism support
 
 Parallelism is applied in [`parallelize.py`](./parallelize.py) (FSDP, HSDP, TP,
-SP, CP, `torch.compile`, and PP). Sharding for Muse Glimmer-specific modules is defined
-in [`sharding.py`](./sharding.py). CP and PP are not yet supported for the
-multimodal path.
+SP, CP, `torch.compile`, and PP). Sharding for Muse Glimmer-specific modules is
+defined in [`sharding.py`](./sharding.py). The multimodal path supports CP,
+TP+CP+SP, and TP+CP+PP+SP while keeping vision computation replicated across
+CP.

torchtitan/models/muse_glimmer/__init__.py

@@ -19,6 +19,7 @@
 )
 from torchtitan.models.common.attention import QKVLinear, VarlenAttention
 from torchtitan.models.common.config_utils import get_attention_config, make_ffn_config
+from torchtitan.models.common.multimodal import VisionScatter
 from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm
 from torchtitan.models.common.param_init import depth_scaled_std
 from torchtitan.models.common.vision_encoder import (
@@ -372,13 +373,15 @@ def _muse_glimmer_config(
     # apply a scaleless norm. Left as None for the text-only model.
     vision_projection = None
     perception_emb_norm = None
+    vision_scatter = None
     if vision_adapter_dim is not None:
         vision_projection = Linear.Config(
             in_features=vision_adapter_dim,
             out_features=dim,
             param_init=_LINEAR_INIT,
         )
         perception_emb_norm = _scaleless_norm(dim, _NORM_EPS)
+        vision_scatter = VisionScatter.Config()
 
     # When the model owns the vision stack, fill the encoder/adapter sharding
     # configs so ``model.parallelize`` applies their TP.
@@ -421,6 +424,7 @@ def _muse_glimmer_config(
         ),
         vision_projection=vision_projection,
         perception_emb_norm=perception_emb_norm,
+        vision_scatter=vision_scatter,
         vision_encoder=vision_encoder,
         vision_adapter=vision_adapter,
     )

torchtitan/models/muse_glimmer/config_registry.py

@@ -152,13 +152,12 @@ def muse_glimmer_debugmodel_mm() -> Trainer.Config:
     Trains the ``debugmodel_mm`` flavor (debug text decoder that owns a
     scaled-down vision encoder + adapter) end-to-end on the ``cc12m-test`` local
     tar fixture. The shared Grain data pipeline emits packed ``pixel_values`` +
-    ``grid_thw`` + ``special_tokens``; the model derives the vision-placeholder
-    mask from ``special_tokens``. Vision-placeholder positions are already
-    ``IGNORE_INDEX`` in the labels, so a standard ``CrossEntropyLoss`` (wrapped in
-    ``ChunkedLossWrapper``) is used.
+    ``grid_thw`` + ``special_tokens``; preprocessing builds packed-bank
+    indices from the image placeholder token. Vision-placeholder positions are
+    already ``IGNORE_INDEX`` in the labels, so a standard ``CrossEntropyLoss``
+    (wrapped in ``ChunkedLossWrapper``) is used.
 
-    The parallelism smoke suite covers FSDP and FSDP+TP+SP (see
-    ``build_muse_glimmer_mm_test_list``); PP and CP are multimodal follow-ups.
+    The integration smoke suite covers the TP+CP+PP+SP path.
     """
     mm_model_spec = model_registry("debugmodel_mm", attn_backend="flex")
     return Trainer.Config(