diff --git a/src/megatron/bridge/recipes/nemotronh/nemotronh.py b/src/megatron/bridge/recipes/nemotronh/nemotronh.py index 3b9ee0cf05..b9b0b187df 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotronh.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotronh.py @@ -139,7 +139,7 @@ def nemotronh_47b_pretrain_config(**user_kwargs: Unpack[NemotronHCommonKwargs]) "pipeline_parallelism": 1, "sequence_parallelism": True, "precision_config": "nemotron_h_bf16_with_fp8_current_scaling_mixed", - "enable_default_comm_overlap": False, + "enable_default_comm_overlap": True, } combined_kwargs: NemotronHCommonKwargs = {**recommended_kwargs, **user_kwargs} return _nemotronh_common(tokenizer_model="nvidia/Nemotron-H-47B-Base-8K", **combined_kwargs) @@ -162,7 +162,7 @@ def nemotronh_56b_pretrain_config(**user_kwargs: Unpack[NemotronHCommonKwargs]) "pipeline_parallelism": 1, "sequence_parallelism": True, "precision_config": "nemotron_h_bf16_with_fp8_current_scaling_mixed", - "enable_default_comm_overlap": False, + "enable_default_comm_overlap": True, } combined_kwargs: NemotronHCommonKwargs = {**recommended_kwargs, **user_kwargs} return _nemotronh_common(tokenizer_model="nvidia/Nemotron-H-8B-Base-8K", **combined_kwargs) diff --git a/src/megatron/bridge/training/initialize.py b/src/megatron/bridge/training/initialize.py index f462ffaac2..bdd63988d2 100644 --- a/src/megatron/bridge/training/initialize.py +++ b/src/megatron/bridge/training/initialize.py @@ -302,7 +302,24 @@ def _initialize_tp_communicators(model_config: GPTModelProvider | T5ModelProvide model_config.hidden_size, ] - if is_te_min_version("1.9.0"): + if is_te_min_version("2.7.0"): + UserBufferQuantizationMode = te_module.base.UserBufferQuantizationMode + quantization_modes = [UserBufferQuantizationMode.FP8 if model_config.fp8 else UserBufferQuantizationMode.NONE] + if ( + model_config.fp8 is not None + and model_config.first_last_layers_bf16 + and (model_config.num_layers_at_start_in_bf16 > 0 or model_config.num_layers_at_end_in_bf16 > 0) + ): + quantization_modes.append(UserBufferQuantizationMode.NONE) + # The process group with the target bootstrap backend is created in Transformer Engine. + te_module.base.initialize_ub( + shape=input_shape, + tp_size=model_config.tensor_model_parallel_size, + quantization_modes=quantization_modes, + ub_cfgs=ub_cfgs, + bootstrap_backend=model_config.tp_comm_bootstrap_backend, + ) + elif is_te_min_version("1.9.0"): # The process group with the target bootstrap backend is created in Transformer Engine. te_module.base.initialize_ub( shape=input_shape, diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotronh.py b/tests/unit_tests/recipes/nemotronh/test_nemotronh.py index 58632320d1..ba11124381 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotronh.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotronh.py @@ -66,6 +66,7 @@ def test_pretrain_config_default_parameters(self): # Check comm overlap assert config.comm_overlap is not None assert config.comm_overlap.tp_comm_overlap is True + assert config.comm_overlap.tp_comm_bootstrap_backend == "nccl" def test_pretrain_config_custom_parallelism(self): """Test pretrain_config with custom parallelism.""" @@ -126,6 +127,7 @@ def test_pretrain_config_default_parameters(self): # Check comm overlap assert config.comm_overlap is not None assert config.comm_overlap.tp_comm_overlap is True + assert config.comm_overlap.tp_comm_bootstrap_backend == "nccl" def test_pretrain_config_custom_parallelism(self): """Test pretrain_config with custom parallelism.""" @@ -166,9 +168,10 @@ def test_pretrain_config_default_parameters(self): # Check logger config assert config.logger.log_interval == 10 - # Check comm overlap is not set by default for 47B - # (due to issues with fp8 current scaling) - assert config.comm_overlap is None + # Check comm overlap + assert config.comm_overlap is not None + assert config.comm_overlap.tp_comm_overlap is True + assert config.comm_overlap.tp_comm_bootstrap_backend == "nccl" def test_pretrain_config_custom_parallelism(self): """Test pretrain_config with custom parallelism.""" @@ -207,8 +210,10 @@ def test_pretrain_config_default_parameters(self): # Check logger config assert config.logger.log_interval == 10 - # Check comm overlap is not set by default for 56B - assert config.comm_overlap is None + # Check comm overlap + assert config.comm_overlap is not None + assert config.comm_overlap.tp_comm_overlap is True + assert config.comm_overlap.tp_comm_bootstrap_backend == "nccl" def test_pretrain_config_custom_parallelism(self): """Test pretrain_config with custom parallelism.""" diff --git a/tests/unit_tests/training/test_initialize.py b/tests/unit_tests/training/test_initialize.py new file mode 100644 index 0000000000..145283e38a --- /dev/null +++ b/tests/unit_tests/training/test_initialize.py @@ -0,0 +1,311 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tempfile +import warnings +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest +import yaml + +from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.training.initialize import _initialize_tp_communicators + + +class TestInitializeTPCommunicators: + """Test suite for _initialize_tp_communicators function.""" + + @pytest.fixture + def mock_gpt_config(self): + """Create a mock GPT model configuration.""" + config = Mock(spec=GPTModelProvider) + config.seq_length = 1024 + config.hidden_size = 768 + config.context_parallel_size = 1 + config.tensor_model_parallel_size = 2 + config.tp_comm_overlap_cfg = None + config.fp8 = None + config.first_last_layers_bf16 = False + config.num_layers_at_start_in_bf16 = 0 + config.num_layers_at_end_in_bf16 = 0 + config.tp_comm_bootstrap_backend = "nccl" + return config + + def test_import_error_transformer_engine_missing(self, mock_gpt_config): + """Test ImportError when transformer_engine is not available.""" + with patch("builtins.__import__", side_effect=ImportError("No module named 'transformer_engine'")): + with pytest.raises( + RuntimeError, + match="Tensor Parallel Communication/GEMM Overlap optimization needs 'yaml' and 'transformer_engine' packages", + ): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + def test_import_error_yaml_missing(self, mock_gpt_config): + """Test ImportError when yaml is not available.""" + with patch("builtins.__import__", side_effect=ImportError("No module named 'yaml'")): + with pytest.raises( + RuntimeError, + match="Tensor Parallel Communication/GEMM Overlap optimization needs 'yaml' and 'transformer_engine' packages", + ): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_config_loading_from_string_file(self, mock_init_ub, mock_gpt_config): + """Test loading tp_comm_overlap_cfg from a string file path.""" + # Create a temporary YAML file + config_data = {"buffer_size": 1024, "overlap_enabled": True} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_data, f) + config_file_path = f.name + + try: + mock_gpt_config.tp_comm_overlap_cfg = config_file_path + + with patch("megatron.bridge.training.initialize.is_te_min_version", return_value=True): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + # Verify that the config was loaded from the file + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["ub_cfgs"] == config_data + finally: + Path(config_file_path).unlink() + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_config_loading_from_dict(self, mock_init_ub, mock_gpt_config): + """Test loading tp_comm_overlap_cfg from a dictionary.""" + config_data = {"buffer_size": 2048, "overlap_enabled": False} + mock_gpt_config.tp_comm_overlap_cfg = config_data + + with patch("megatron.bridge.training.initialize.is_te_min_version", return_value=True): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + # Verify that the config dict was used directly + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["ub_cfgs"] == config_data + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_config_loading_none(self, mock_init_ub, mock_gpt_config): + """Test when tp_comm_overlap_cfg is None.""" + mock_gpt_config.tp_comm_overlap_cfg = None + + with patch("megatron.bridge.training.initialize.is_te_min_version", return_value=True): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + # Verify that empty dict was used + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["ub_cfgs"] == {} + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_input_shape_calculation_gpt(self, mock_init_ub, mock_gpt_config): + """Test input_shape calculation for GPT model.""" + mock_gpt_config.seq_length = 1024 + mock_gpt_config.hidden_size = 768 + mock_gpt_config.context_parallel_size = 2 + micro_batch_size = 8 + + expected_shape = [ + (1024 * 8) // 2, # seq_length * micro_batch_size // context_parallel_size + 768, # hidden_size + ] + + with patch("megatron.bridge.training.initialize.is_te_min_version", return_value=True): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size) + + call_args = mock_init_ub.call_args + assert call_args[1]["shape"] == expected_shape + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("transformer_engine.pytorch.module.base.UserBufferQuantizationMode") + def test_te_version_2_7_0_fp8_disabled(self, mock_quant_mode, mock_init_ub, mock_gpt_config): + """Test TE version 2.7.0+ path with FP8 disabled.""" + mock_gpt_config.fp8 = None + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "2.7.0"): + mock_quant_mode.FP8 = "FP8" + mock_quant_mode.NONE = "NONE" + + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["quantization_modes"] == ["NONE"] + assert call_args[1]["tp_size"] == mock_gpt_config.tensor_model_parallel_size + assert call_args[1]["bootstrap_backend"] == mock_gpt_config.tp_comm_bootstrap_backend + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("transformer_engine.pytorch.module.base.UserBufferQuantizationMode") + def test_te_version_2_7_0_fp8_enabled(self, mock_quant_mode, mock_init_ub, mock_gpt_config): + """Test TE version 2.7.0+ path with FP8 enabled.""" + mock_gpt_config.fp8 = "e4m3" + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "2.7.0"): + mock_quant_mode.FP8 = "FP8" + mock_quant_mode.NONE = "NONE" + + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["quantization_modes"] == ["FP8"] + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("transformer_engine.pytorch.module.base.UserBufferQuantizationMode") + def test_te_version_2_7_0_fp8_with_bf16_layers(self, mock_quant_mode, mock_init_ub, mock_gpt_config): + """Test TE version 2.7.0+ path with FP8 and BF16 first/last layers.""" + mock_gpt_config.fp8 = "e4m3" + mock_gpt_config.first_last_layers_bf16 = True + mock_gpt_config.num_layers_at_start_in_bf16 = 2 + mock_gpt_config.num_layers_at_end_in_bf16 = 1 + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "2.7.0"): + mock_quant_mode.FP8 = "FP8" + mock_quant_mode.NONE = "NONE" + + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["quantization_modes"] == ["FP8", "NONE"] + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("transformer_engine.pytorch.module.base.UserBufferQuantizationMode") + def test_te_version_2_7_0_fp8_with_bf16_layers_no_layers(self, mock_quant_mode, mock_init_ub, mock_gpt_config): + """Test TE version 2.7.0+ path with FP8 and BF16 flag but no BF16 layers.""" + mock_gpt_config.fp8 = "e4m3" + mock_gpt_config.first_last_layers_bf16 = True + mock_gpt_config.num_layers_at_start_in_bf16 = 0 + mock_gpt_config.num_layers_at_end_in_bf16 = 0 + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "2.7.0"): + mock_quant_mode.FP8 = "FP8" + mock_quant_mode.NONE = "NONE" + + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["quantization_modes"] == ["FP8"] + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_te_version_1_9_0_fp8_disabled(self, mock_init_ub, mock_gpt_config): + """Test TE version 1.9.0+ path with FP8 disabled.""" + mock_gpt_config.fp8 = None + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "1.9.0"): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["use_fp8"] is False + assert call_args[1]["tp_size"] == mock_gpt_config.tensor_model_parallel_size + assert call_args[1]["bootstrap_backend"] == mock_gpt_config.tp_comm_bootstrap_backend + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + def test_te_version_1_9_0_fp8_enabled(self, mock_init_ub, mock_gpt_config): + """Test TE version 1.9.0+ path with FP8 enabled.""" + mock_gpt_config.fp8 = "e4m3" + + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "1.9.0"): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["use_fp8"] is True + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("torch.distributed.new_group") + def test_te_version_legacy_mpi_backend(self, mock_new_group, mock_init_ub, mock_gpt_config): + """Test legacy TE version path with MPI backend.""" + mock_gpt_config.tp_comm_bootstrap_backend = "mpi" + + with ( + patch("megatron.bridge.training.initialize.is_te_min_version", return_value=False), + patch("megatron.bridge.training.initialize.get_te_version", return_value="1.8.0"), + ): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + # Should not create new group for MPI backend + mock_new_group.assert_called_once_with(backend="mpi") + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["use_fp8"] is False + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("torch.distributed.new_group") + def test_te_version_legacy_non_mpi_backend_warning(self, mock_new_group, mock_init_ub, mock_gpt_config): + """Test legacy TE version path with non-MPI backend shows warning.""" + mock_gpt_config.tp_comm_bootstrap_backend = "nccl" + + with ( + patch("megatron.bridge.training.initialize.is_te_min_version", return_value=False), + patch("megatron.bridge.training.initialize.get_te_version", return_value="1.8.0"), + ): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + # Check that warning was issued + assert len(w) == 1 + assert "Transformer Engine v1.8.0 supports only MPI bootstrap backend" in str(w[0].message) + + # Should create MPI group for non-MPI backend + mock_new_group.assert_called_once_with(backend="mpi") + mock_init_ub.assert_called_once() + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("torch.distributed.new_group") + def test_te_version_legacy_fp8_enabled(self, mock_new_group, mock_init_ub, mock_gpt_config): + """Test legacy TE version path with FP8 enabled.""" + mock_gpt_config.fp8 = "e4m3" + mock_gpt_config.tp_comm_bootstrap_backend = "mpi" + + with ( + patch("megatron.bridge.training.initialize.is_te_min_version", return_value=False), + patch("megatron.bridge.training.initialize.get_te_version", return_value="1.8.0"), + ): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + + mock_init_ub.assert_called_once() + call_args = mock_init_ub.call_args + assert call_args[1]["use_fp8"] is True + + @patch("transformer_engine.pytorch.module.base.initialize_ub") + @patch("torch.distributed.new_group") + def test_version_checking_logic(self, mock_new_group, mock_init_ub, mock_gpt_config): + """Test that version checking logic works correctly.""" + # Test 2.7.0+ path + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "2.7.0"): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + call_args = mock_init_ub.call_args + assert "quantization_modes" in call_args[1] + assert "use_fp8" not in call_args[1] + + # Test 1.9.0+ path + with patch("megatron.bridge.training.initialize.is_te_min_version", side_effect=lambda v: v == "1.9.0"): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + call_args = mock_init_ub.call_args + assert "use_fp8" in call_args[1] + assert "quantization_modes" not in call_args[1] + + # Test legacy path + with patch("megatron.bridge.training.initialize.is_te_min_version", return_value=False): + _initialize_tp_communicators(mock_gpt_config, micro_batch_size=4) + call_args = mock_init_ub.call_args + assert "use_fp8" in call_args[1] + assert "quantization_modes" not in call_args[1] + assert "bootstrap_backend" not in call_args[1]