[training] fix: use startswith for Blackwell GPU name validation in DeepEP (#2731)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>

Author: yaoyu-33

src/megatron/bridge/training/flex_dispatcher_backend.py

@@ -45,10 +45,13 @@ def apply_flex_dispatcher_backend(
 
     device_properties = torch.cuda.get_device_properties(0)
     if moe_flex_dispatcher_backend == "deepep":
-        if not (device_properties.major in [8, 9] or device_properties.name in ["NVIDIA B200", "NVIDIA B300"]):
+        if not (
+            device_properties.major in [8, 9] or device_properties.name.startswith(("NVIDIA B200", "NVIDIA B300"))
+        ):
             if get_rank_safe() == 0:
                 logger.warning(
-                    "DeepEP is only applicable to Ampere, Hopper, and Blackwell (only B200 and B300) GPUs. Skipping DeepEP configuration."
+                    f"DeepEP is only applicable to Ampere, Hopper, and Blackwell (B200/B300) GPUs. "
+                    f"Current GPU: {device_properties.name}. Skipping DeepEP configuration."
                 )
             return
     elif moe_flex_dispatcher_backend == "hybridep":
@@ -74,8 +77,13 @@ def validate_flex_dispatcher_backend(model_config: TransformerConfig) -> None:
     if model_config.moe_token_dispatcher_type == "flex":
         device_properties = torch.cuda.get_device_properties(0)
         if model_config.moe_flex_dispatcher_backend == "deepep":
-            if not (device_properties.major in (8, 9) or device_properties.name in ["NVIDIA B200", "NVIDIA B300"]):
-                raise ValueError("DeepEP is supported for Ampere, Hopper, and Blackwell (only B200 and B300) GPUs")
+            if not (
+                device_properties.major in (8, 9) or device_properties.name.startswith(("NVIDIA B200", "NVIDIA B300"))
+            ):
+                raise ValueError(
+                    f"DeepEP is supported for Ampere, Hopper, and Blackwell (B200/B300) GPUs. "
+                    f"Current GPU: {device_properties.name}"
+                )
 
         if model_config.moe_flex_dispatcher_backend == "hybridep":
             if not device_properties.major in [8, 9, 10]:

tests/unit_tests/training/test_config.py

@@ -1035,25 +1035,34 @@ def test_context_parallel_finetuning_validations(
             restore_get_world_size_safe(og_ws, cfg_mod)
 
     @pytest.mark.parametrize(
-        "gpu_major, moe_enable_deepep, expect_error",
+        "gpu_major, gpu_name, moe_enable_deepep, expect_error",
         [
-            (8, True, False),  # Ampere GPU with DeepEP enabled - should pass
-            (9, True, False),  # Hopper GPU with DeepEP enabled - should pass
-            (7, True, True),  # Volta GPU with DeepEP enabled - should raise ValueError
-            (6, True, True),  # Pascal GPU with DeepEP enabled - should raise ValueError
-            (10, True, True),  # Future unsupported GPU with DeepEP enabled - should raise ValueError
-            (7, False, False),  # Volta GPU with DeepEP disabled - should pass
-            (6, False, False),  # Pascal GPU with DeepEP disabled - should pass
+            (8, "NVIDIA A100", True, False),  # Ampere GPU with DeepEP enabled - should pass
+            (9, "NVIDIA H100", True, False),  # Hopper GPU with DeepEP enabled - should pass
+            (10, "NVIDIA B200", True, False),  # Blackwell B200 GPU with DeepEP enabled - should pass
+            (10, "NVIDIA B200 SXM6 AC", True, False),  # Blackwell B200 variant with DeepEP enabled - should pass
+            (10, "NVIDIA B300", True, False),  # Blackwell B300 GPU with DeepEP enabled - should pass
+            (7, "NVIDIA V100", True, True),  # Volta GPU with DeepEP enabled - should raise ValueError
+            (6, "NVIDIA P100", True, True),  # Pascal GPU with DeepEP enabled - should raise ValueError
+            (
+                10,
+                "NVIDIA B100",
+                True,
+                True,
+            ),  # Unsupported Blackwell variant with DeepEP enabled - should raise ValueError
+            (7, "NVIDIA V100", False, False),  # Volta GPU with DeepEP disabled - should pass
+            (6, "NVIDIA P100", False, False),  # Pascal GPU with DeepEP disabled - should pass
         ],
     )
     @patch("torch.cuda.get_device_properties")
     def test_deepep_validation(
-        self, mock_get_device_properties, monkeypatch, gpu_major, moe_enable_deepep, expect_error
+        self, mock_get_device_properties, monkeypatch, gpu_major, gpu_name, moe_enable_deepep, expect_error
     ):
         """Test DeepEP validation during config container validation."""
         # Mock GPU device properties
         mock_properties = MagicMock()
         mock_properties.major = gpu_major
+        mock_properties.name = gpu_name
         mock_get_device_properties.return_value = mock_properties
 
         # Create a GPT model config with MoE settings
@@ -1069,7 +1078,7 @@ def test_deepep_validation(
 
         try:
             if expect_error:
-                with pytest.raises(ValueError, match="DeepEP is supported for Ampere"):
+                with pytest.raises(ValueError, match="DeepEP is supported for Ampere, Hopper, and Blackwell"):
                     container.validate()
             else:
                 container.validate()  # Should pass without error

tests/unit_tests/training/test_deepep.py

@@ -133,6 +133,7 @@ def test_apply_flex_dispatcher_backend_warns_for_unsupported_gpu_volta(
         # Mock Volta GPU (compute capability 7.x)
         mock_properties = MagicMock()
         mock_properties.major = 7
+        mock_properties.name = "NVIDIA V100"
         mock_get_device_properties.return_value = mock_properties
 
         # Create a mock TransformerConfig with MoE enabled
@@ -145,7 +146,7 @@ def test_apply_flex_dispatcher_backend_warns_for_unsupported_gpu_volta(
         # Verify warning was logged
         mock_logger.warning.assert_called_once()
         assert (
-            "DeepEP is only applicable to Ampere, Hopper, and Blackwell (only B200 and B300) GPUs"
+            "DeepEP is only applicable to Ampere, Hopper, and Blackwell (B200/B300) GPUs"
             in mock_logger.warning.call_args[0][0]
         )
 
@@ -162,6 +163,7 @@ def test_apply_flex_dispatcher_backend_warns_for_unsupported_gpu_pascal(
         # Mock Pascal GPU (compute capability 6.x)
         mock_properties = MagicMock()
         mock_properties.major = 6
+        mock_properties.name = "NVIDIA P100"
         mock_get_device_properties.return_value = mock_properties
 
         # Create a mock TransformerConfig with MoE enabled
@@ -174,7 +176,7 @@ def test_apply_flex_dispatcher_backend_warns_for_unsupported_gpu_pascal(
         # Verify warning was logged
         mock_logger.warning.assert_called_once()
         assert (
-            "DeepEP is only applicable to Ampere, Hopper, and Blackwell (only B200 and B300) GPUs"
+            "DeepEP is only applicable to Ampere, Hopper, and Blackwell (B200/B300) GPUs"
             in mock_logger.warning.call_args[0][0]
         )
 
@@ -249,6 +251,7 @@ def test_validate_flex_dispatcher_backend_volta_gpu_raises_error(self, mock_get_
         # Mock Volta GPU (compute capability 7.x)
         mock_properties = MagicMock()
         mock_properties.major = 7
+        mock_properties.name = "NVIDIA V100"
         mock_get_device_properties.return_value = mock_properties
 
         # Create a mock TransformerConfig with DeepEP enabled
@@ -259,7 +262,7 @@ def test_validate_flex_dispatcher_backend_volta_gpu_raises_error(self, mock_get_
         # Should raise ValueError
         with pytest.raises(
             ValueError,
-            match="DeepEP is supported for Ampere, Hopper, and Blackwell \\(only B200 and B300\\) GPUs",
+            match="DeepEP is supported for Ampere, Hopper, and Blackwell \\(B200/B300\\) GPUs",
         ):
             validate_flex_dispatcher_backend(config)
 
@@ -272,6 +275,7 @@ def test_validate_flex_dispatcher_backend_future_gpu_raises_error(self, mock_get
         # Mock future GPU
         mock_properties = MagicMock()
         mock_properties.major = 200
+        mock_properties.name = "NVIDIA Future GPU"
         mock_get_device_properties.return_value = mock_properties
 
         # Create a mock TransformerConfig with DeepEP enabled