fix(compat): adapt Bridge to Megatron-Core dev

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Author: svcnemo-autobot

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

@@ -1286,10 +1286,16 @@ def _forward_mlp(
         hidden_states: Tensor,
         inference_context: BaseInferenceContext | None = None,
         padding_mask: Tensor | None = None,
+        input_ids: Tensor | None = None,
         packed_seq_params: PackedSeqParams | None = None,
     ) -> Tensor:
-        """Run HF's separate shared-expert, routed-expert, and router inputs."""
-        del inference_context
+        """Run HF's separate shared-expert, routed-expert, and router inputs.
+
+        ``input_ids`` is accepted for compatibility with Megatron-Core ``dev``'s
+        ``TransformerLayer.forward``, which forwards it for hash-based MoE routing;
+        Gemma 4's separate-input MoE path does not use it.
+        """
+        del inference_context, input_ids
         residual = hidden_states.float() if self.config.fp32_residual_connection else hidden_states
 
         moe_input = residual
@@ -1377,9 +1383,10 @@ def routing(
         logits: Tensor,
         padding_mask: Tensor | None = None,
         input_ids: Tensor | None = None,
+        **kwargs: object,
     ) -> tuple[Tensor, Tensor | None]:
-        # Token identities do not affect Gemma 4 routing; retain the argument for compatibility with existing callers.
-        del input_ids
+        """Route Gemma 4 tokens with arguments accepted by both MCore refs."""
+        del input_ids, kwargs
         routing_probs, routing_map = super().routing(logits, padding_mask=padding_mask)
         if routing_map is not None:
             prob_sums = routing_probs.sum(dim=-1, keepdim=True).clamp(min=1e-20)

tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py

@@ -144,7 +144,8 @@ def copy(src: str, dst: str) -> None:
 
         indexer_prefix = f"{compressor_prefix}.indexer"
         copy(f"{indexer_prefix}.q_b_proj.weight", f"{ckpt_prefix}.attn.indexer.wq_b.weight")
-        copy(f"{indexer_prefix}.weights_proj.weight", f"{ckpt_prefix}.attn.indexer.weights_proj.weight")
+        # HF transformers nests the indexer score projection under `.scorer.`.
+        copy(f"{indexer_prefix}.scorer.weights_proj.weight", f"{ckpt_prefix}.attn.indexer.weights_proj.weight")
         copy(f"{indexer_prefix}.kv_proj.weight", f"{ckpt_prefix}.attn.indexer.compressor.wkv.weight")
         copy(f"{indexer_prefix}.gate_proj.weight", f"{ckpt_prefix}.attn.indexer.compressor.wgate.weight")
         copy(f"{indexer_prefix}.position_bias", f"{ckpt_prefix}.attn.indexer.compressor.ape")

tests/unit_tests/models/gemma/test_gemma4_modeling.py

@@ -1910,7 +1910,11 @@ def fake_routing(self, logits, padding_mask=None, input_ids=None):
         router = object.__new__(Gemma4TopKRouter)
         router.per_expert_scale = torch.tensor([1.0, 2.0, 3.0])
 
-        out_probs, out_map = Gemma4TopKRouter.routing(router, torch.zeros(2, 3))
+        out_probs, out_map = Gemma4TopKRouter.routing(
+            router,
+            torch.zeros(2, 3),
+            packed_seq_params=object(),
+        )
 
         assert out_map is routing_map
         torch.testing.assert_close(out_probs[0], torch.tensor([0.4, 1.2, 0.0]))

tests/unit_tests/recipes/test_glm5_perf_recipes.py

@@ -20,9 +20,6 @@
 from types import SimpleNamespace
 
 import pytest
-from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
-    _validate_dsa_index_share_pipeline_split,
-)
 from megatron.core.transformer.enums import LayerType
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 
@@ -90,6 +87,20 @@ def _build_recipe(recipe_func: Callable[[], ConfigContainer], monkeypatch: pytes
     return recipe_func()
 
 
+def _dsa_source_layer_id(layer_id: int, *, skip_topk_offset: int, topk_freq: int) -> int:
+    """Return the zero-based source layer defined by MCore's DSA sharing contract."""
+    # Mirrors the private MCore `_validate_dsa_index_share_pipeline_split`
+    # helper that guarded this recipe before removal from the pinned revision.
+    # MCore defines DSA sharing with one-based layers: layers through
+    # max(skip_topk_offset, 1) compute their own indices, then each topk_freq
+    # group reuses the indices computed by its first layer.
+    layer_number = layer_id + 1
+    sharing_offset = max(skip_topk_offset, 1)
+    if layer_number <= sharing_offset:
+        return layer_id
+    return layer_number - ((layer_number - sharing_offset) % topk_freq) - 1
+
+
 @pytest.mark.parametrize("recipe_func", _RECIPES, ids=lambda recipe: recipe.__name__)
 def test_glm5_perf_recipes_are_flat_and_preserve_bridge_dsa_fields(
     recipe_func: Callable[[], ConfigContainer], monkeypatch: pytest.MonkeyPatch
@@ -183,7 +194,14 @@ def test_glm52_h100_pipeline_layout_keeps_dsa_index_sharing_within_each_vpp_chun
             decoder_count = stage.count(LayerType.decoder)
             if decoder_count:
                 local_layer_ids = range(decoder_offset, decoder_offset + decoder_count)
-                _validate_dsa_index_share_pipeline_split(cfg.model, local_layer_ids)
+                local_layer_id_set = set(local_layer_ids)
+                for layer_id in local_layer_ids:
+                    source_layer_id = _dsa_source_layer_id(
+                        layer_id,
+                        skip_topk_offset=cfg.model.dsa_indexer_skip_topk_offset,
+                        topk_freq=cfg.model.dsa_indexer_topk_freq,
+                    )
+                    assert source_layer_id in local_layer_id_set
                 decoder_offset += decoder_count
 
     assert decoder_offset == cfg.model.num_layers