Update megatron/tests/unit_tests/transformer/moe/test_a2a_token_dispa…

…tcher.py
FlagOpen · Dec 9, 2024 · 9705478 · 9705478
1 parent 6f9a32a
commit 9705478
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/megatron/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py b/megatron/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
+import os
 import pytest
 import torch
 
@@ -66,10 +67,11 @@ def test_capacity_forward_backward(self, tp_size, ep_size):
         )
         container.dispacher_capacity_test()
 
+    # Skip because not running in internal and flaky
+    @pytest.mark.skipif(os.getenv('flagscale_skip') == '1', reason="flagscale_skip is enabled, skipping test.")
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.internal
     @pytest.mark.timeout(120)
-    @pytest.mark.internal
     @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2), (1, 1)])
     @pytest.mark.flaky
     def test_capacity_padding_forward_backward(self, tp_size, ep_size):