Skip test_moe_deepseek_unscanned_bf16 test

AI-Hypercomputer · Mar 6, 2025 · af00004 · af00004
1 parent df7813f
commit af00004
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/MaxText/tests/train_compile_test.py b/MaxText/tests/train_compile_test.py
@@ -458,6 +458,7 @@ def test_moe_deepseek_scanned_bf16(self):
         )
     )
 
+  @pytest.mark.skip(reason="Fix sharding issue of all layers of DeepSeek")
   @pytest.mark.tpu_only
   def test_moe_deepseek_unscanned_bf16(self):
     compiled_trainstep_file = "/tmp/test_moe_deepseek_unscanned_bf16.pickle"
@@ -472,7 +473,7 @@ def test_moe_deepseek_unscanned_bf16(self):
             "model_name=deepseek3-671b",
             "sparse_matmul=True",
             "megablox=False",
-            "per_device_batch_size=2",
+            "per_device_batch_size=1",
             "max_target_length=1024",
             "attention=dot_product",  # Change to flush attention once it works for MLA
             "dtype=bfloat16",