From 626eb60b12c84eca606d794c2227a7f24db82875 Mon Sep 17 00:00:00 2001 From: Yu Yao Date: Mon, 17 Aug 2026 00:38:04 -0700 Subject: [PATCH 1/2] [model] fix: Finalize providers before distributed model construction Signed-off-by: Yu Yao --- src/megatron/bridge/models/model_provider.py | 2 ++ .../models/test_model_provider_mixin.py | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index 1de2dd49f5..6303ccf846 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -247,6 +247,8 @@ def provide_distributed_model( if wrap_with_ddp and not ddp_config: raise ValueError("ddp_config is required when wrap_with_ddp is True") + self.finalize() + if not torch.distributed.is_initialized(): os.environ["RANK"] = os.environ.get("RANK", "0") os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1") diff --git a/tests/unit_tests/models/test_model_provider_mixin.py b/tests/unit_tests/models/test_model_provider_mixin.py index 0ddfad2cb3..3bc5e1f0fe 100644 --- a/tests/unit_tests/models/test_model_provider_mixin.py +++ b/tests/unit_tests/models/test_model_provider_mixin.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from unittest.mock import Mock, call, patch import pytest @@ -20,6 +21,7 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.bridge.models.llama_nemotron.llama_nemotron_provider import LlamaNemotronHeterogeneousProvider from megatron.bridge.models.model_provider import ModelProviderMixin @@ -92,6 +94,31 @@ def test_apply_overrides_rejects_unknown_attributes(provider): provider.finalize.assert_not_called() +@patch("megatron.bridge.models.model_provider.get_model") +@patch("megatron.bridge.models.model_provider.torch.distributed") +def test_provide_distributed_model_finalizes_heterogeneous_config_before_model_build(mock_dist, mock_get_model): + """Distributed model construction must finalize deferred heterogeneous fields.""" + block = { + "attention": {"no_op": False, "replace_with_linear": False, "num_query_groups": 4}, + "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 256}, + } + provider = LlamaNemotronHeterogeneousProvider( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + heterogeneous_layers_config_encoded_json=json.dumps({"block_configs": [block]}), + ) + mock_dist.is_initialized.return_value = True + + def build_model(model_provider, **_kwargs): + assert len(model_provider.per_block_parameters) == 1 + return [] + + mock_get_model.side_effect = build_model + + provider.provide_distributed_model(pg_collection=Mock(), wrap_with_ddp=False) + + @patch("megatron.bridge.models.model_provider.ProcessGroupCollection.use_mpu_process_groups") @patch("megatron.bridge.models.model_provider.get_model") @patch("megatron.bridge.models.model_provider.torch.distributed") From 26d2fb10846cdab066e807d3bac2df3d772cf093 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 17 Aug 2026 09:32:47 -0700 Subject: [PATCH 2/2] [model] fix: finalize provider before distributed model entry Signed-off-by: yaoyu-33 --- src/megatron/bridge/models/model_provider.py | 2 -- src/megatron/bridge/training/setup.py | 1 + .../models/test_model_provider_mixin.py | 27 -------------- tests/unit_tests/training/test_setup.py | 35 +++++++++++++++++++ 4 files changed, 36 insertions(+), 29 deletions(-) diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index 6303ccf846..1de2dd49f5 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -247,8 +247,6 @@ def provide_distributed_model( if wrap_with_ddp and not ddp_config: raise ValueError("ddp_config is required when wrap_with_ddp is True") - self.finalize() - if not torch.distributed.is_initialized(): os.environ["RANK"] = os.environ.get("RANK", "0") os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1") diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 7355b9adca..59c458c094 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -596,6 +596,7 @@ def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCo data_parallel_random_init=cfg.rng.data_parallel_random_init, ) else: + model_config.finalize() return model_config.provide_distributed_model( ddp_config=cfg.ddp, use_megatron_fsdp=cfg.dist.use_megatron_fsdp, diff --git a/tests/unit_tests/models/test_model_provider_mixin.py b/tests/unit_tests/models/test_model_provider_mixin.py index 3bc5e1f0fe..0ddfad2cb3 100644 --- a/tests/unit_tests/models/test_model_provider_mixin.py +++ b/tests/unit_tests/models/test_model_provider_mixin.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from unittest.mock import Mock, call, patch import pytest @@ -21,7 +20,6 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.bridge.models.llama_nemotron.llama_nemotron_provider import LlamaNemotronHeterogeneousProvider from megatron.bridge.models.model_provider import ModelProviderMixin @@ -94,31 +92,6 @@ def test_apply_overrides_rejects_unknown_attributes(provider): provider.finalize.assert_not_called() -@patch("megatron.bridge.models.model_provider.get_model") -@patch("megatron.bridge.models.model_provider.torch.distributed") -def test_provide_distributed_model_finalizes_heterogeneous_config_before_model_build(mock_dist, mock_get_model): - """Distributed model construction must finalize deferred heterogeneous fields.""" - block = { - "attention": {"no_op": False, "replace_with_linear": False, "num_query_groups": 4}, - "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 256}, - } - provider = LlamaNemotronHeterogeneousProvider( - num_layers=1, - hidden_size=64, - num_attention_heads=4, - heterogeneous_layers_config_encoded_json=json.dumps({"block_configs": [block]}), - ) - mock_dist.is_initialized.return_value = True - - def build_model(model_provider, **_kwargs): - assert len(model_provider.per_block_parameters) == 1 - return [] - - mock_get_model.side_effect = build_model - - provider.provide_distributed_model(pg_collection=Mock(), wrap_with_ddp=False) - - @patch("megatron.bridge.models.model_provider.ProcessGroupCollection.use_mpu_process_groups") @patch("megatron.bridge.models.model_provider.get_model") @patch("megatron.bridge.models.model_provider.torch.distributed") diff --git a/tests/unit_tests/training/test_setup.py b/tests/unit_tests/training/test_setup.py index eeefa23f69..7aeb8c695a 100644 --- a/tests/unit_tests/training/test_setup.py +++ b/tests/unit_tests/training/test_setup.py @@ -13,6 +13,7 @@ # limitations under the License. +import json from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -25,6 +26,7 @@ from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hybrid.hybrid_builder import HybridModelConfig from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.models.llama_nemotron.llama_nemotron_provider import LlamaNemotronHeterogeneousProvider from megatron.bridge.models.transformer_config import TransformerConfig from megatron.bridge.training.callbacks import CallbackManager from megatron.bridge.training.checkpointing import load_checkpoint @@ -556,9 +558,42 @@ def test_build_with_provider(self): result = _build_distributed_model(cfg, pg_collection=MagicMock()) + mock_provider.finalize.assert_called_once_with() mock_provider.provide_distributed_model.assert_called_once() assert result == mock_dist_model + def test_finalizes_heterogeneous_provider_before_entering_distributed_model_build(self): + """Finalize deferred heterogeneous fields before entering provider model construction.""" + block = { + "attention": {"no_op": False, "replace_with_linear": False, "num_query_groups": 4}, + "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 256}, + } + provider = LlamaNemotronHeterogeneousProvider( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + heterogeneous_layers_config_encoded_json=json.dumps({"block_configs": [block]}), + ) + expected_model = [MagicMock()] + + def provide_distributed_model(**_kwargs): + assert len(provider.per_block_parameters) == 1 + return expected_model + + provider.provide_distributed_model = Mock(side_effect=provide_distributed_model) + cfg = SimpleNamespace( + model=provider, + ddp=MagicMock(), + optimizer=SimpleNamespace(overlap_param_gather_with_optimizer_step=False), + dist=SimpleNamespace(use_megatron_fsdp=False, use_torch_fsdp2=False), + rng=SimpleNamespace(data_parallel_random_init=False), + ) + + result = _build_distributed_model(cfg, pg_collection=MagicMock()) + + provider.provide_distributed_model.assert_called_once() + assert result == expected_model + def test_restart_rebinds_overlap_callbacks_to_rebuilt_model(): """A restart must replace callbacks bound to the discarded model."""