fix(model): pass pp_rank to callable transformer layer specs

GPTModelProvider.provide() forwarded only vp_stage to a callable
transformer_layer_spec, so Megatron-Core block-spec builders fell back to
parallel_state.get_pipeline_model_parallel_rank(). Under
use_decentralized_pg=True the MPU globals are never initialized, so building
any block-spec model asserted with "pipeline_model parallel group is not
initialized" before the first forward. mtp_block_spec() had the same omission.

Adds unit tests covering the change (red-green verified).
Detected by: megatron-bridge QA (test_eval_cp_gdn_metadata_e2e)

Signed-off-by: Pruthviraj Prakash <pruprakash@nvidia.com>

Author: pruprakash

src/megatron/bridge/models/gpt_provider.py

@@ -36,6 +36,7 @@
 from megatron.core.transformer import ModuleSpec
 from megatron.core.transformer.dot_product_attention import DotProductAttention as MCoreDotProductAttention
 from megatron.core.transformer.enums import AttnBackend
+from megatron.core.utils import get_pg_rank
 
 from megatron.bridge.models.logit_dtype import logit_dtype_kwarg
 from megatron.bridge.models.model_provider import ModelProviderMixin
@@ -256,11 +257,9 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGP
 
         transformer_layer_spec = self.transformer_layer_spec
         if not isinstance(transformer_layer_spec, ModuleSpec):
-            # Check if the transformer_layer_spec function accepts vp_stage parameter
-            if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:
-                transformer_layer_spec = transformer_layer_spec(self, vp_stage=vp_stage)
-            else:
-                transformer_layer_spec = transformer_layer_spec(self)
+            transformer_layer_spec = transformer_layer_spec(
+                self, **_callable_spec_kwargs(transformer_layer_spec, self, vp_stage)
+            )
 
         assert self.vocab_size is not None, "vocab_size must be configured before calling provide()"
         if self.should_pad_vocab:
@@ -342,6 +341,21 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGP
         return model
 
 
+def _callable_spec_kwargs(spec_fn: Callable, config: "GPTModelProvider", vp_stage: Optional[int]) -> dict[str, Any]:
+    """Build the keyword arguments a callable transformer layer spec declares."""
+    params = inspect.signature(spec_fn).parameters
+    kwargs: dict[str, Any] = {}
+    if "vp_stage" in params:
+        kwargs["vp_stage"] = vp_stage
+    if "pp_rank" in params:
+        # Resolve the pipeline rank from the provider's own groups so that block spec
+        # builders do not fall back to the MPU globals, which decentralized runs never set.
+        pg_collection = getattr(config, "_pg_collection", None)
+        if pg_collection is not None:
+            kwargs["pp_rank"] = get_pg_rank(pg_collection.pp)
+    return kwargs
+
+
 def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -> Optional[ModuleSpec]:
     """Pass in the MTP block spec if model has MTP layers.
 
@@ -355,10 +369,9 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
         from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
 
         if isinstance(config.transformer_layer_spec, Callable):
-            if "vp_stage" in inspect.signature(config.transformer_layer_spec).parameters:
-                spec = config.transformer_layer_spec(config, vp_stage=vp_stage)
-            else:
-                spec = config.transformer_layer_spec(config)
+            spec = config.transformer_layer_spec(
+                config, **_callable_spec_kwargs(config.transformer_layer_spec, config, vp_stage)
+            )
         else:
             spec = config.transformer_layer_spec
         if hasattr(spec, "layer_specs") and len(spec.layer_specs) == 0:

tests/unit_tests/models/test_gpt_provider.py

@@ -635,3 +635,95 @@ def fake_spec_unsupported(
 
         assert result == "te_spec_unsupported"
         assert "use_grouped_gemm_for_dense_mlp" not in captured
+
+
+class TestCallableSpecPpRank:
+    """Tests for forwarding pp_rank to callable transformer layer specs."""
+
+    @staticmethod
+    def _provider(**kwargs) -> GPTModelProvider:
+        return GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            vocab_size=1000,
+            **kwargs,
+        )
+
+    def test_provide_forwards_pp_rank_resolved_from_the_pg_collection(self):
+        """A spec callable that accepts pp_rank receives the rank of _pg_collection.pp."""
+        captured: dict = {}
+
+        def spec_fn(config, vp_stage=None, pp_rank=None):
+            captured["vp_stage"] = vp_stage
+            captured["pp_rank"] = pp_rank
+            return "block_spec"
+
+        pp_group = object()
+        provider = self._provider(transformer_layer_spec=spec_fn)
+        provider._pg_collection = type("PG", (), {"pp": pp_group, "tp": object(), "cp": object()})()
+
+        with patch("megatron.bridge.models.gpt_provider.get_pg_rank", return_value=3) as mock_rank:
+            with patch("megatron.bridge.models.gpt_provider.MCoreGPTModel"):
+                provider.provide(pre_process=True, post_process=True, vp_stage=1)
+
+        mock_rank.assert_called_once_with(pp_group)
+        assert captured["pp_rank"] == 3, "pp_rank must come from the provider's pipeline group"
+        assert captured["vp_stage"] == 1
+
+    def test_provide_omits_pp_rank_for_specs_that_do_not_accept_it(self):
+        """Regression: a spec callable without a pp_rank parameter is called unchanged."""
+        captured: dict = {}
+
+        def spec_fn(config, vp_stage=None):
+            captured["vp_stage"] = vp_stage
+            return "block_spec"
+
+        provider = self._provider(transformer_layer_spec=spec_fn)
+        provider._pg_collection = type("PG", (), {"pp": object(), "tp": object(), "cp": object()})()
+
+        with patch("megatron.bridge.models.gpt_provider.MCoreGPTModel") as mock_model:
+            provider.provide(pre_process=True, post_process=True, vp_stage=2)
+
+        assert captured["vp_stage"] == 2
+        assert mock_model.call_args.kwargs["transformer_layer_spec"] == "block_spec"
+
+    def test_provide_omits_pp_rank_when_no_pg_collection_is_set(self):
+        """Without a process group collection the spec keeps its own pp_rank default."""
+        captured: dict = {}
+
+        def spec_fn(config, vp_stage=None, pp_rank="unset"):
+            captured["pp_rank"] = pp_rank
+            return "block_spec"
+
+        provider = self._provider(transformer_layer_spec=spec_fn)
+
+        with patch("megatron.bridge.models.gpt_provider.MCoreGPTModel"):
+            provider.provide(pre_process=True, post_process=True)
+
+        assert captured["pp_rank"] == "unset"
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_mtp_block_spec_forwards_pp_rank_to_callable_spec(self, mock_get_mtp):
+        """mtp_block_spec resolves pp_rank the same way when re-invoking the spec callable."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        captured: dict = {}
+        block_spec = Mock()
+        block_spec.layer_specs = ["layer_a"]
+
+        def spec_fn(config, vp_stage=None, pp_rank=None):
+            captured["pp_rank"] = pp_rank
+            return block_spec
+
+        pp_group = object()
+        provider = self._provider(mtp_num_layers=1, transformer_layer_spec=spec_fn)
+        provider._pg_collection = type("PG", (), {"pp": pp_group, "tp": object(), "cp": object()})()
+        mock_get_mtp.return_value = "mtp_spec"
+
+        with patch("megatron.bridge.models.gpt_provider.get_pg_rank", return_value=5) as mock_rank:
+            result = mtp_block_spec(provider, vp_stage=0)
+
+        mock_rank.assert_called_once_with(pp_group)
+        assert captured["pp_rank"] == 5
+        assert result == "mtp_spec"