From 4bf28b94acfad0274fe42d3e7128bd4d0fcc6bef Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:42:03 +0000 Subject: [PATCH] [conversion] fix: skip absent auto mappings across PP Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../bridge/models/conversion/param_mapping.py | 9 +++++++-- tests/unit_tests/models/test_param_mapping.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index 63b4cc371d..39df84b82b 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -1366,9 +1366,14 @@ def megatron_to_hf( 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") + try: + self._detected_type = self.broadcast_obj_from_pp_rank(None, "detected_type") + except ValueError as error: + if str(error) != "Object must exist on at least one PP rank": + raise + self._detected_type = None if self._detected_type is None: - # PP group likely has 1 member - skipping. + # No PP stage owns this optional parameter. return {} self._mapping = self._get_or_create_mapping(self._detected_type) diff --git a/tests/unit_tests/models/test_param_mapping.py b/tests/unit_tests/models/test_param_mapping.py index 6050ea28ff..332da158a4 100644 --- a/tests/unit_tests/models/test_param_mapping.py +++ b/tests/unit_tests/models/test_param_mapping.py @@ -273,6 +273,17 @@ class MyCustomRow(torch.nn.Module): with pytest.raises(ValueError): mapping._detect_parallelism_type(torch.nn.Linear(5, 5)) + def test_megatron_to_hf_skips_parameter_missing_from_all_pp_stages(self, mock_distributed_env): + _, mock_dist = mock_distributed_env(pp_size=2, pp_rank=0) + mapping = AutoMapping(megatron_param="optional.weight", hf_param="hf.optional.weight") + + mock_dist.all_gather_object.side_effect = lambda output, obj, group: output.__setitem__( + slice(None), [False, False] + ) + + assert mapping.megatron_to_hf(None, None) == {} + mock_dist.broadcast_object_list.assert_not_called() + def test_detect_parallelism_type_dynamic_module(self): mtq = pytest.importorskip("modelopt.torch.quantization") DynamicModule = pytest.importorskip("modelopt.torch.opt.dynamic").DynamicModule