fix(ckpt): skip export of parameters no PP rank owns

src/megatron/bridge/models/conversion/param_mapping.py

@@ -432,7 +432,9 @@ def broadcast_from_pp_rank(
 
         return tensor
 
-    def broadcast_obj_from_pp_rank(self, obj: Optional[Any], cache_key: Optional[str] = None) -> Any:
+    def broadcast_obj_from_pp_rank(
+        self, obj: Optional[Any], cache_key: Optional[str] = None, allow_missing: bool = False
+    ) -> Any:
         """Broadcast any Python object from the PP rank that owns it.
 
         This method is useful for broadcasting configuration objects or
@@ -443,12 +445,16 @@ def broadcast_obj_from_pp_rank(self, obj: Optional[Any], cache_key: Optional[str
             obj (Optional[Any]): Object to broadcast (None on non-owning ranks).
             cache_key (Optional[str]): Optional cache key. If not provided,
                 no caching will be performed.
+            allow_missing (bool): Return None instead of raising when no PP rank
+                owns the object. Must be passed by every rank of the PP group.
 
         Returns:
-            Any: Broadcasted object on all ranks.
+            Any: Broadcasted object on all ranks, or None if no rank owns it and
+            ``allow_missing`` is set.
 
         Raises:
-            ValueError: If object does not exist on any rank.
+            ValueError: If object does not exist on any rank and ``allow_missing``
+                is not set.
         """
         if self.pp_size == 1:
             return obj
@@ -476,6 +482,8 @@ def broadcast_obj_from_pp_rank(self, obj: Optional[Any], cache_key: Optional[str
                 break
 
         if src_rank is None:
+            if allow_missing:
+                return None
             raise ValueError("Object must exist on at least one PP rank")
 
         # ------------------------------------------------------------------
@@ -1581,11 +1589,8 @@ def megatron_to_hf(
                 # Broadcast to other ranks
                 self._detected_type = self.broadcast_obj_from_pp_rank(self._detected_type, "detected_type")
             else:
-                # Receive from owning rank
-                self._detected_type = self.broadcast_obj_from_pp_rank(None, "detected_type")
-                if self._detected_type is None:
-                    # PP group likely has 1 member - skipping.
-                    return {}
+                # Receive from owning rank, or None when no PP rank owns the module
+                self._detected_type = self.broadcast_obj_from_pp_rank(None, "detected_type", allow_missing=True)
 
             # If no PP rank detected a type (e.g. Megatron parameter without an
             # HF counterpart, such as MoE modules on dense layers created by

tests/unit_tests/models/test_param_mapping.py

@@ -301,6 +301,52 @@ def forward(self, x):
         finally:
             AutoMapping._MODULE_TYPE_REGISTRY["column"].discard("Linear")
 
+    def test_megatron_to_hf_skips_param_no_pp_rank_owns(self, mock_distributed_env):
+        """A parameter whose module is absent from every PP rank exports as no weights."""
+        _, mock_dist = mock_distributed_env(pp_size=2, pp_rank=1)
+        mapping = AutoMapping(megatron_param="some.weight", hf_param="hf.weight")
+
+        mock_dist.all_gather_object.side_effect = lambda output, obj, group: output.__setitem__(
+            slice(None), [False, False]
+        )
+
+        result = mapping.megatron_to_hf(None, None)
+
+        assert result == {}, f"expected an empty export for an unowned parameter, got {result}"
+        assert mock_dist.broadcast_object_list.call_count == 0, "nothing should be broadcast when no rank owns it"
+
+    def test_megatron_to_hf_uses_type_broadcast_by_owning_pp_rank(self, mock_distributed_env):
+        """Regression: a non-owning rank still receives the parallelism type from the owner."""
+        _, mock_dist = mock_distributed_env(pp_size=2, pp_rank=1)
+        mapping = AutoMapping(megatron_param="some.weight", hf_param="hf.weight")
+
+        mock_dist.all_gather_object.side_effect = lambda output, obj, group: output.__setitem__(
+            slice(None), [True, False]
+        )
+        mock_dist.broadcast_object_list.side_effect = lambda obj_list, src, group: obj_list.__setitem__(
+            0, "replicated"
+        )
+
+        with patch.object(AutoMapping, "_get_or_create_mapping") as mock_get_mapping:
+            mock_get_mapping.return_value.megatron_to_hf.return_value = {"hf.weight": torch.zeros(2)}
+            result = mapping.megatron_to_hf(None, None)
+
+        assert mapping._detected_type == "replicated"
+        assert mock_get_mapping.call_args[0][0] == "replicated"
+        assert set(result) == {"hf.weight"}
+
+    def test_broadcast_obj_from_pp_rank_raises_when_unowned_by_default(self, mock_distributed_env):
+        """Without allow_missing, an object owned by no PP rank is still an error."""
+        _, mock_dist = mock_distributed_env(pp_size=2, pp_rank=1)
+        mapping = AutoMapping(megatron_param="some.weight", hf_param="hf.weight")
+
+        mock_dist.all_gather_object.side_effect = lambda output, obj, group: output.__setitem__(
+            slice(None), [False, False]
+        )
+
+        with pytest.raises(ValueError, match="Object must exist on at least one PP rank"):
+            mapping.broadcast_obj_from_pp_rank(None, "detected_type")
+
 
 class TestHelperFunctions:
     def test_qkv_merge_split(self, transformer_config):