From 6732c3b55876f0b4f6ce0e7c619b94c0ca85a474 Mon Sep 17 00:00:00 2001
From: Jin Soo Ihm <jinsooihm@meta.com>
Date: Mon, 31 Aug 2026 15:01:36 -0700
Subject: [PATCH 1/3] [muse_glimmer] Add multimodal context parallelism

Build packed vision source indices before context-parallel token sharding and carry them through the same load-balancing permutation. Keep vision compute replicated across CP, then use a stateless sharding boundary to gather each local token shard while preserving TP and FSDP gradient semantics.

Generalize the shared vision attention layouts, keep multimodal embedding preparation on pipeline stage 0, and support the combined TP+CP+PP+SP topology without changing existing non-CP behavior or checkpoint keys.

Test Plan: 71 focused CPU tests; 2 distributed GPU tests; 10-step CP2xPP2 and TP2xCP2xPP2+SP real-process-group runs; exact 10-step deterministic loss and grad-norm parity; repository pre-commit hooks excluding the separately documented optional-dependency Pyrefly blocker.
---
 tests/integration_tests/models.py             |   7 +
 torchtitan/models/common/multimodal.py        |  46 ++++-
 .../models/common/vision_encoder_sharding.py  | 180 ++++++++++++----
 torchtitan/models/muse_glimmer/README.md      |  11 +-
 torchtitan/models/muse_glimmer/__init__.py    |   4 +
 .../models/muse_glimmer/config_registry.py    |  11 +-
 torchtitan/models/muse_glimmer/model.py       | 178 ++++++++--------
 torchtitan/models/muse_glimmer/parallelize.py |  25 +--
 torchtitan/models/muse_glimmer/sharding.py    | 192 +++++++-----------
 .../models/muse_glimmer/vision_encoder.py     |   6 +-
 torchtitan_recipes/tests/models.py            |  18 ++
 11 files changed, 403 insertions(+), 275 deletions(-)

diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py
index 0cb83176fc..139340b3cb 100755
--- a/tests/integration_tests/models.py
+++ b/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,
+        ),
     ]
diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py
index f318eaada2..d07443ecef 100644
--- a/torchtitan/models/common/multimodal.py
+++ b/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,
     *,
diff --git a/torchtitan/models/common/vision_encoder_sharding.py b/torchtitan/models/common/vision_encoder_sharding.py
index 52944a7fe8..7088dbf294 100644
--- a/torchtitan/models/common/vision_encoder_sharding.py
+++ b/torchtitan/models/common/vision_encoder_sharding.py
@@ -20,10 +20,32 @@
 
 
 DP = MeshAxisName.DP
+CP = MeshAxisName.CP
 TP = MeshAxisName.TP
 
 
-def multimodal_input_sharding() -> dict[str, SpmdType]:
+def _vision_state_placement(
+    *,
+    tp: spmd.PerMeshAxisSpmdType,
+    include_cp_axis: bool = False,
+) -> SpmdType:
+    if include_cp_axis:
+        return SpmdType({DP: spmd.R, CP: spmd.R, TP: tp})
+    return SpmdType({DP: spmd.R, TP: tp})
+
+
+def _vision_activation_placement(
+    *,
+    dp: spmd.PerMeshAxisSpmdType = spmd.V,
+    tp: spmd.PerMeshAxisSpmdType = spmd.I,
+    include_cp_axis: bool = False,
+) -> SpmdType:
+    if include_cp_axis:
+        return SpmdType({DP: dp, CP: spmd.R, TP: tp})
+    return SpmdType({DP: dp, TP: tp})
+
+
+def multimodal_input_sharding(*, include_cp_axis: bool = False) -> dict[str, SpmdType]:
     """SPMD layouts for VLM vision inputs (folded into a model's input_sharding).
 
     The vision tensors are DP-local (``V@DP``) -- each DP rank owns its own
@@ -31,7 +53,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(include_cp_axis=include_cp_axis)
     return {
         "pixel_values": layout,
         "pixel_values_videos": layout,
@@ -40,68 +62,95 @@ def multimodal_input_sharding() -> dict[str, SpmdType]:
     }
 
 
-def invariant_norm_config() -> ShardingConfig:
+def invariant_norm_config(*, include_cp_axis: bool = False) -> 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, include_cp_axis=include_cp_axis
+            ),
+            "bias": _vision_state_placement(tp=spmd.I, include_cp_axis=include_cp_axis),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(include_cp_axis=include_cp_axis),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(include_cp_axis=include_cp_axis),
         },
-        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(include_cp_axis=include_cp_axis),
+        out_dst_shardings=_vision_activation_placement(include_cp_axis=include_cp_axis),
     )
 
 
-def vision_invariant_linear_config() -> ShardingConfig:
+def vision_invariant_linear_config(*, include_cp_axis: bool = False) -> 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, include_cp_axis=include_cp_axis
+            ),
+            "bias": _vision_state_placement(tp=spmd.I, include_cp_axis=include_cp_axis),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(include_cp_axis=include_cp_axis),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.I}),
+            "input": _vision_activation_placement(include_cp_axis=include_cp_axis),
         },
-        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(include_cp_axis=include_cp_axis),
+        out_dst_shardings=_vision_activation_placement(include_cp_axis=include_cp_axis),
     )
 
 
 def vision_colwise_config(
-    *, input_tp: spmd.PerMeshAxisSpmdType = spmd.I
+    *,
+    input_tp: spmd.PerMeshAxisSpmdType = spmd.I,
+    include_cp_axis: bool = False,
 ) -> ShardingConfig:
     """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), include_cp_axis=include_cp_axis
+            ),
+            "bias": _vision_state_placement(
+                tp=spmd.S(0), include_cp_axis=include_cp_axis
+            ),
         },
         in_src_shardings={
-            "input": SpmdType({DP: spmd.V, TP: input_tp}),
+            "input": _vision_activation_placement(
+                tp=input_tp, include_cp_axis=include_cp_axis
+            ),
         },
         in_dst_shardings={
-            "input": SpmdType({DP: spmd.V, TP: spmd.R}),
+            "input": _vision_activation_placement(
+                tp=spmd.R, include_cp_axis=include_cp_axis
+            ),
         },
-        out_src_shardings=SpmdType({DP: spmd.V, TP: spmd.S(-1)}),
+        out_src_shardings=_vision_activation_placement(
+            tp=spmd.S(-1), include_cp_axis=include_cp_axis
+        ),
     )
 
 
-def vision_scaled_bias_rowwise_config() -> ShardingConfig:
+def vision_scaled_bias_rowwise_config(
+    *, include_cp_axis: bool = False
+) -> 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), include_cp_axis=include_cp_axis
+    )
+    input_grad_layout = (
+        SpmdType({DP: spmd.V, CP: spmd.P, TP: spmd.S(1)})
+        if include_cp_axis
+        else input_layout
+    )
     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), include_cp_axis=include_cp_axis
+            ),
+            "bias": _vision_state_placement(tp=spmd.R, include_cp_axis=include_cp_axis),
         },
         in_src_shardings={
             "input": input_layout,
@@ -109,9 +158,11 @@ def vision_scaled_bias_rowwise_config() -> ShardingConfig:
         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, include_cp_axis=include_cp_axis
+        ),
+        out_dst_shardings=_vision_activation_placement(include_cp_axis=include_cp_axis),
+        local_map=LocalMapConfig(in_grad_placements=(input_grad_layout,)),
     )
 
 
@@ -119,26 +170,69 @@ def set_vision_transformer_block_sharding_config(
     block: "VisionTransformerBlock.Config",
     *,
     rope_cache_dp: spmd.PerMeshAxisSpmdType,
+    include_cp_axis: bool = False,
 ) -> None:
     """Set TP sharding for the common vision transformer block."""
-    block.norm1.sharding_config = invariant_norm_config()
-    block.norm2.sharding_config = invariant_norm_config()
+    block.norm1.sharding_config = invariant_norm_config(include_cp_axis=include_cp_axis)
+    block.norm2.sharding_config = invariant_norm_config(include_cp_axis=include_cp_axis)
 
     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(include_cp_axis=include_cp_axis),
+            "rope_cache": _vision_activation_placement(
+                dp=rope_cache_dp, include_cp_axis=include_cp_axis
+            ),
         },
         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, include_cp_axis=include_cp_axis
+            ),
+            "rope_cache": _vision_activation_placement(
+                dp=rope_cache_dp,
+                tp=spmd.R,
+                include_cp_axis=include_cp_axis,
+            ),
         },
     )
-    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)
-
-    block.mlp.fc1.sharding_config = vision_colwise_config()
-    block.mlp.fc2.sharding_config = vision_scaled_bias_rowwise_config()
+    block.attn.wq.sharding_config = vision_colwise_config(
+        input_tp=spmd.R, include_cp_axis=include_cp_axis
+    )
+    block.attn.wk.sharding_config = vision_colwise_config(
+        input_tp=spmd.R, include_cp_axis=include_cp_axis
+    )
+    block.attn.wv.sharding_config = vision_colwise_config(
+        input_tp=spmd.R, include_cp_axis=include_cp_axis
+    )
+    block.attn.proj.sharding_config = vision_scaled_bias_rowwise_config(
+        include_cp_axis=include_cp_axis
+    )
+    if include_cp_axis:
+        attention_layout = _vision_activation_placement(
+            tp=spmd.S(1), include_cp_axis=True
+        )
+        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,
+            ),
+        )
+    else:
+        set_gqa_inner_attention_local_map(block.attn.inner_attention)
+
+    block.mlp.fc1.sharding_config = vision_colwise_config(
+        include_cp_axis=include_cp_axis
+    )
+    block.mlp.fc2.sharding_config = vision_scaled_bias_rowwise_config(
+        include_cp_axis=include_cp_axis
+    )
diff --git a/torchtitan/models/muse_glimmer/README.md b/torchtitan/models/muse_glimmer/README.md
index 0906884a6a..2cac465068 100644
--- a/torchtitan/models/muse_glimmer/README.md
+++ b/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.
diff --git a/torchtitan/models/muse_glimmer/__init__.py b/torchtitan/models/muse_glimmer/__init__.py
index 0868762c3f..0d784fff23 100644
--- a/torchtitan/models/muse_glimmer/__init__.py
+++ b/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,6 +373,7 @@ 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,
@@ -379,6 +381,7 @@ def _muse_glimmer_config(
             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,
     )
diff --git a/torchtitan/models/muse_glimmer/config_registry.py b/torchtitan/models/muse_glimmer/config_registry.py
index 85bb676c38..0f8c883adf 100644
--- a/torchtitan/models/muse_glimmer/config_registry.py
+++ b/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(
diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py
index ea7c4cb65c..b13e746527 100644
--- a/torchtitan/models/muse_glimmer/model.py
+++ b/torchtitan/models/muse_glimmer/model.py
@@ -7,7 +7,6 @@
 from dataclasses import dataclass
 from typing import Any, cast
 
-import spmd_types as spmd
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -16,7 +15,7 @@
 from torchtitan.config import ParallelismConfig
 from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.distributed.spmd_types import annotate_input_spmd_types
-from torchtitan.distributed.utils import get_spmd_backend, is_in_batch_invariant_mode
+from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.models.common.attention import (
     AttentionMasksType,
     create_attention_mask,
@@ -29,13 +28,16 @@
     VarlenAttention,
 )
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
-from torchtitan.models.common.decoder_sharding import decoder_input_sharding
+from torchtitan.models.common.decoder_sharding import (
+    decoder_input_sharding,
+    token_id_placement,
+)
 from torchtitan.models.common.embedding import Embedding
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.multimodal import (
-    get_vision_positions,
+    build_vision_bank_indices,
     multimodal_context,
-    scatter_vision_embeds,
+    VisionScatter,
 )
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.vision_encoder_sharding import multimodal_input_sharding
@@ -278,13 +280,12 @@ class Config(Decoder.Config):
         # Dataclass fields are invariant, so pyrefly flags the (intentional) override.
         # pyrefly: ignore [bad-override]
         tok_embeddings: EmbeddingWithNorm.Config
-        # Optional LLM-side multimodal injection. When set, encoded vision
-        # features (already adapter-projected to ``vision_projection`` in_features)
-        # are projected to ``dim``, scaleless-normed, and scattered into the token
-        # embeddings at masked positions. Both default to None for the text-only
-        # model, leaving the text path untouched.
+        # Optional LLM-side multimodal injection. Preprocessing builds absolute
+        # packed-bank indices before CP; VisionScatter gathers the rows needed
+        # by each local token shard after projection and normalization.
         vision_projection: Linear.Config | None = None
         perception_emb_norm: RMSNorm.Config | None = None
+        vision_scatter: VisionScatter.Config | None = None
         # Optional owned vision stack. When set, ``MuseGlimmerModel`` builds the encoder
         # + adapter as submodules and runs them inside ``forward`` (from padded
         # ``pixel_values`` + ``grid_thw``), mirroring qwen3_5's
@@ -309,7 +310,6 @@ def update_from_config(
                     "Context Parallel only supports SDPA and FlexAttention. "
                     "Varlen attention is not supported with CP."
                 )
-
             from .sharding import set_muse_glimmer_sharding_config
 
             set_muse_glimmer_sharding_config(
@@ -361,6 +361,9 @@ def __init__(self, config: "MuseGlimmerModel.Config") -> None:
             if config.perception_emb_norm is not None
             else None
         )
+        self.vision_scatter = (
+            config.vision_scatter.build() if config.vision_scatter is not None else None
+        )
         # Owned vision stack (None unless a multimodal flavor configured it). When
         # present, ``forward`` runs encoder->adapter on packed pixel_values.
         self.vision_encoder = (
@@ -377,20 +380,58 @@ def preprocess_inputs(
         parallel_dims: ParallelDims,
         parallelism: ParallelismConfig,
     ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
-        """Build masks, CP-shard, SPMD-annotate, and return the batch."""
+        """Build first-stage vision-bank indices and masks, then shard the batch."""
         # Function-local import avoids a circular import.
         from torchtitan.distributed.context_parallel.api import (
             prepare_context_parallel_input,
         )
 
         batch: dict[str, Any] = dict(input_dict)
+        pixel_values = batch.get("pixel_values")
+        grid_thw = batch.get("grid_thw")
+        pixel_values_videos = batch.get("pixel_values_videos")
+        grid_thw_videos = batch.get("grid_thw_videos")
+        special_tokens = batch.get("special_tokens")
+        if pixel_values_videos is not None or grid_thw_videos is not None:
+            raise NotImplementedError(
+                "Muse Glimmer vision encoder does not support video inputs."
+            )
+        if pixel_values is not None:
+            vision_encoder_config = cast(
+                MuseGlimmerModel.Config, self.config
+            ).vision_encoder
+            if vision_encoder_config is None:
+                raise ValueError(
+                    "pixel_values were provided but the model config has no "
+                    "vision_encoder configured."
+                )
+            if grid_thw is None:
+                raise ValueError(
+                    "pixel_values were provided but grid_thw was not provided."
+                )
+            if special_tokens is None or "image_id" not in special_tokens:
+                raise ValueError(
+                    "pixel_values were provided but special_tokens with an "
+                    "'image_id' entry was not provided."
+                )
+            if self.tok_embeddings is not None:
+                batch["vision_bank_indices_T"] = build_vision_bank_indices(
+                    batch["input"],
+                    placeholder_id=special_tokens["image_id"],
+                )
+        batch.pop("special_tokens", None)
+
         positions = batch.get("positions", None)
         if positions is not None:
             inner = getattr(self.config.first_attention, "inner_attention", None)
             if isinstance(inner, (FlexAttention.Config, VarlenAttention.Config)):
                 batch["attention_masks"] = self.get_attention_masks(positions=positions)
 
-        input_sharding = {**decoder_input_sharding(), **multimodal_input_sharding()}
+        input_sharding = {
+            **decoder_input_sharding(),
+            **multimodal_input_sharding(include_cp_axis=True),
+        }
+        input_sharding["vision_bank_indices_T"] = token_id_placement()
         if parallel_dims.cp_enabled:
             batch = prepare_context_parallel_input(
                 batch,
@@ -422,6 +463,29 @@ def _get_vision_features(
         )
         return feats
 
+    def _prepare_multimodal_embeds(
+        self,
+        h_TD: torch.Tensor,
+        *,
+        pixel_values: torch.Tensor | None,
+        grid_thw: torch.Tensor | None,
+        vision_bank_indices_T: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Build and inject image embeddings on the embedding pipeline stage."""
+        if pixel_values is None:
+            return h_TD
+        assert grid_thw is not None
+        assert vision_bank_indices_T is not None
+        assert self.vision_projection is not None
+        assert self.perception_emb_norm is not None
+        assert self.vision_scatter is not None
+
+        vision_features_VD = self._get_vision_features(pixel_values, grid_thw)
+        vision_bank_VD = self.perception_emb_norm(
+            self.vision_projection(vision_features_VD)
+        )
+        return self.vision_scatter(h_TD, vision_bank_VD, vision_bank_indices_T)
+
     def forward(
         self,
         tokens: torch.Tensor,
@@ -432,93 +496,37 @@ def forward(
         grid_thw: torch.Tensor | None = None,
         pixel_values_videos: torch.Tensor | None = None,
         grid_thw_videos: torch.Tensor | None = None,
-        special_tokens: dict[str, int] | None = None,
+        vision_bank_indices_T: torch.Tensor | None = None,
     ):
+        # Video inputs are rejected by preprocess_inputs.
+        del pixel_values_videos, grid_thw_videos
+
         # Embedding stage: embed tokens (the scaleless norm is bundled inside
         # tok_embeddings) and inject vision features before the decoder layers.
         # On non-embedding pipeline stages tok_embeddings is None and the input
         # is already hidden states, so injection is skipped there.
         with multimodal_context():
             if self.tok_embeddings is not None:
-                h = self.tok_embeddings(tokens)
-                # The model owns the encoder: when packed pixel_values are passed,
-                # run encoder->adapter here to produce the features for injection.
-                # The placeholder mask is derived from tokens + special_tokens.
-                # TODO: Video is not implemented in the training forward. The
-                # encoder itself is video-capable; this path just lacks the
-                # video-specific glue that the image path (above) doesn't need:
-                #   1. Temporal frame packing -- group `patch_temporal` frames per
-                #      patch.
-                #   2. Spatial avg-pool compression between encoder and adapter
-                #      (pool_factor from compression_ratio); the image path goes
-                #      encoder->adapter directly with no compression.
-                #   3. Video grid sizing (with compression_ratio / max_num_tokens)
-                #      vs image grid sizing.
-                #   4. A separate video placeholder token + mask instead of the
-                #      image_id used below.
-                if pixel_values_videos is not None or grid_thw_videos is not None:
-                    raise NotImplementedError(
-                        "Muse Glimmer vision encoder does not support video inputs."
-                    )
-                if pixel_values is not None:
-                    if self.vision_encoder is None:
-                        raise ValueError(
-                            "pixel_values were provided but the model has no "
-                            "vision_encoder configured."
-                        )
-                    if grid_thw is None:
-                        raise ValueError(
-                            "pixel_values were provided but grid_thw was not provided."
-                        )
-                    if special_tokens is None or "image_id" not in special_tokens:
-                        raise ValueError(
-                            "pixel_values were provided but special_tokens with an "
-                            "'image_id' entry was not provided."
-                        )
-                    vision_features = self._get_vision_features(pixel_values, grid_thw)
-                    if (
-                        self.vision_projection is None
-                        or self.perception_emb_norm is None
-                    ):
-                        raise ValueError(
-                            "pixel_values were provided but the model has no "
-                            "vision_projection/perception_emb_norm configured."
-                        )
-                    vision_embeds = self.perception_emb_norm(
-                        self.vision_projection(vision_features)
-                    )
-                    downsample_factor = self.vision_encoder.downsample_factor
-                    num_tokens_per_item = (grid_thw[:, 1] // downsample_factor) * (
-                        grid_thw[:, 2] // downsample_factor
-                    )
-                    vision_positions = get_vision_positions(
-                        tokens,
-                        num_tokens_per_item,
-                        special_tokens["image_id"],
-                    )
-                    h = scatter_vision_embeds(
-                        h,
-                        vision_embeds=vision_embeds,
-                        vision_positions=vision_positions,
-                    )
+                h_TD = self.tok_embeddings(tokens)
+                h_TD = self._prepare_multimodal_embeds(
+                    h_TD,
+                    pixel_values=pixel_values,
+                    grid_thw=grid_thw,
+                    vision_bank_indices_T=vision_bank_indices_T,
+                )
             else:
-                h = tokens
-
-        if get_spmd_backend() == "spmd_types":
-            # The scatter restores a token-aligned tensor, so text-model DP
-            # resumes sharding the leading token dimension.
-            spmd.assert_type(h, {"dp": spmd.S(0), "tp": spmd.R})
+                h_TD = tokens
 
         for layer in self.layers.values():
-            h = layer(h, attention_masks, positions)
+            h_TD = layer(h_TD, attention_masks, positions)
 
-        h = self.norm(h) if self.norm is not None else h
+        h_TD = self.norm(h_TD) if self.norm is not None else h_TD
 
         # _skip_lm_head is an attribute (not a kwarg) because PP backward calls
         # .requires_grad on all stage inputs, which fails on bool kwargs.
         if self._skip_lm_head:
-            return h
-        return self.lm_head(h) if self.lm_head is not None else h
+            return h_TD
+        return self.lm_head(h_TD) if self.lm_head is not None else h_TD
 
     def get_attention_masks(
         self,
diff --git a/torchtitan/models/muse_glimmer/parallelize.py b/torchtitan/models/muse_glimmer/parallelize.py
index d8410a4ef5..46ca6830ed 100644
--- a/torchtitan/models/muse_glimmer/parallelize.py
+++ b/torchtitan/models/muse_glimmer/parallelize.py
@@ -50,10 +50,6 @@ def parallelize_muse_glimmer(
     has_vision = model.vision_encoder is not None
     if has_vision:
         assert model.vision_adapter is not None
-        if parallel_dims.cp_enabled:
-            raise NotImplementedError(
-                "context parallel is not supported for the Muse Glimmer vision encoder."
-            )
         if parallel_dims.tp_enabled:
             # pyrefly: ignore [missing-attribute]
             vision_num_heads = model.vision_encoder.num_heads
@@ -156,12 +152,10 @@ def pipeline_muse_glimmer(
     Muse Glimmer's vision modules, so without this they would be pruned to ``None`` on
     every stage.
 
-    Muse Glimmer's owned vision stack runs inside ``MuseGlimmerModel.forward`` on the embedding
-    stage (where ``tok_embeddings`` lives): the encoder + adapter encode raw
-    images, and ``vision_projection`` + ``perception_emb_norm`` scatter the result
-    into the token embeddings. All present vision modules must therefore live on
-    stage 0. (For the standalone-encoder flavor only ``vision_projection`` +
-    ``perception_emb_norm`` exist; the per-module presence check handles that.)
+    Muse Glimmer's owned vision stack runs inside ``MuseGlimmerModel.forward``
+    on the embedding stage: the encoder and adapter build raw features, the
+    projection and norm build the packed bank, and the stateless scatter fuses
+    it into token embeddings. All present vision modules live on stage 0.
     """
     import dataclasses
 
@@ -173,11 +167,9 @@ def pipeline_muse_glimmer(
 
     # NOTE: We cannot delegate to the generic ``pipeline_vlm`` here. That helper
     # only injects a single ``vision_encoder`` FQN into stage 0; Muse Glimmer owns
-    # a multi-module vision stack (vision_encoder, vision_adapter,
-    # vision_projection, perception_emb_norm) whose membership varies by flavor
-    # (the standalone-encoder flavor has only vision_projection +
-    # perception_emb_norm). We therefore replicate ``pipeline_vlm``'s structure but
-    # inject the full, per-module presence-checked stack instead.
+    # a multi-module vision stack whose membership varies by flavor. We therefore
+    # replicate ``pipeline_vlm``'s structure but inject the full,
+    # per-module presence-checked stack instead.
     if parallelism.module_fqns_per_model_part is None:
         (
             num_virtual_stages,
@@ -194,7 +186,7 @@ def pipeline_muse_glimmer(
         # vision encoder, bump
         # parallelism.pipeline_parallel_first_stage_less_layers to rebalance.
         # Prepend in data-flow order so the resulting stage-0 list reads
-        # encoder -> adapter -> projection -> emb_norm -> tok_embeddings.
+        # encoder -> adapter -> projection -> emb_norm -> scatter -> embeddings.
         vision_fqns = [
             fqn
             for fqn in (
@@ -202,6 +194,7 @@ def pipeline_muse_glimmer(
                 "vision_adapter",
                 "vision_projection",
                 "perception_emb_norm",
+                "vision_scatter",
             )
             if getattr(model, fqn, None) is not None
         ]
diff --git a/torchtitan/models/muse_glimmer/sharding.py b/torchtitan/models/muse_glimmer/sharding.py
index ed09f61ed6..711d62523b 100644
--- a/torchtitan/models/muse_glimmer/sharding.py
+++ b/torchtitan/models/muse_glimmer/sharding.py
@@ -36,9 +36,46 @@
 
 
 DP = MeshAxisName.DP
+CP = MeshAxisName.CP
 TP = MeshAxisName.TP
 
 
+def _sequence_parallel_index_placement() -> SpmdType:
+    return SpmdType(
+        {DP: spmd.V, CP: spmd.V, TP: spmd.V},
+        partition_spec=spmd.PartitionSpec((DP, CP, TP)),
+    )
+
+
+def _vision_scatter_config(*, enable_sp: bool) -> ShardingConfig:
+    vision_bank_src = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.I})
+    vision_bank_dst = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.R})
+    vision_bank_indices = token_id_placement()
+    if enable_sp:
+        hidden_src = dense_sequence_parallel_placement()
+        hidden_dst = hidden_src
+        local_vision_bank_indices = _sequence_parallel_index_placement()
+    else:
+        hidden_src = dense_activation_placement(tp=spmd.I, cp=spmd.S(0))
+        hidden_dst = dense_activation_placement(tp=spmd.R, cp=spmd.S(0))
+        local_vision_bank_indices = vision_bank_indices
+
+    return ShardingConfig(
+        in_src_shardings={
+            "inputs_TD": hidden_src,
+            "vision_bank_VD": vision_bank_src,
+            "vision_bank_indices_T": vision_bank_indices,
+        },
+        in_dst_shardings={
+            "inputs_TD": hidden_dst,
+            "vision_bank_VD": vision_bank_dst,
+            "vision_bank_indices_T": local_vision_bank_indices,
+        },
+        out_src_shardings=hidden_dst,
+        out_dst_shardings=hidden_src,
+    )
+
+
 def set_muse_glimmer_sharding_config(
     config: "MuseGlimmerModel.Config",
     *,
@@ -46,21 +83,10 @@ def set_muse_glimmer_sharding_config(
 ) -> None:
     """Fill ``sharding_config`` on all Muse Glimmer sub-configs.
 
-    The text-only and multimodal models share the same decoder sharding; they
-    differ only in how the token embeddings flow into the first decoder layer:
-
-    * **Text-only**: the base decoder sharding is the whole story. The token
-      embeddings emit sequence-parallel (``Shard(0)``) activations that flow
-      straight into the decoder layers.
-    * **Multimodal** (``config.vision_encoder is not None``): ``MuseGlimmerModel.forward``
-      scatters vision features into the token embeddings over the full
-      ``[num_tokens]`` sequence *between* ``tok_embeddings`` and the decoder
-      layers, so the embedding output must be ``Replicate`` -- not
-      ``Shard(0)``/SP.
-      :func:`_set_multimodal_sharding` overrides the embedding + norm (and the
-      vision-injection modules) to ``Replicate``, and the layer loop gives the
-      first decoder layer a ``Replicate`` input that its attention reduce-scatters
-      back to SP. Everything else is identical to the text path.
+    Text-only and multimodal models use the same token-sharded decoder path.
+    The multimodal path keeps its packed vision bank replicated across CP and
+    TP-invariant until ``VisionScatter`` gathers rows into each local token
+    shard. The scatter boundary performs the required TP ``I -> R`` transition.
 
     All sub-configs are populated unconditionally -- ``Module.parallelize``
     filters disabled axes at runtime.
@@ -71,9 +97,7 @@ def set_muse_glimmer_sharding_config(
     for layer_cfg in config.layers:
         _set_muse_glimmer_layer_sharding(layer_cfg, enable_sp=enable_sp)
 
-    # Multimodal-only override: re-point the embedding (and the first decoder
-    # layer) at Replicate so the vision scatter can index the full sequence.
-    # No-op for the text model.
+    # Configure the replicated vision bank and token-local scatter boundary.
     if config.vision_encoder is not None:
         _set_multimodal_sharding(config, enable_sp=enable_sp)
 
@@ -112,87 +136,18 @@ def _set_multimodal_sharding(
     *,
     enable_sp: bool,
 ) -> None:
-    """Override the text-path sharding for the multimodal (vision) model.
-
-    ``MuseGlimmerModel.forward`` scatters vision features into the token embeddings
-    (masked index over the full ``[num_tokens]``) between ``tok_embeddings`` and
-    the decoder layers, so that activation must be ``Replicate`` -- not
-    ``Shard(0)``/SP. This re-points the embedding children at ``Replicate``
-    outputs, marks the vision-injection modules ``Replicate`` so the whole vision
-    path stays DTensor-consistent, and (under SP) re-shards the first decoder
-    layer to take that ``Replicate`` input -- its rowwise ``wo`` reduce-scatters
-    back to ``Shard(0)``, restoring SP activations for every later layer. Mirrors
-    qwen3_5's multimodal sharding overrides.
-    """
-    replicate = dense_activation_placement(tp=spmd.R, cp=spmd.S(0))
-    # The vision encoder + adapter emit TP-invariant activations (the common
-    # vision_encoder_sharding helpers flow {DP: V, TP: I}). The LLM-side injection
-    # modules only promote the TP axis I->R so the features become TP-Replicate for
-    # the scatter; the DP axis stays V (per-image, local under multimodal_context).
-    # DP is deliberately NOT redistributed to S(0): config-based redistribution
-    # cannot move an axis to/from spmd.V, so (like qwen3_5's scatter helper) the
-    # raw scatter below writes the V vision rows into the S(0) text positions
-    # per-rank instead.
-    vision_invariant = SpmdType({DP: spmd.V, TP: spmd.I})
-    vision_tp_replicate = SpmdType({DP: spmd.V, TP: spmd.R})
-    emb_cfg = config.tok_embeddings
-
-    # Embedding output Replicate (vs Shard(0)/SP): the vision scatter needs the
-    # full sequence. Vocab-parallel Embedding.forward runs a manual local masked
-    # lookup on a Shard(0) weight and emits a Partial sum; local_map localizes the
-    # Replicate DTensor input so the manual path sees a plain tensor (otherwise it
-    # mixes a DTensor input with the local weight), and the Partial output is
-    # all-reduced to Replicate. Mirrors the text-path tok_embeddings config in
-    # set_decoder_sharding_config.
-    emb_cfg.embedding.sharding_config = ShardingConfig(
-        state_shardings={"weight": dense_param_placement(tp=spmd.S(0))},
-        in_src_shardings={"input": token_id_placement()},
-        in_dst_shardings={"input": token_id_placement()},
-        out_src_shardings=dense_activation_placement(tp=spmd.P, cp=spmd.S(0)),
-        out_dst_shardings=replicate,
-        local_map=LocalMapConfig(in_grad_placements=None),
-    )
-    emb_cfg.norm.sharding_config = ShardingConfig(
-        in_src_shardings={"input": replicate},
-        in_dst_shardings={"input": replicate},
-        out_src_shardings=replicate,
-        out_dst_shardings=replicate,
-    )
-
-    # LLM-side vision injection (vision_projection + perception_emb_norm) consumes
-    # the adapter output, which flows TP-invariant ({DP: V, TP: I}). vision_projection
-    # promotes the TP axis I->R (single-axis) so the features are TP-Replicate; the
-    # norm keeps them there. The DP axis stays V through both, matching the vision
-    # features that the scatter writes into the Replicate text stream.
+    """Configure the vision bank and its token-sharded scatter boundary."""
     if config.vision_projection is not None:
-        config.vision_projection.sharding_config = ShardingConfig(
-            state_shardings={
-                "weight": dense_param_placement(tp=spmd.R),
-                # Harmless no-op if the projection has no bias.
-                "bias": dense_param_placement(tp=spmd.R),
-            },
-            in_src_shardings={"input": vision_invariant},
-            in_dst_shardings={"input": vision_tp_replicate},
-            out_src_shardings=vision_tp_replicate,
+        config.vision_projection.sharding_config = vision_invariant_linear_config(
+            include_cp_axis=True
         )
     if config.perception_emb_norm is not None:
-        config.perception_emb_norm.sharding_config = ShardingConfig(
-            in_src_shardings={"input": vision_tp_replicate},
-            in_dst_shardings={"input": vision_tp_replicate},
-            out_src_shardings=vision_tp_replicate,
+        config.perception_emb_norm.sharding_config = invariant_norm_config(
+            include_cp_axis=True
         )
-
-    # First-layer SP bridge: tok_embeddings now emits Replicate activations, but the
-    # rest of the decoder expects sequence-parallel activations (Shard(0)). Re-shard
-    # the first layer to take a Replicate input; its rowwise ``wo`` (output_sp)
-    # reduce-scatters back to Shard(0), restoring SP activations for every later layer.
-    # Only needed under SP -- without SP the whole decoder already uses Replicate
-    # activations. Mirrors qwen3_5's first-layer Replicate input layout.
-    if enable_sp and config.layers:
-        config.layers[0].sharding_config = ShardingConfig(
-            in_src_shardings={"x": replicate},
-            in_dst_shardings={"x": dense_sequence_parallel_placement()},
-            out_src_shardings=dense_sequence_parallel_placement(),
+    if config.vision_scatter is not None:
+        config.vision_scatter.sharding_config = _vision_scatter_config(
+            enable_sp=enable_sp
         )
 
 
@@ -258,12 +213,10 @@ def set_muse_glimmer_vision_sharding_config(
 ) -> None:
     """Fill ``sharding_config`` on the Muse Glimmer vision encoder (+ optional adapter).
 
-    All vision activations flow as TP-invariant (no sequence parallelism), exactly
-    like qwen3_5's vision encoder. The shared block/linear/norm helpers in
-    :mod:`torchtitan.models.common.vision_encoder_sharding` carry the actual TP
-    sharding (colwise q/k/v + rowwise proj, colwise fc1 + rowwise fc2, invariant
-    norms, inner-attention local_map); only the Muse-Glimmer-specific learned
-    positional grid, RoPE frequencies, and patch ``conv1`` are declared here.
+    Vision activations are invariant across TP and replicated across CP. The
+    shared block/linear/norm helpers carry TP sharding and explicit CP layouts;
+    only the Muse-specific learned positional grid, RoPE frequencies, patch
+    ``conv1``, and local permutation boundaries are declared here.
 
     Must be called BEFORE the configs are built (``config.build()``): the built
     modules copy these configs into ``Module.parallelize``.
@@ -272,44 +225,49 @@ def set_muse_glimmer_vision_sharding_config(
     # Replicate (it is bilinearly resampled per image, see _get_pos_emb).
     encoder_cfg.sharding_config = ShardingConfig(
         state_shardings={
-            "positional_embedding_vlm": SpmdType({DP: spmd.R, TP: spmd.I}),
+            "positional_embedding_vlm": SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I}),
         },
     )
     encoder_cfg.rope_freq.sharding_config = ShardingConfig(
         state_shardings={
-            "inv_freq": SpmdType({DP: spmd.R, TP: spmd.I}),
+            "inv_freq": SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I}),
         },
     )
 
     # conv1 builds ``self.conv1_linear``; sharding goes on the *config* field
     # ``conv1``. Plain pixel patches enter invariant; the (bias-free) weight stays
     # Replicate. Mirrors qwen3_5's patch_embed_proj (vision_invariant_linear_config).
-    encoder_cfg.conv1.sharding_config = vision_invariant_linear_config()
-    encoder_cfg.ln_pre.sharding_config = invariant_norm_config()
-    encoder_cfg.ln_post.sharding_config = invariant_norm_config()
+    encoder_cfg.conv1.sharding_config = vision_invariant_linear_config(
+        include_cp_axis=True
+    )
+    encoder_cfg.ln_pre.sharding_config = invariant_norm_config(include_cp_axis=True)
+    encoder_cfg.ln_post.sharding_config = invariant_norm_config(include_cp_axis=True)
 
     # Per-block TP via the shared helper (norms, q/k/v/proj, fc1/fc2, and the
     # inner-attention local_map), same as qwen3_5/kimi_k2_7. ``rope_cache`` is a
-    # per-image vision activation, so it flows {DP: V, TP: I} like kimi_k2_7.
+    # per-image vision activation, so it flows {DP: V, CP: R, TP: I}.
     set_vision_transformer_block_sharding_config(
         encoder_cfg.block,
         rope_cache_dp=spmd.V,
+        include_cp_axis=True,
     )
 
-    vision_invariant = SpmdType({DP: spmd.V, TP: spmd.I})
-    pos_param_invariant = SpmdType({DP: spmd.R, TP: spmd.I})
+    vision_invariant = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.I})
+    vision_invariant_grad = SpmdType({DP: spmd.V, CP: spmd.P, TP: spmd.I})
+    pos_param_invariant = SpmdType({DP: spmd.R, CP: spmd.R, TP: spmd.I})
+    pos_param_grad = SpmdType({DP: spmd.P, CP: spmd.P, TP: spmd.I})
     encoder_cfg.pos_embed.sharding_config = ShardingConfig(
         in_src_shardings={"pos_param": pos_param_invariant},
         in_dst_shardings={"pos_param": pos_param_invariant},
         out_src_shardings=vision_invariant,
-        local_map=LocalMapConfig(in_grad_placements=(pos_param_invariant,)),
+        local_map=LocalMapConfig(in_grad_placements=(pos_param_grad,)),
     )
     encoder_cfg.token_permute.sharding_config = ShardingConfig(
         in_src_shardings={"x": vision_invariant, "index": vision_invariant},
         in_dst_shardings={"x": vision_invariant, "index": vision_invariant},
         out_src_shardings=vision_invariant,
         local_map=LocalMapConfig(
-            in_grad_placements=(vision_invariant, vision_invariant)
+            in_grad_placements=(vision_invariant_grad, vision_invariant)
         ),
     )
 
@@ -317,7 +275,11 @@ def set_muse_glimmer_vision_sharding_config(
         # The adapter runs on the flattened 2D [tokens, dim] vision features (the
         # encoder cats per-image tokens), so the block helpers' fixed Shard(2)
         # layout is out of bounds. Keep both linears TP-invariant (dimension-
-        # agnostic); the adapter output then stays {DP: V, TP: I}, matching the
-        # LLM-side vision_projection input.
-        adapter_cfg.c_fc.sharding_config = vision_invariant_linear_config()
-        adapter_cfg.c_proj.sharding_config = vision_invariant_linear_config()
+        # agnostic); the adapter output stays {DP: V, CP: R, TP: I}, matching
+        # the LLM-side vision_projection input.
+        adapter_cfg.c_fc.sharding_config = vision_invariant_linear_config(
+            include_cp_axis=True
+        )
+        adapter_cfg.c_proj.sharding_config = vision_invariant_linear_config(
+            include_cp_axis=True
+        )
diff --git a/torchtitan/models/muse_glimmer/vision_encoder.py b/torchtitan/models/muse_glimmer/vision_encoder.py
index c6f764dd73..821162ff7b 100644
--- a/torchtitan/models/muse_glimmer/vision_encoder.py
+++ b/torchtitan/models/muse_glimmer/vision_encoder.py
@@ -50,7 +50,11 @@
 def _annotate_vision_activation_type(tensor: torch.Tensor) -> torch.Tensor:
     """Annotate a tensor created inside the vision forward."""
     if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
-        return spmd.mutate_type(tensor, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I})
+        return spmd.mutate_type(
+            tensor,
+            src=spmd.R,
+            dst={"dp": spmd.V, "cp": spmd.R, "tp": spmd.I},
+        )
     return tensor
 
 
diff --git a/torchtitan_recipes/tests/models.py b/torchtitan_recipes/tests/models.py
index 79bc268b60..25935b919c 100644
--- a/torchtitan_recipes/tests/models.py
+++ b/torchtitan_recipes/tests/models.py
@@ -343,3 +343,21 @@ def muse_glimmer_debugmodel_mm_fsdp2_tp2() -> Trainer.Config:
     config.parallelism.tensor_parallel_degree = 2
     config.training.disable_cuda_graphs = True
     return config
+
+
+def muse_glimmer_debugmodel_mm_tp2_cp2_pp2() -> Trainer.Config:
+    from torchtitan.models.muse_glimmer.config_registry import (
+        muse_glimmer_debugmodel_mm,
+    )
+
+    config = muse_glimmer_debugmodel_mm()
+    _use_spmd_types(config, typechecking=False)
+    config.parallelism.data_parallel_shard_degree = 1
+    config.parallelism.context_parallel_degree = 2
+    config.parallelism.tensor_parallel_degree = 2
+    config.parallelism.enable_sequence_parallel = True
+    config.parallelism.pipeline_parallel_degree = 2
+    config.parallelism.num_pp_microbatches = 2
+    config.parallelism.pipeline_parallel_schedule = "1F1B"
+    config.training.disable_cuda_graphs = True
+    return config

From 454216f3f926af0513392219735a934fab6a5d7b Mon Sep 17 00:00:00 2001
From: Jin Soo Ihm <jinsooihm@meta.com>
Date: Tue, 1 Sep 2026 12:38:21 -0700
Subject: [PATCH 2/3] update

---
 tests/unit_tests/cpu/test_packed_vision.py    | 49 ++++++++++-
 torchtitan/models/common/multimodal.py        | 42 ++++-----
 torchtitan/models/muse_glimmer/README.md      |  4 +-
 torchtitan/models/muse_glimmer/__init__.py    |  4 -
 torchtitan/models/muse_glimmer/model.py       | 34 +++++---
 torchtitan/models/muse_glimmer/parallelize.py |  7 +-
 torchtitan/models/muse_glimmer/sharding.py    | 86 ++++++++-----------
 7 files changed, 131 insertions(+), 95 deletions(-)

diff --git a/tests/unit_tests/cpu/test_packed_vision.py b/tests/unit_tests/cpu/test_packed_vision.py
index a0851eb4d5..a5e1607d09 100644
--- a/tests/unit_tests/cpu/test_packed_vision.py
+++ b/tests/unit_tests/cpu/test_packed_vision.py
@@ -15,7 +15,10 @@
 from torchtitan.components.loss import IGNORE_INDEX
 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 +233,50 @@ 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)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py
index d07443ecef..01a93957d2 100644
--- a/torchtitan/models/common/multimodal.py
+++ b/torchtitan/models/common/multimodal.py
@@ -8,19 +8,17 @@
 
 ``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.
+``gather_vision_embeds`` support gather-based fusion 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]:
@@ -109,29 +107,19 @@ def build_vision_bank_indices(
     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 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)]
+    return torch.where(is_vision_T1, gathered_TD, inputs_TD)
 
 
 def scatter_vision_embeds(
diff --git a/torchtitan/models/muse_glimmer/README.md b/torchtitan/models/muse_glimmer/README.md
index 2cac465068..2e41e8c3a5 100644
--- a/torchtitan/models/muse_glimmer/README.md
+++ b/torchtitan/models/muse_glimmer/README.md
@@ -47,7 +47,9 @@ 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`,
 building packed vision features and gathering them into token embeddings using
-bank indices prepared before token sharding.
+bank indices prepared before token sharding. Multimodal fusion is TP-replicated;
+the first decoder layer restores the standard decoder layout (sequence-parallel
+when enabled).
 
 ## Parallelism support
 
diff --git a/torchtitan/models/muse_glimmer/__init__.py b/torchtitan/models/muse_glimmer/__init__.py
index 0d784fff23..0868762c3f 100644
--- a/torchtitan/models/muse_glimmer/__init__.py
+++ b/torchtitan/models/muse_glimmer/__init__.py
@@ -19,7 +19,6 @@
 )
 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 (
@@ -373,7 +372,6 @@ 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,
@@ -381,7 +379,6 @@ def _muse_glimmer_config(
             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.
@@ -424,7 +421,6 @@ 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,
     )
diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py
index b13e746527..2bf5ec10b1 100644
--- a/torchtitan/models/muse_glimmer/model.py
+++ b/torchtitan/models/muse_glimmer/model.py
@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 from typing import Any, cast
 
+import spmd_types as spmd
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -15,7 +16,7 @@
 from torchtitan.config import ParallelismConfig
 from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.distributed.spmd_types import annotate_input_spmd_types
-from torchtitan.distributed.utils import is_in_batch_invariant_mode
+from torchtitan.distributed.utils import get_spmd_backend, is_in_batch_invariant_mode
 from torchtitan.models.common.attention import (
     AttentionMasksType,
     create_attention_mask,
@@ -36,8 +37,8 @@
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.multimodal import (
     build_vision_bank_indices,
+    gather_vision_embeds,
     multimodal_context,
-    VisionScatter,
 )
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.vision_encoder_sharding import multimodal_input_sharding
@@ -281,11 +282,10 @@ class Config(Decoder.Config):
         # pyrefly: ignore [bad-override]
         tok_embeddings: EmbeddingWithNorm.Config
         # Optional LLM-side multimodal injection. Preprocessing builds absolute
-        # packed-bank indices before CP; VisionScatter gathers the rows needed
-        # by each local token shard after projection and normalization.
+        # packed-bank indices before CP; forward gathers the corresponding vision
+        # rows into the TP-replicated token embeddings.
         vision_projection: Linear.Config | None = None
         perception_emb_norm: RMSNorm.Config | None = None
-        vision_scatter: VisionScatter.Config | None = None
         # Optional owned vision stack. When set, ``MuseGlimmerModel`` builds the encoder
         # + adapter as submodules and runs them inside ``forward`` (from padded
         # ``pixel_values`` + ``grid_thw``), mirroring qwen3_5's
@@ -361,9 +361,6 @@ def __init__(self, config: "MuseGlimmerModel.Config") -> None:
             if config.perception_emb_norm is not None
             else None
         )
-        self.vision_scatter = (
-            config.vision_scatter.build() if config.vision_scatter is not None else None
-        )
         # Owned vision stack (None unless a multimodal flavor configured it). When
         # present, ``forward`` runs encoder->adapter on packed pixel_values.
         self.vision_encoder = (
@@ -478,13 +475,16 @@ def _prepare_multimodal_embeds(
         assert vision_bank_indices_T is not None
         assert self.vision_projection is not None
         assert self.perception_emb_norm is not None
-        assert self.vision_scatter is not None
 
         vision_features_VD = self._get_vision_features(pixel_values, grid_thw)
         vision_bank_VD = self.perception_emb_norm(
             self.vision_projection(vision_features_VD)
         )
-        return self.vision_scatter(h_TD, vision_bank_VD, vision_bank_indices_T)
+        return gather_vision_embeds(
+            h_TD,
+            vision_bank_VD=vision_bank_VD,
+            vision_bank_indices_T=vision_bank_indices_T,
+        )
 
     def forward(
         self,
@@ -517,6 +517,20 @@ def forward(
             else:
                 h_TD = tokens
 
+        # torch.where can erase the token PartitionSpec. Restore it before the
+        # layer-0 FSDP pre-forward hook runs ahead of input redistribution.
+        if (
+            self.tok_embeddings is not None
+            and self.vision_projection is not None
+            and get_spmd_backend() == "spmd_types"
+            and spmd.is_type_checking()
+        ):
+            spmd.assert_type(
+                h_TD,
+                {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R},
+                spmd.PartitionSpec(("dp", "cp"), None),
+            )
+
         for layer in self.layers.values():
             h_TD = layer(h_TD, attention_masks, positions)
 
diff --git a/torchtitan/models/muse_glimmer/parallelize.py b/torchtitan/models/muse_glimmer/parallelize.py
index 46ca6830ed..07799ab9e1 100644
--- a/torchtitan/models/muse_glimmer/parallelize.py
+++ b/torchtitan/models/muse_glimmer/parallelize.py
@@ -154,8 +154,8 @@ def pipeline_muse_glimmer(
 
     Muse Glimmer's owned vision stack runs inside ``MuseGlimmerModel.forward``
     on the embedding stage: the encoder and adapter build raw features, the
-    projection and norm build the packed bank, and the stateless scatter fuses
-    it into token embeddings. All present vision modules live on stage 0.
+    projection and norm build the packed bank, and the model fuses it into token
+    embeddings. All present vision modules live on stage 0.
     """
     import dataclasses
 
@@ -186,7 +186,7 @@ def pipeline_muse_glimmer(
         # vision encoder, bump
         # parallelism.pipeline_parallel_first_stage_less_layers to rebalance.
         # Prepend in data-flow order so the resulting stage-0 list reads
-        # encoder -> adapter -> projection -> emb_norm -> scatter -> embeddings.
+        # encoder -> adapter -> projection -> emb_norm -> embeddings.
         vision_fqns = [
             fqn
             for fqn in (
@@ -194,7 +194,6 @@ def pipeline_muse_glimmer(
                 "vision_adapter",
                 "vision_projection",
                 "perception_emb_norm",
-                "vision_scatter",
             )
             if getattr(model, fqn, None) is not None
         ]
diff --git a/torchtitan/models/muse_glimmer/sharding.py b/torchtitan/models/muse_glimmer/sharding.py
index 711d62523b..3e40807585 100644
--- a/torchtitan/models/muse_glimmer/sharding.py
+++ b/torchtitan/models/muse_glimmer/sharding.py
@@ -40,42 +40,6 @@
 TP = MeshAxisName.TP
 
 
-def _sequence_parallel_index_placement() -> SpmdType:
-    return SpmdType(
-        {DP: spmd.V, CP: spmd.V, TP: spmd.V},
-        partition_spec=spmd.PartitionSpec((DP, CP, TP)),
-    )
-
-
-def _vision_scatter_config(*, enable_sp: bool) -> ShardingConfig:
-    vision_bank_src = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.I})
-    vision_bank_dst = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.R})
-    vision_bank_indices = token_id_placement()
-    if enable_sp:
-        hidden_src = dense_sequence_parallel_placement()
-        hidden_dst = hidden_src
-        local_vision_bank_indices = _sequence_parallel_index_placement()
-    else:
-        hidden_src = dense_activation_placement(tp=spmd.I, cp=spmd.S(0))
-        hidden_dst = dense_activation_placement(tp=spmd.R, cp=spmd.S(0))
-        local_vision_bank_indices = vision_bank_indices
-
-    return ShardingConfig(
-        in_src_shardings={
-            "inputs_TD": hidden_src,
-            "vision_bank_VD": vision_bank_src,
-            "vision_bank_indices_T": vision_bank_indices,
-        },
-        in_dst_shardings={
-            "inputs_TD": hidden_dst,
-            "vision_bank_VD": vision_bank_dst,
-            "vision_bank_indices_T": local_vision_bank_indices,
-        },
-        out_src_shardings=hidden_dst,
-        out_dst_shardings=hidden_src,
-    )
-
-
 def set_muse_glimmer_sharding_config(
     config: "MuseGlimmerModel.Config",
     *,
@@ -83,10 +47,9 @@ def set_muse_glimmer_sharding_config(
 ) -> None:
     """Fill ``sharding_config`` on all Muse Glimmer sub-configs.
 
-    Text-only and multimodal models use the same token-sharded decoder path.
-    The multimodal path keeps its packed vision bank replicated across CP and
-    TP-invariant until ``VisionScatter`` gathers rows into each local token
-    shard. The scatter boundary performs the required TP ``I -> R`` transition.
+    Text-only models use the standard decoder layout. Multimodal models keep
+    token embeddings and the projected vision bank TP-replicated for fusion;
+    the first decoder layer restores the standard SP or invariant layout.
 
     All sub-configs are populated unconditionally -- ``Module.parallelize``
     filters disabled axes at runtime.
@@ -97,7 +60,7 @@ def set_muse_glimmer_sharding_config(
     for layer_cfg in config.layers:
         _set_muse_glimmer_layer_sharding(layer_cfg, enable_sp=enable_sp)
 
-    # Configure the replicated vision bank and token-local scatter boundary.
+    # Configure the TP-replicated multimodal fusion path.
     if config.vision_encoder is not None:
         _set_multimodal_sharding(config, enable_sp=enable_sp)
 
@@ -136,18 +99,45 @@ def _set_multimodal_sharding(
     *,
     enable_sp: bool,
 ) -> None:
-    """Configure the vision bank and its token-sharded scatter boundary."""
+    """Keep multimodal fusion TP-replicated until the first decoder layer."""
+    fusion_layout = dense_activation_placement(tp=spmd.R, cp=spmd.S(0))
+    decoder_layout = (
+        dense_sequence_parallel_placement()
+        if enable_sp
+        else dense_activation_placement(tp=spmd.I, cp=spmd.S(0))
+    )
+    vision_replicated = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.R})
+
+    emb_cfg = config.tok_embeddings
+    emb_cfg.embedding.sharding_config = ShardingConfig(
+        state_shardings={"weight": dense_param_placement(tp=spmd.S(0))},
+        in_src_shardings={"input": token_id_placement()},
+        in_dst_shardings={"input": token_id_placement()},
+        out_src_shardings=dense_activation_placement(tp=spmd.P, cp=spmd.S(0)),
+        out_dst_shardings=fusion_layout,
+        local_map=LocalMapConfig(in_grad_placements=None),
+    )
+    emb_cfg.norm.sharding_config = ShardingConfig(
+        in_src_shardings={"input": fusion_layout},
+        in_dst_shardings={"input": fusion_layout},
+        out_src_shardings=fusion_layout,
+        out_dst_shardings=fusion_layout,
+    )
+
     if config.vision_projection is not None:
         config.vision_projection.sharding_config = vision_invariant_linear_config(
             include_cp_axis=True
         )
     if config.perception_emb_norm is not None:
-        config.perception_emb_norm.sharding_config = invariant_norm_config(
-            include_cp_axis=True
-        )
-    if config.vision_scatter is not None:
-        config.vision_scatter.sharding_config = _vision_scatter_config(
-            enable_sp=enable_sp
+        vision_norm = invariant_norm_config(include_cp_axis=True)
+        vision_norm.out_dst_shardings = vision_replicated
+        config.perception_emb_norm.sharding_config = vision_norm
+
+    if config.layers:
+        config.layers[0].sharding_config = ShardingConfig(
+            in_src_shardings={"x": fusion_layout},
+            in_dst_shardings={"x": decoder_layout},
+            out_src_shardings=decoder_layout,
         )
 
 

From ee49a5750cab3a39b4b8d9795a398ea9900bceb2 Mon Sep 17 00:00:00 2001
From: Jin Soo Ihm <jinsooihm@meta.com>
Date: Tue, 1 Sep 2026 17:13:15 -0700
Subject: [PATCH 3/3] update

---
 tests/unit_tests/cpu/test_packed_vision.py | 43 ++++++++++++++++
 torchtitan/models/common/multimodal.py     |  9 +++-
 torchtitan/models/muse_glimmer/model.py    | 48 +++++++++---------
 torchtitan/models/muse_glimmer/sharding.py | 59 ++++++++--------------
 4 files changed, 95 insertions(+), 64 deletions(-)

diff --git a/tests/unit_tests/cpu/test_packed_vision.py b/tests/unit_tests/cpu/test_packed_vision.py
index a5e1607d09..b8e80eb096 100644
--- a/tests/unit_tests/cpu/test_packed_vision.py
+++ b/tests/unit_tests/cpu/test_packed_vision.py
@@ -8,11 +8,14 @@
 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 (
@@ -277,6 +280,46 @@ def test_gather_vision_embeds_accepts_empty_bank(self) -> None:
 
         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()
diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py
index 01a93957d2..eb4e00cb14 100644
--- a/torchtitan/models/common/multimodal.py
+++ b/torchtitan/models/common/multimodal.py
@@ -119,7 +119,14 @@ def gather_vision_embeds(
     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)
+    # 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(
diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py
index 2bf5ec10b1..1134e933a2 100644
--- a/torchtitan/models/muse_glimmer/model.py
+++ b/torchtitan/models/muse_glimmer/model.py
@@ -16,7 +16,7 @@
 from torchtitan.config import ParallelismConfig
 from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.distributed.spmd_types import annotate_input_spmd_types
-from torchtitan.distributed.utils import get_spmd_backend, is_in_batch_invariant_mode
+from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.models.common.attention import (
     AttentionMasksType,
     create_attention_mask,
@@ -29,10 +29,7 @@
     VarlenAttention,
 )
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
-from torchtitan.models.common.decoder_sharding import (
-    decoder_input_sharding,
-    token_id_placement,
-)
+from torchtitan.models.common.decoder_sharding import decoder_input_sharding
 from torchtitan.models.common.embedding import Embedding
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.multimodal import (
@@ -383,6 +380,8 @@ def preprocess_inputs(
             prepare_context_parallel_input,
         )
 
+        from .sharding import vision_bank_indices_placement
+
         batch: dict[str, Any] = dict(input_dict)
         pixel_values = batch.get("pixel_values")
         grid_thw = batch.get("grid_thw")
@@ -428,7 +427,9 @@ def preprocess_inputs(
             **decoder_input_sharding(),
             **multimodal_input_sharding(include_cp_axis=True),
         }
-        input_sharding["vision_bank_indices_T"] = token_id_placement()
+        input_sharding["vision_bank_indices_T"] = vision_bank_indices_placement(
+            enable_sp=parallelism.enable_sequence_parallel
+        )
         if parallel_dims.cp_enabled:
             batch = prepare_context_parallel_input(
                 batch,
@@ -438,6 +439,17 @@ def preprocess_inputs(
                 parallelism.context_parallel_ptrr_mask_key,
             )
         if parallelism.spmd_backend == "spmd_types":
+            if (
+                parallelism.enable_sequence_parallel
+                and parallel_dims.tp_enabled
+                and "vision_bank_indices_T" in batch
+            ):
+                batch["vision_bank_indices_T"] = spmd.shard(
+                    batch["vision_bank_indices_T"],
+                    parallel_dims.get_dense_tp_mesh().get_group(),
+                    src=spmd.I,
+                    dst=spmd.S(0),
+                )
             batch = annotate_input_spmd_types(parallel_dims, batch, input_sharding)
 
         inputs = batch.pop("input")
@@ -505,31 +517,17 @@ def forward(
         # tok_embeddings) and inject vision features before the decoder layers.
         # On non-embedding pipeline stages tok_embeddings is None and the input
         # is already hidden states, so injection is skipped there.
-        with multimodal_context():
-            if self.tok_embeddings is not None:
-                h_TD = self.tok_embeddings(tokens)
+        if self.tok_embeddings is not None:
+            h_TD = self.tok_embeddings(tokens)
+            with multimodal_context():
                 h_TD = self._prepare_multimodal_embeds(
                     h_TD,
                     pixel_values=pixel_values,
                     grid_thw=grid_thw,
                     vision_bank_indices_T=vision_bank_indices_T,
                 )
-            else:
-                h_TD = tokens
-
-        # torch.where can erase the token PartitionSpec. Restore it before the
-        # layer-0 FSDP pre-forward hook runs ahead of input redistribution.
-        if (
-            self.tok_embeddings is not None
-            and self.vision_projection is not None
-            and get_spmd_backend() == "spmd_types"
-            and spmd.is_type_checking()
-        ):
-            spmd.assert_type(
-                h_TD,
-                {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R},
-                spmd.PartitionSpec(("dp", "cp"), None),
-            )
+        else:
+            h_TD = tokens
 
         for layer in self.layers.values():
             h_TD = layer(h_TD, attention_masks, positions)
diff --git a/torchtitan/models/muse_glimmer/sharding.py b/torchtitan/models/muse_glimmer/sharding.py
index 3e40807585..8dad673158 100644
--- a/torchtitan/models/muse_glimmer/sharding.py
+++ b/torchtitan/models/muse_glimmer/sharding.py
@@ -21,7 +21,6 @@
     set_dense_ffn_sharding,
     set_gqa_attention_sharding,
     set_gqa_inner_attention_local_map,
-    token_id_placement,
 )
 from torchtitan.models.common.vision_encoder_sharding import (
     invariant_norm_config,
@@ -40,6 +39,19 @@
 TP = MeshAxisName.TP
 
 
+def vision_bank_indices_placement(*, enable_sp: bool) -> SpmdType:
+    """Placement for token-aligned indices into the packed vision bank."""
+    token_axes = (DP, CP, TP) if enable_sp else (DP, CP)
+    return SpmdType(
+        {
+            DP: spmd.V,
+            CP: spmd.V,
+            TP: spmd.V if enable_sp else spmd.I,
+        },
+        partition_spec=spmd.PartitionSpec(token_axes),
+    )
+
+
 def set_muse_glimmer_sharding_config(
     config: "MuseGlimmerModel.Config",
     *,
@@ -47,9 +59,9 @@ def set_muse_glimmer_sharding_config(
 ) -> None:
     """Fill ``sharding_config`` on all Muse Glimmer sub-configs.
 
-    Text-only models use the standard decoder layout. Multimodal models keep
-    token embeddings and the projected vision bank TP-replicated for fusion;
-    the first decoder layer restores the standard SP or invariant layout.
+    Text-only and multimodal models use the standard decoder activation layout.
+    With SP, each TP rank gathers vision rows for its local token shard, so the
+    vision bank becomes TP-replicated before fusion.
 
     All sub-configs are populated unconditionally -- ``Module.parallelize``
     filters disabled axes at runtime.
@@ -60,7 +72,6 @@ def set_muse_glimmer_sharding_config(
     for layer_cfg in config.layers:
         _set_muse_glimmer_layer_sharding(layer_cfg, enable_sp=enable_sp)
 
-    # Configure the TP-replicated multimodal fusion path.
     if config.vision_encoder is not None:
         _set_multimodal_sharding(config, enable_sp=enable_sp)
 
@@ -99,47 +110,19 @@ def _set_multimodal_sharding(
     *,
     enable_sp: bool,
 ) -> None:
-    """Keep multimodal fusion TP-replicated until the first decoder layer."""
-    fusion_layout = dense_activation_placement(tp=spmd.R, cp=spmd.S(0))
-    decoder_layout = (
-        dense_sequence_parallel_placement()
-        if enable_sp
-        else dense_activation_placement(tp=spmd.I, cp=spmd.S(0))
-    )
-    vision_replicated = SpmdType({DP: spmd.V, CP: spmd.R, TP: spmd.R})
-
-    emb_cfg = config.tok_embeddings
-    emb_cfg.embedding.sharding_config = ShardingConfig(
-        state_shardings={"weight": dense_param_placement(tp=spmd.S(0))},
-        in_src_shardings={"input": token_id_placement()},
-        in_dst_shardings={"input": token_id_placement()},
-        out_src_shardings=dense_activation_placement(tp=spmd.P, cp=spmd.S(0)),
-        out_dst_shardings=fusion_layout,
-        local_map=LocalMapConfig(in_grad_placements=None),
-    )
-    emb_cfg.norm.sharding_config = ShardingConfig(
-        in_src_shardings={"input": fusion_layout},
-        in_dst_shardings={"input": fusion_layout},
-        out_src_shardings=fusion_layout,
-        out_dst_shardings=fusion_layout,
-    )
-
+    """Configure token-local multimodal fusion."""
     if config.vision_projection is not None:
         config.vision_projection.sharding_config = vision_invariant_linear_config(
             include_cp_axis=True
         )
     if config.perception_emb_norm is not None:
         vision_norm = invariant_norm_config(include_cp_axis=True)
-        vision_norm.out_dst_shardings = vision_replicated
+        if enable_sp:
+            vision_norm.out_dst_shardings = SpmdType(
+                {DP: spmd.V, CP: spmd.R, TP: spmd.R}
+            )
         config.perception_emb_norm.sharding_config = vision_norm
 
-    if config.layers:
-        config.layers[0].sharding_config = ShardingConfig(
-            in_src_shardings={"x": fusion_layout},
-            in_dst_shardings={"x": decoder_layout},
-            out_src_shardings=decoder_layout,
-        )
-
 
 def _set_muse_glimmer_layer_sharding(
     layer_cfg: "MuseGlimmerTransformerBlock.Config",