fix(ckpt): initialize model parallel for builder loads (#5942)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/model_load_save.py

@@ -161,6 +161,27 @@ def temporary_distributed_context(backend: str = "gloo") -> Generator[None, None
         dist.destroy_process_group()
 
 
+def _get_or_initialize_pg_collection(
+    model_cfg: TransformerConfig | ModelConfig,
+) -> ProcessGroupCollection:
+    """Return MPU process groups, initializing model-parallel state when needed."""
+    if not parallel_state.is_initialized():
+        parallel_state.initialize_model_parallel(
+            tensor_model_parallel_size=model_cfg.tensor_model_parallel_size,
+            pipeline_model_parallel_size=model_cfg.pipeline_model_parallel_size,
+            virtual_pipeline_model_parallel_size=model_cfg.virtual_pipeline_model_parallel_size,
+            context_parallel_size=model_cfg.context_parallel_size or 1,
+            expert_model_parallel_size=model_cfg.expert_model_parallel_size or 1,
+            expert_tensor_parallel_size=model_cfg.expert_tensor_parallel_size,
+        )
+        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
+            from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed
+
+            model_parallel_cuda_manual_seed(0)
+
+    return ProcessGroupCollection.use_mpu_process_groups()
+
+
 def load_tokenizer(checkpoint_path: str, **kwargs) -> MegatronTokenizer:
     """Create a tokenizer from a training checkpoint.
 
@@ -355,13 +376,18 @@ def build_and_load_model(
 
     def _call_model_provider(model_cfg):
         """Handles provider call for both MBridge and MLM providers."""
-        if isinstance(model_cfg, ModelProviderMixin):
-            if hasattr(model_cfg, "finalize"):
-                model_cfg.finalize()
-            return model_cfg.provide_distributed_model(wrap_with_ddp=False, use_cpu_initialization=use_cpu_init)
-        elif isinstance(model_cfg, ModelConfig):
+        if isinstance(model_cfg, (ModelProviderMixin, ModelConfig)):
             if hasattr(model_cfg, "finalize"):
                 model_cfg.finalize()
+            pg_collection = _get_or_initialize_pg_collection(model_cfg)
+
+            if isinstance(model_cfg, ModelProviderMixin):
+                return model_cfg.provide_distributed_model(
+                    wrap_with_ddp=False,
+                    use_cpu_initialization=use_cpu_init,
+                    pg_collection=pg_collection,
+                )
+
             builder_cls = model_cfg.get_builder_cls()
             builder = builder_cls(model_cfg)
             # Note: `use_cpu_initialization` is not passed as an explicit kwarg here,
@@ -371,9 +397,7 @@ def _call_model_provider(model_cfg):
             # the flag on the config object. This is intentional — we do not want
             # to duplicate TransformerConfig fields like `use_cpu_initialization`
             # as kwargs on `build_distributed_models`.
-            return builder.build_distributed_models(
-                ProcessGroupCollection.use_mpu_process_groups(), wrap_with_ddp=False
-            )
+            return builder.build_distributed_models(pg_collection, wrap_with_ddp=False)
         else:
             assert model_type in ("gpt", "hybrid", "mamba"), f"model type {model_type} not supported."
             assert megatron_args is not None, "megatron_args must be provided if the checkpoint is from MegatronLM."

tests/unit_tests/training/test_model_load_save.py

@@ -20,6 +20,8 @@
 
 import pytest
 import torch
+import torch.distributed as dist
+from megatron.core import parallel_state
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 from megatron.training.models.base import ModelConfig
 
@@ -238,7 +240,11 @@ def test_temporary_distributed_context_nccl(self, mock_seed, mock_os, mock_socke
         mock_socket_instance.getsockname.return_value = ("localhost", 12345)
         mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance
 
-        with temporary_distributed_context(backend="nccl"):
+        with (
+            patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=True),
+            patch("megatron.bridge.training.model_load_save.torch.cuda.device_count", return_value=1),
+            temporary_distributed_context(backend="nccl"),
+        ):
             pass
 
         mock_dist.init_process_group.assert_called_once_with(
@@ -250,6 +256,74 @@ def test_temporary_distributed_context_nccl(self, mock_seed, mock_os, mock_socke
         mock_dist.destroy_process_group.assert_called_once()
 
 
+class TestGetOrInitializePgCollection:
+    """Test shared model-parallel process-group setup for Bridge model configs."""
+
+    @patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed")
+    @patch("megatron.bridge.training.model_load_save.torch.cuda.device_count", return_value=1)
+    @patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=True)
+    @patch("megatron.bridge.training.model_load_save.ProcessGroupCollection")
+    @patch("megatron.bridge.training.model_load_save.parallel_state")
+    def test_initializes_model_parallel_and_cuda_rng(
+        self,
+        mock_parallel_state,
+        mock_pg_collection,
+        mock_cuda_available,
+        mock_cuda_device_count,
+        mock_seed,
+    ):
+        """Initialize missing MPU state before returning its process groups."""
+        model_cfg = SimpleNamespace(
+            tensor_model_parallel_size=2,
+            pipeline_model_parallel_size=4,
+            virtual_pipeline_model_parallel_size=2,
+            context_parallel_size=0,
+            expert_model_parallel_size=None,
+            expert_tensor_parallel_size=1,
+        )
+        expected_pg_collection = Mock()
+        mock_parallel_state.is_initialized.return_value = False
+        mock_pg_collection.use_mpu_process_groups.return_value = expected_pg_collection
+
+        result = model_load_save._get_or_initialize_pg_collection(model_cfg)
+
+        assert result is expected_pg_collection
+        mock_parallel_state.initialize_model_parallel.assert_called_once_with(
+            tensor_model_parallel_size=2,
+            pipeline_model_parallel_size=4,
+            virtual_pipeline_model_parallel_size=2,
+            context_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+        )
+        mock_seed.assert_called_once_with(0)
+        mock_pg_collection.use_mpu_process_groups.assert_called_once_with()
+
+    @patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed")
+    @patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=True)
+    @patch("megatron.bridge.training.model_load_save.ProcessGroupCollection")
+    @patch("megatron.bridge.training.model_load_save.parallel_state")
+    def test_reuses_initialized_model_parallel_state(
+        self,
+        mock_parallel_state,
+        mock_pg_collection,
+        mock_cuda_available,
+        mock_seed,
+    ):
+        """Do not reinitialize or reseed an existing MPU state."""
+        model_cfg = Mock()
+        expected_pg_collection = Mock()
+        mock_parallel_state.is_initialized.return_value = True
+        mock_pg_collection.use_mpu_process_groups.return_value = expected_pg_collection
+
+        result = model_load_save._get_or_initialize_pg_collection(model_cfg)
+
+        assert result is expected_pg_collection
+        mock_parallel_state.initialize_model_parallel.assert_not_called()
+        mock_seed.assert_not_called()
+        mock_pg_collection.use_mpu_process_groups.assert_called_once_with()
+
+
 class TestLoadMegatronModel:
     """Test load_megatron_model function."""
 
@@ -361,6 +435,7 @@ def _finalized_layout(checkpoint_path, model_cfg, *args):
         else:
             assert result is None
 
+    @patch("megatron.bridge.training.model_load_save._get_or_initialize_pg_collection")
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.checkpointing._load_model_weights_from_checkpoint")
     @patch("megatron.bridge.utils.instantiate_utils.instantiate")
@@ -377,6 +452,7 @@ def test_load_mbridge_saved_model(
         mock_instantiate,
         mock_load_weights,
         mock_temp_dist,
+        mock_get_pg_collection,
     ):
         # Setup mocks
         mock_dist.is_available.return_value = False
@@ -394,6 +470,8 @@ def test_load_mbridge_saved_model(
         mock_model_cfg.use_cpu_initialization = False
 
         mock_instantiate.return_value = mock_model_cfg
+        expected_pg_collection = Mock()
+        mock_get_pg_collection.return_value = expected_pg_collection
         expected_result = {"layer.weight": torch.randn(2, 2)}
         mock_load_weights.return_value = expected_result
 
@@ -408,32 +486,37 @@ def test_load_mbridge_saved_model(
         mock_run_config.assert_called_once()
         mock_instantiate.assert_called_once_with(mock_run_cfg_dict["model"])
         mock_cpu_context.assert_called_once()
-        mock_model_cfg.provide_distributed_model.assert_called_once()
+        mock_get_pg_collection.assert_called_once_with(mock_model_cfg)
+        mock_model_cfg.provide_distributed_model.assert_called_once_with(
+            wrap_with_ddp=False,
+            use_cpu_initialization=True,
+            pg_collection=expected_pg_collection,
+        )
         mock_load_weights.assert_called_once_with(ckpt_path, [mock_model], return_state_dict=True)
         assert mock_model_cfg.params_dtype == torch.bfloat16
 
         result = load_megatron_model(ckpt_path, return_state_dict=False, use_cpu_init=True)
         assert result == [mock_model]
         mock_load_weights.assert_called_with(ckpt_path, [mock_model], return_state_dict=False)
 
+    @patch("megatron.bridge.training.model_load_save._get_or_initialize_pg_collection")
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.checkpointing._load_model_weights_from_checkpoint")
     @patch("megatron.bridge.training.checkpointing.read_run_config")
     @patch("megatron.bridge.training.checkpointing.get_checkpoint_run_config_filename")
     @patch("megatron.bridge.training.model_load_save.megatron_cpu_init_context")
     @patch("megatron.bridge.training.model_load_save.dist")
-    @patch("megatron.bridge.training.model_load_save.ProcessGroupCollection")
     @patch("megatron.bridge.training.model_load_save.ModelConfig.from_dict")
     def test_load_mbridge_saved_model_config(
         self,
         mock_from_dict,
-        mock_pg_collection,
         mock_dist,
         mock_cpu_context,
         mock_run_config_fname,
         mock_run_config,
         mock_load_weights,
         mock_temp_dist,
+        mock_get_pg_collection,
     ):
         """Test loading a model when config yaml contains a serialized ModelConfig instance."""
         # Setup mocks
@@ -464,7 +547,7 @@ def test_load_mbridge_saved_model_config(
         mock_from_dict.return_value = mock_model_cfg
 
         mock_mpu_pgs = Mock()
-        mock_pg_collection.use_mpu_process_groups.return_value = mock_mpu_pgs
+        mock_get_pg_collection.return_value = mock_mpu_pgs
 
         expected_result = {"layer.weight": torch.randn(2, 2)}
         mock_load_weights.return_value = expected_result
@@ -482,6 +565,7 @@ def test_load_mbridge_saved_model_config(
         mock_cpu_context.assert_called_once()
         mock_model_cfg.finalize.assert_called_once()
         mock_model_cfg.get_builder_cls.assert_called_once()
+        mock_get_pg_collection.assert_called_once_with(mock_model_cfg)
         mock_builder_cls.assert_called_once_with(mock_model_cfg)
         mock_builder.build_distributed_models.assert_called_once_with(
             mock_mpu_pgs,
@@ -494,6 +578,53 @@ def test_load_mbridge_saved_model_config(
         assert result == [mock_model]
         mock_load_weights.assert_called_with(ckpt_path, [mock_model], return_state_dict=False)
 
+    @patch("megatron.bridge.training.checkpointing._load_model_weights_from_checkpoint")
+    @patch("megatron.bridge.training.checkpointing.read_run_config")
+    @patch("megatron.bridge.training.checkpointing.get_checkpoint_run_config_filename")
+    @patch("megatron.bridge.training.model_load_save.ModelConfig.from_dict")
+    def test_load_builder_config_initializes_model_parallel_with_existing_default_group(
+        self,
+        mock_from_dict,
+        mock_run_config_fname,
+        mock_run_config,
+        mock_load_weights,
+        tmp_path,
+    ):
+        """Load a builder checkpoint when the caller owns only the default process group."""
+        mock_run_config_fname.return_value = tmp_path / "run_config.yaml"
+        mock_run_config.return_value = {
+            "model": {"tensor_model_parallel_size": 1, "_builder_": "import.path.to.SomeModelBuilder"}
+        }
+
+        mock_model = Mock()
+        mock_model_cfg = Mock(spec=GPTModelConfig)
+        mock_model_cfg.params_dtype = torch.float32
+        mock_model_cfg.bf16 = False
+        mock_model_cfg.fp16 = False
+        mock_model_cfg.use_cpu_initialization = False
+        mock_model_cfg.finalize = Mock()
+        mock_builder = Mock()
+        mock_builder.build_distributed_models.return_value = [mock_model]
+        mock_model_cfg.get_builder_cls.return_value = Mock(return_value=mock_builder)
+        mock_from_dict.return_value = mock_model_cfg
+
+        (tmp_path / "run_config.yaml").touch()
+        rendezvous = tmp_path / "gloo_rendezvous"
+        dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1)
+        try:
+            assert not parallel_state.is_initialized()
+
+            result = load_megatron_model(tmp_path, use_cpu_init=True)
+
+            assert result == [mock_model]
+            assert parallel_state.is_initialized()
+            mock_builder.build_distributed_models.assert_called_once()
+            mock_load_weights.assert_called_once_with(tmp_path, [mock_model], return_state_dict=False)
+        finally:
+            if parallel_state.is_initialized():
+                parallel_state.destroy_model_parallel()
+            dist.destroy_process_group()
+
     @pytest.mark.parametrize("model_type", ["gpt", "hybrid", "mamba", "resnet"])
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.mlm_compat.model._hybrid_provider")
@@ -582,6 +713,7 @@ def test_load_mlm_saved_model(
             with pytest.raises(AssertionError, match=f"model type {model_type} not supported."):
                 load_megatron_model(ckpt_path, model_type=model_type, return_state_dict=True, use_cpu_init=True)
 
+    @patch("megatron.bridge.training.model_load_save._get_or_initialize_pg_collection")
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.checkpointing._load_model_weights_from_checkpoint")
     @patch("megatron.bridge.utils.instantiate_utils.instantiate")
@@ -598,6 +730,7 @@ def test_load_megatron_model_skip_temp_dist_context(
         mock_instantiate,
         mock_load_weights,
         mock_temp_dist,
+        mock_get_pg_collection,
     ):
         """Test loading model when distributed is already initialized."""
 
@@ -617,6 +750,7 @@ def test_load_megatron_model_skip_temp_dist_context(
         mock_model_cfg.use_cpu_initialization = False
 
         mock_instantiate.return_value = mock_model_cfg
+        mock_get_pg_collection.return_value = Mock()
 
         with tempfile.TemporaryDirectory() as ckpt_path:
             config_file = Path(ckpt_path) / "run_config.yaml"
@@ -626,6 +760,7 @@ def test_load_megatron_model_skip_temp_dist_context(
         assert result == mock_model
         mock_temp_dist.assert_not_called()
 
+    @patch("megatron.bridge.training.model_load_save._get_or_initialize_pg_collection")
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.post_training.checkpointing.load_modelopt_state")
     @patch("megatron.bridge.training.post_training.checkpointing.has_modelopt_state")
@@ -646,6 +781,7 @@ def test_load_mbridge_saved_model_with_modelopt_state(
         mock_has_modelopt_state,
         mock_load_modelopt_state,
         mock_temp_dist,
+        mock_get_pg_collection,
     ):
         """Test loading model when modelopt state exists and model supports it."""
         # Setup mocks
@@ -665,6 +801,7 @@ def test_load_mbridge_saved_model_with_modelopt_state(
         mock_model_cfg.restore_modelopt_state = False  # Initially False
 
         mock_instantiate.return_value = mock_model_cfg
+        mock_get_pg_collection.return_value = Mock()
         expected_result = {"layer.weight": torch.randn(2, 2)}
         mock_load_weights.return_value = expected_result