[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,
+        ),
     ]

tests/unit_tests/cpu/test_packed_vision.py

@@ -8,14 +8,20 @@
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import spmd_types as spmd
 import torch
 import torch.nn as nn
+from spmd_types.checker import typecheck
 from torch.nn.attention.flex_attention import create_mask
 
 from torchtitan.components.loss import IGNORE_INDEX
+from torchtitan.distributed.utils import get_spmd_backend, set_spmd_backend
 from torchtitan.hf_datasets.multimodal.mm_collator import MultiModalCollator
 from torchtitan.models.common.linear import Linear
-from torchtitan.models.common.multimodal import scatter_vision_embeds
+from torchtitan.models.common.multimodal import (
+    gather_vision_embeds,
+    scatter_vision_embeds,
+)
 from torchtitan.models.common.nn_modules import LayerNorm
 from torchtitan.models.common.vision_encoder import create_block_diagonal_mask
 from torchtitan.models.kimi_k2_7.vision_encoder import (
@@ -230,6 +236,90 @@ def test_scatter_vision_embeds_uses_packed_layout(self) -> None:
         torch.testing.assert_close(result_TD[3:], vision_TD[2:])
         torch.testing.assert_close(result_TD[2], torch.zeros(2))
 
+    def test_gather_vision_embeds_uses_packed_bank_indices(self) -> None:
+        inputs_TD = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]],
+            requires_grad=True,
+        )
+        vision_bank_VD = torch.tensor(
+            [[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]],
+            dtype=torch.float64,
+            requires_grad=True,
+        )
+        vision_bank_indices_T = torch.tensor([-1, 2, 0, 2])
+
+        result_TD = gather_vision_embeds(
+            inputs_TD,
+            vision_bank_VD=vision_bank_VD,
+            vision_bank_indices_T=vision_bank_indices_T,
+        )
+
+        expected_TD = torch.tensor(
+            [[1.0, 2.0], [30.0, 31.0], [10.0, 11.0], [30.0, 31.0]]
+        )
+        self.assertEqual(result_TD.dtype, inputs_TD.dtype)
+        torch.testing.assert_close(result_TD, expected_TD)
+        result_TD.sum().backward()
+        torch.testing.assert_close(
+            inputs_TD.grad,
+            torch.tensor([[1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]),
+        )
+        torch.testing.assert_close(
+            vision_bank_VD.grad,
+            torch.tensor([[1.0, 1.0], [0.0, 0.0], [2.0, 2.0]], dtype=torch.float64),
+        )
+
+    def test_gather_vision_embeds_accepts_empty_bank(self) -> None:
+        inputs_TD = torch.randn(4, 3)
+
+        result_TD = gather_vision_embeds(
+            inputs_TD,
+            vision_bank_VD=torch.empty(0, 3),
+            vision_bank_indices_T=torch.full((4,), -1),
+        )
+
+        torch.testing.assert_close(result_TD, inputs_TD)
+
+    def test_gather_vision_embeds_preserves_token_sharding(self) -> None:
+        dp_axis = spmd.MeshAxis.of(2, 4)
+        cp_axis = spmd.MeshAxis.of(2, 2)
+        tp_axis = spmd.MeshAxis.of(2, 1)
+        inputs_TD = torch.zeros(2, 2)
+        vision_bank_VD = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        vision_bank_indices_T = torch.tensor([0, 1])
+        token_type = {dp_axis: spmd.V, cp_axis: spmd.V, tp_axis: spmd.V}
+        token_spec = spmd.PartitionSpec((dp_axis, cp_axis, tp_axis), None)
+        index_spec = spmd.PartitionSpec((dp_axis, cp_axis, tp_axis))
+
+        previous_backend = get_spmd_backend()
+        set_spmd_backend("spmd_types")
+        try:
+            with spmd.set_current_mesh(
+                {"dp": dp_axis, "cp": cp_axis, "tp": tp_axis},
+                local_axes=(dp_axis,),
+            ):
+                spmd.assert_type(inputs_TD, token_type, token_spec)
+                spmd.assert_type(
+                    vision_bank_VD,
+                    {dp_axis: spmd.V, cp_axis: spmd.R, tp_axis: spmd.R},
+                )
+                spmd.assert_type(
+                    vision_bank_indices_T,
+                    token_type,
+                    index_spec,
+                )
+                with typecheck(strict_mode="strict", local=False):
+                    result_TD = gather_vision_embeds(
+                        inputs_TD,
+                        vision_bank_VD=vision_bank_VD,
+                        vision_bank_indices_T=vision_bank_indices_T,
+                    )
+        finally:
+            set_spmd_backend(previous_backend)
+
+        self.assertEqual(spmd.get_local_type(result_TD), token_type)
+        self.assertEqual(spmd.get_partition_spec(result_TD), token_spec)
+
 
 if __name__ == "__main__":
     unittest.main()

torchtitan/models/common/multimodal.py

@@ -6,10 +6,10 @@
 
 """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
+``gather_vision_embeds`` support gather-based fusion by carrying an absolute
+packed-bank row for every placeholder token.
 """
 
 import contextlib
@@ -96,6 +96,39 @@ 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)
+
+
+def gather_vision_embeds(
+    inputs_TD: torch.Tensor,
+    *,
+    vision_bank_VD: torch.Tensor,
+    vision_bank_indices_T: torch.Tensor,
+) -> torch.Tensor:
+    """Gather packed vision features into their placeholder token positions."""
+    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)]
+    # The vision bank is DP-local, so global propagation through where omits
+    # DP from the token PartitionSpec. Validate locally, then restore the exact
+    # token layout at the fusion boundary.
+    with spmd.local():
+        fused_TD = torch.where(is_vision_T1, gathered_TD, inputs_TD)
+    if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
+        spmd.assert_type_like(fused_TD, inputs_TD)
+    return fused_TD
+
+
 def scatter_vision_embeds(
     inputs_embeds: torch.Tensor,
     *,