From dd4f9cd8c33859cdfd7f21dd2c45b79d8e8663f6 Mon Sep 17 00:00:00 2001 From: Yu Yao Date: Thu, 3 Sep 2026 01:24:01 -0700 Subject: [PATCH 1/2] [ckpt] fix: initialize model parallel for builder loads Signed-off-by: Yu Yao --- .../bridge/training/model_load_save.py | 9 ++++ .../training/test_model_load_save.py | 49 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index 8015a296a5..8fad8bc104 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -364,6 +364,15 @@ def _call_model_provider(model_cfg): model_cfg.finalize() builder_cls = model_cfg.get_builder_cls() builder = builder_cls(model_cfg) + if dist.is_initialized() and 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, + ) # Note: `use_cpu_initialization` is not passed as an explicit kwarg here, # unlike the ModelProviderMixin path which passes it directly to # `provide_distributed_model`. Instead, it is handled by the diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index 8773a3f7c9..34733b74c7 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/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 @@ -444,6 +446,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") From 1b5cb0b6135662526ee45c25979520a886bc605f Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 10:12:06 -0700 Subject: [PATCH 2/2] refactor(ckpt): share model parallel initialization Signed-off-by: yaoyu-33 --- .../bridge/training/model_load_save.py | 49 ++++++---- .../training/test_model_load_save.py | 98 ++++++++++++++++++- 2 files changed, 125 insertions(+), 22 deletions(-) diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index 8fad8bc104..6db02fd3e8 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/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,24 +376,20 @@ 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) - if dist.is_initialized() and 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, - ) # Note: `use_cpu_initialization` is not passed as an explicit kwarg here, # unlike the ModelProviderMixin path which passes it directly to # `provide_distributed_model`. Instead, it is handled by the @@ -380,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." diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index 34733b74c7..88ba7cb0dd 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -240,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( @@ -252,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.""" @@ -313,6 +385,7 @@ def _built_layer_count(checkpoint_path, model_cfg, *args): assert built_layer_count == provider.num_layers + @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") @@ -329,6 +402,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 @@ -346,6 +420,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 @@ -360,7 +436,12 @@ 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 @@ -368,24 +449,24 @@ def test_load_mbridge_saved_model( 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 @@ -416,7 +497,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 @@ -434,6 +515,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, @@ -581,6 +663,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") @@ -597,6 +680,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.""" @@ -616,6 +700,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" @@ -625,6 +710,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") @@ -645,6 +731,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 @@ -664,6 +751,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