[fix] add testcase with microbatch 4;

hpcaitech · Sep 9, 2024 · 6c2a120 · 6c2a120
1 parent 8366a78
commit 6c2a120
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/tests/test_pipeline/test_schedule/test_zerobubble_pp.py b/tests/test_pipeline/test_schedule/test_zerobubble_pp.py
@@ -509,15 +509,15 @@ def criterion(x, *args, **kwargs):
             "precision": "bf16",
             "num_model_chunk": 2,
         },
-        # {
-        #     "batch_size": 8,
-        #     "tp_size": 1,
-        #     "pp_size": 4,
-        #     "num_microbatches": 8,
-        #     "zero_stage": 1,
-        #     "precision": "bf16",
-        #     "num_model_chunk": 2,
-        # },
+        {
+            "batch_size": 8,
+            "tp_size": 1,
+            "pp_size": 4,
+            "num_microbatches": 8,
+            "zero_stage": 1,
+            "precision": "bf16",
+            "num_model_chunk": 2,
+        },
     ],
 )
 def run_fwd_bwd_vschedule_with_optim(test_config):