diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 0a38392ce0..653b5b6e45 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, parse_hybrid_pattern from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.core.tensor_parallel import param_is_not_tensor_parallel_duplicate from megatron.core.transformer.module import MegatronModule @@ -594,6 +595,42 @@ def _build_moe_metric_writer( return _MoeMetricFanoutWriter(tb_writer, comet_logger, mlflow_logger) +def _get_num_moe_layers(model_config: Any) -> int: + """Count MoE aux-loss contributors for MCore metric averaging.""" + num_layers = model_config.num_layers + mtp_num_layers = getattr(model_config, "mtp_num_layers", None) or 0 + repeated_mtp = getattr(model_config, "mtp_use_repeated_layer", False) + + if getattr(model_config, "is_hybrid_model", False): + pattern = parse_hybrid_pattern(getattr(model_config, "hybrid_layer_pattern", None)) + main_moe_layers = (pattern.main_pattern or "").count(Symbols.MOE) + mtp_moe_layers = (pattern.mtp_pattern or "").count(Symbols.MOE) + mtp_depth = pattern.mtp_num_depths + else: + moe_layer_freq = getattr(model_config, "moe_layer_freq", None) + if moe_layer_freq is None: + main_moe_layers = num_layers + last_layer_is_moe = True + elif isinstance(moe_layer_freq, int): + main_moe_layers = sum(i % moe_layer_freq == 0 for i in range(num_layers)) + last_layer_is_moe = (num_layers - 1) % moe_layer_freq == 0 + elif isinstance(moe_layer_freq, list): + main_moe_layers = sum(moe_layer_freq) + last_layer_is_moe = bool(moe_layer_freq[-1]) + else: + raise ValueError(f"Invalid moe_layer_freq: {moe_layer_freq}") + + # Non-hybrid MTP copies the final main layer's dense/MoE layout. + mtp_moe_layers = int(last_layer_is_moe and mtp_num_layers > 0) + mtp_depth = mtp_num_layers + + # A reused MTP block contributes once; distinct blocks contribute at each depth. + if not repeated_mtp: + mtp_moe_layers *= mtp_depth + + return main_moe_layers + mtp_moe_layers + + def training_log( loss_dict: dict[str, torch.Tensor], total_loss_dict: dict[str, Any], @@ -986,11 +1023,6 @@ def training_log( if getattr(config.model, "moe_z_loss_coeff", None) is not None: track_names.append("z_loss") - if getattr(config.model, "is_hybrid_model", False): - layers = getattr(config.model, "hybrid_layer_pattern", "").count("E") - else: - layers = getattr(config.model, "num_layers", None) - # Wrap the TB writer so MoE/MTP metrics also reach MLFlow / Comet (issue #2989). # No-op when neither logger is configured: the original writer is returned as-is. moe_metric_writer = _build_moe_metric_writer(writer, comet_logger, mlflow_logger) @@ -1003,7 +1035,8 @@ def training_log( per_layer_logging=getattr(config.model, "moe_per_layer_logging", False), force_initialize=True, track_names=track_names, - num_layers=layers, + num_layers=config.model.num_layers, + num_moe_layers=_get_num_moe_layers(config.model), moe_layer_freq=getattr(config.model, "moe_layer_freq", None), mtp_num_layers=getattr(config.model, "mtp_num_layers", None), pg_collection=pg_collection, diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index 10345bb917..1041799eed 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -28,6 +28,7 @@ from megatron.bridge.training.state import GlobalState from megatron.bridge.training.utils.train_utils import ( LinearForLastLayer, + _get_num_moe_layers, calc_params_l2_norm, create_value_head_hook, freeze_moe_router, @@ -83,6 +84,8 @@ def make_default_model_config(): num_layers=24, moe_layer_freq=1, mtp_num_layers=None, + mtp_use_repeated_layer=False, + hybrid_layer_pattern=None, kv_channels=128, num_attention_heads=32, hidden_size=4096, @@ -111,6 +114,62 @@ def make_default_model_config(): ) +@pytest.mark.parametrize( + ("num_layers", "mtp_num_layers", "hybrid_pattern", "moe_layer_freq", "repeated_mtp", "expected"), + [ + pytest.param( + 52, + 2, + "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME/*E/*E", + None, + True, + 24, + id="nemotron_3_5", + ), + pytest.param( + 88, + 2, + "MEMEMEM*EMEMEMEM*EMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEM*EMEMEMEME/*E/*E", + None, + True, + 41, + id="nemotron_3_super", + ), + pytest.param( + 108, + 2, + ( + "MEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*" + "EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME/*E/*E" + ), + None, + True, + 49, + id="nemotron_3_ultra", + ), + pytest.param(61, 1, None, [0] * 3 + [1] * 58, False, 59, id="deepseek_v3"), + pytest.param(45, 0, None, [0] * 3 + [1] * 42, False, 42, id="step_3_7"), + pytest.param(4, 2, "MEME/*E/*E", None, False, 4, id="hybrid_distinct_mtp"), + pytest.param(4, 2, "MEME/**/**", None, False, 2, id="hybrid_dense_mtp"), + pytest.param(4, 2, None, [1, 0, 1, 1], True, 4, id="non_hybrid_repeated_mtp"), + pytest.param(4, 2, None, [1, 0, 1, 1], False, 5, id="non_hybrid_distinct_mtp"), + pytest.param(4, 2, None, [1, 0, 1, 0], False, 2, id="non_hybrid_dense_mtp"), + ], +) +def test_get_num_moe_layers(num_layers, mtp_num_layers, hybrid_pattern, moe_layer_freq, repeated_mtp, expected): + """Test representative architecture and edge-case contributor counts.""" + model_config = SimpleNamespace( + num_layers=num_layers, + moe_layer_freq=moe_layer_freq, + mtp_num_layers=mtp_num_layers, + mtp_use_repeated_layer=repeated_mtp, + is_hybrid_model=hybrid_pattern is not None, + hybrid_layer_pattern=hybrid_pattern, + ) + + assert _get_num_moe_layers(model_config) == expected + + class TestTrainingLog: """Test suite for the training_log function.""" @@ -991,6 +1050,13 @@ def test_memory_reporting_checkpoint_resume( assert result is True mock_report_memory.assert_called_once() + @pytest.mark.parametrize( + ("model_num_layers", "mtp_num_layers", "hybrid_pattern", "moe_layer_freq", "expected_num_moe_layers"), + [ + pytest.param(12, None, None, 2, 6, id="non_hybrid"), + pytest.param(4, 2, "MMME/*E/*E", None, 3, id="hybrid"), + ], + ) @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group") @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe") @@ -1012,6 +1078,11 @@ def test_moe_logging( mock_config, mock_global_state, loss_dict, + model_num_layers, + mtp_num_layers, + hybrid_pattern, + moe_layer_freq, + expected_num_moe_layers, ): """Test MoE (Mixture of Experts) logging when enabled.""" # Get fresh total_loss_dict for this test @@ -1030,9 +1101,11 @@ def test_moe_logging( mock_config.model.moe_router_load_balancing_type = "aux_loss" mock_config.model.moe_z_loss_coeff = 0.1 mock_config.model.moe_per_layer_logging = True - mock_config.model.num_layers = 12 - mock_config.model.moe_layer_freq = 2 - mock_config.model.mtp_num_layers = None + mock_config.model.num_layers = model_num_layers + mock_config.model.moe_layer_freq = moe_layer_freq + mock_config.model.mtp_num_layers = mtp_num_layers + mock_config.model.is_hybrid_model = hybrid_pattern is not None + mock_config.model.hybrid_layer_pattern = hybrid_pattern training_log( loss_dict=loss_dict, @@ -1056,6 +1129,9 @@ def test_moe_logging( call_args = mock_track_moe.call_args assert "load_balancing_loss" in call_args.kwargs["track_names"] assert "z_loss" in call_args.kwargs["track_names"] + assert call_args.kwargs["num_layers"] == model_num_layers + assert call_args.kwargs["num_moe_layers"] == expected_num_moe_layers + assert call_args.kwargs["mtp_num_layers"] == mtp_num_layers @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group")