Merge commit 'c931bed154d09469a745bc89b8f57931f35d3444' into chore/bump-mcore-main-260904

Author: svcnemo-autobot

tests/unit_tests/training/utils/test_train_utils.py

@@ -1060,12 +1060,22 @@ def test_memory_reporting_checkpoint_resume(
         mock_report_memory.assert_called_once()
 
     @pytest.mark.parametrize(
-        ("model_num_layers", "mtp_num_layers", "hybrid_pattern", "moe_layer_freq", "expected_num_moe_layers"),
+        (
+            "model_num_layers",
+            "mtp_num_layers",
+            "hybrid_pattern",
+            "moe_layer_freq",
+            "expected_num_moe_layers",
+            "supports_num_moe_layers",
+        ),
         [
-            pytest.param(12, None, None, 2, 6, id="non_hybrid"),
-            pytest.param(4, 2, "MMME/*E/*E", None, 3, id="hybrid"),
+            pytest.param(12, None, None, 2, 6, True, id="non_hybrid-supported"),
+            pytest.param(12, None, None, 2, 6, False, id="non_hybrid-unsupported"),
+            pytest.param(4, 2, "MMME/*E/*E", None, 3, True, id="hybrid-supported"),
+            pytest.param(4, 2, "MMME/*E/*E", None, 3, False, id="hybrid-unsupported"),
         ],
     )
+    @mock.patch("megatron.bridge.training.utils.train_utils._track_moe_metrics_supports_num_moe_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")
     @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe")
@@ -1084,6 +1094,7 @@ def test_moe_logging(
         mock_get_world_size,
         mock_reduce_lr,
         mock_get_microbatches,
+        mock_supports_num_moe_layers,
         mock_config,
         mock_global_state,
         loss_dict,
@@ -1092,12 +1103,14 @@ def test_moe_logging(
         hybrid_pattern,
         moe_layer_freq,
         expected_num_moe_layers,
+        supports_num_moe_layers,
     ):
         """Test MoE (Mixture of Experts) logging when enabled."""
         # Get fresh total_loss_dict for this test
         total_loss_dict = self.get_fresh_total_loss_dict()
 
         # Setup mocks
+        mock_supports_num_moe_layers.return_value = supports_num_moe_layers
         mock_report_l2_norm_grad.return_value = {}
         mock_report_throughput.return_value = {}
         mock_report_runtime.return_value = {}
@@ -1139,7 +1152,10 @@ def test_moe_logging(
         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
+        if supports_num_moe_layers:
+            assert call_args.kwargs["num_moe_layers"] == expected_num_moe_layers
+        else:
+            assert "num_moe_layers" not in call_args.kwargs
         assert call_args.kwargs["mtp_num_layers"] == mtp_num_layers
 
     @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches")