fix(model): preserve Gemma 4 image attention across pipeline stages (NVIDIA-NeMo#6187)

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/models/gemma/modeling_gemma4.py

@@ -1741,8 +1741,21 @@ def __init__(
         **kwargs,
     ):
         config = copy.deepcopy(config)
-        if _is_local_attn_layer(layer_number, config.interleaved_attn_pattern):
-            config.window_size = (config.window_size - 1, 0)
+        is_local = _is_local_attn_layer(layer_number, config.interleaved_attn_pattern)
+        text_config = getattr(config, "text_config", None)
+        self._image_bidirectional_attention = (
+            is_local and getattr(text_config, "use_bidirectional_attention", None) == "vision"
+        )
+        self._image_attention_window = config.window_size if self._image_bidirectional_attention else None
+        if is_local:
+            if self._image_bidirectional_attention:
+                # Compose the left window into the explicit mask. TE cannot express
+                # a finite left window with an unrestricted right window, and a zero
+                # right window would remove future tokens in the same image block.
+                config.window_size = None
+                attn_mask_type = AttnMaskType.arbitrary
+            else:
+                config.window_size = (config.window_size - 1, 0)
         else:
             config.window_size = None
 
@@ -1755,6 +1768,28 @@ def __init__(
             **kwargs,
         )
 
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        attention_mask: Tensor | None,
+        attn_mask_type: AttnMaskType,
+        **kwargs,
+    ) -> Tensor:
+        """Honor the image mask in sliding layers while keeping global layers causal."""
+        if self._image_bidirectional_attention:
+            # TE ignores explicit masks in causal mode. Global Gemma 4 attention
+            # remains causal, including image tokens, as in the HF implementation.
+            query_positions = torch.arange(query.size(0), device=query.device) + key.size(0) - query.size(0)
+            key_positions = torch.arange(key.size(0), device=key.device)
+            outside_window = key_positions[None, :] <= query_positions[:, None] - self._image_attention_window
+            if attention_mask is None:
+                attention_mask = key_positions[None, :] > query_positions[:, None]
+            attention_mask = attention_mask | outside_window[None, None, :, :]
+            attn_mask_type = AttnMaskType.arbitrary
+        return super().forward(query, key, value, attention_mask, attn_mask_type, **kwargs)
+
 
 class Gemma4RotaryEmbedding(RotaryEmbedding):
     """Gemma 4 MoE position RoPE — dual local/global embeddings."""

src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py

@@ -365,8 +365,7 @@ def freeze(
 
     def _compute_attention_mask(self, input_ids: torch.Tensor) -> Optional[torch.Tensor]:
         """Compute HF-style attention masks for full and sliding Gemma4 layers."""
-        if not self.pre_process:
-            return None
+        # Every pipeline stage receives token IDs and must preserve image attention.
         batch_size, seq_len = input_ids.shape
         causal_mask = torch.tril(
             torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool, device=input_ids.device)

tests/unit_tests/models/gemma/test_gemma4_modeling.py

@@ -838,6 +838,77 @@ def fake_forward(self, **kwargs):
 
 
 class TestGemma4TEDotProductAttention:
+    def test_image_attention_keeps_future_image_tokens_in_window(self, monkeypatch):
+        from megatron.core.transformer.enums import AttnMaskType
+
+        calls = {}
+
+        def fake_init(self, **kwargs):
+            calls.update(kwargs)
+
+        monkeypatch.setattr("megatron.bridge.models.gemma.modeling_gemma4.TEDotProductAttention.__init__", fake_init)
+        config = SimpleNamespace(
+            interleaved_attn_pattern=(1, 1),
+            window_size=512,
+            text_config=SimpleNamespace(use_bidirectional_attention="vision"),
+        )
+        Gemma4TEDotProductAttention(
+            config=config, layer_number=1, attn_mask_type=AttnMaskType.causal, attention_type="self"
+        )
+
+        assert calls["config"].window_size is None
+        assert calls["attn_mask_type"] == AttnMaskType.arbitrary
+        assert config.window_size == 512
+
+    @pytest.mark.parametrize("image_bidirectional", [True, False])
+    def test_forward_selects_explicit_image_mask_mode(self, monkeypatch, image_bidirectional):
+        from megatron.core.transformer.enums import AttnMaskType
+
+        calls = {}
+
+        def fake_forward(self, query, key, value, attention_mask, attn_mask_type, **kwargs):
+            calls.update(mask=attention_mask, mask_type=attn_mask_type, **kwargs)
+            return value
+
+        monkeypatch.setattr("megatron.bridge.models.gemma.modeling_gemma4.TEDotProductAttention.forward", fake_forward)
+        attention = object.__new__(Gemma4TEDotProductAttention)
+        attention._image_bidirectional_attention = image_bidirectional
+        attention._image_attention_window = 512 if image_bidirectional else None
+        mask = torch.tensor([[[[False, False], [False, False]]]])
+        value = torch.ones(2, 1, 1, 4)
+
+        output = attention.forward(value, value, value, mask, AttnMaskType.causal, attention_bias=None)
+
+        assert output is value
+        torch.testing.assert_close(calls["mask"], mask)
+        assert calls["mask_type"] == (AttnMaskType.arbitrary if image_bidirectional else AttnMaskType.causal)
+        assert calls["attention_bias"] is None
+
+    @pytest.mark.parametrize("with_image_mask", [True, False])
+    def test_image_attention_applies_left_window_without_losing_future_image_tokens(
+        self, monkeypatch, with_image_mask
+    ):
+        from megatron.core.transformer.enums import AttnMaskType
+
+        calls = {}
+
+        def fake_forward(self, query, key, value, attention_mask, attn_mask_type, **kwargs):
+            calls["mask"] = attention_mask
+            return value
+
+        monkeypatch.setattr("megatron.bridge.models.gemma.modeling_gemma4.TEDotProductAttention.forward", fake_forward)
+        attention = object.__new__(Gemma4TEDotProductAttention)
+        attention._image_bidirectional_attention = True
+        attention._image_attention_window = 2
+        value = torch.ones(4, 1, 1, 4)
+        mask = torch.zeros(1, 1, 4, 4, dtype=torch.bool) if with_image_mask else None
+
+        attention.forward(value, value, value, mask, AttnMaskType.causal)
+
+        assert calls["mask"][0, 0, 3, 0].item() is True
+        assert calls["mask"][0, 0, 3, 2].item() is False
+        assert calls["mask"][0, 0, 0, 1].item() is (not with_image_mask)
+
     def test_init_sets_local_window_size(self, monkeypatch):
         calls = []
 
@@ -870,7 +941,11 @@ def fake_init(self, **kwargs):
             "megatron.bridge.models.gemma.modeling_gemma4.TEDotProductAttention.__init__",
             fake_init,
         )
-        cfg = SimpleNamespace(interleaved_attn_pattern=(1, 1), window_size=512)
+        cfg = SimpleNamespace(
+            interleaved_attn_pattern=(1, 1),
+            window_size=512,
+            text_config=SimpleNamespace(use_bidirectional_attention="vision"),
+        )
 
         Gemma4TEDotProductAttention(
             config=cfg,

tests/unit_tests/models/gemma_vl/test_gemma4_vl_modeling.py

@@ -151,13 +151,16 @@ def test_two_separate_image_blocks(self):
         # The key point: block1 CANNOT attend forward to block2 (causal blocks it)
         assert mask[0, 0, 0, 3].item() is True, "block1 pos 0 should be blocked from future block2 pos 3"
 
-    def test_not_pre_process_returns_none(self):
-        """Returns None when pre_process=False (PP pipeline stage)."""
+    @pytest.mark.parametrize("pre_process", [True, False])
+    def test_image_attention_on_every_pipeline_stage(self, pre_process):
+        """Later stages preserve image bidirectionality and text causality."""
         model = _make_model()
-        model.pre_process = False
-        input_ids = self._make_ids([self.TEXT_TOKEN] * 4)
+        model.pre_process = pre_process
+        input_ids = self._make_ids([self.TEXT_TOKEN, self.IMAGE_TOKEN, self.IMAGE_TOKEN, self.TEXT_TOKEN])
         result = model._compute_attention_mask(input_ids)
-        assert result is None
+        assert result[0, 0, 1, 2].item() is False
+        assert result[0, 0, 0, 1].item() is True
+        assert result[0, 0, 2, 3].item() is True
 
     def test_output_shape_batch_size_2(self):
         """Mask shape is [B, 1, S, S] for batch_size=2."""