feat: add virtual_pipeline_parallel test scripts

Author: JYMiracle305

infini_train/src/nn/parallel/pp/pipeline_schedule.cc

@@ -194,8 +194,10 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
 
     auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule(n, num_stages, vpp_size);
 
-    if (stage_idx == 0) {
+    static bool has_printed = false;
+    if (!has_printed && stage_idx == 0) {
         PrintScheduleTable(schedule, n, num_stages, vpp_size);
+        has_printed = true;
     }
 
     float total_loss = 0.0f;

scripts/test_config.json

@@ -141,6 +141,30 @@
         },
         {
             "id": "7",
+            "args": {
+                "dtype": "float32",
+                "nthread_per_process": 4,
+                "num_iteration": 10,
+                "batch_size": 10,
+                "total_batch_size": 5120,
+                "pipeline_parallel": 4,
+                "virtual_pipeline_parallel": 2
+            }
+        },
+        {
+            "id": "7_bfloat16",
+            "args": {
+                "dtype": "bfloat16",
+                "nthread_per_process": 4,
+                "num_iteration": 10,
+                "batch_size": 10,
+                "total_batch_size": 5120,
+                "pipeline_parallel": 4,
+                "virtual_pipeline_parallel": 2
+            }
+        },
+        {
+            "id": "8",
             "args": {
                 "dtype": "float32",
                 "nthread_per_process": 8,
@@ -149,11 +173,12 @@
                 "total_batch_size": 5120,
                 "tensor_parallel": 2,
                 "sequence_parallel": true,
-                "pipeline_parallel": 2
+                "pipeline_parallel": 2,
+                "virtual_pipeline_parallel": 2
             }
         },
         {
-            "id": "7_bfloat16",
+            "id": "8_bfloat16",
             "args": {
                 "dtype": "bfloat16",
                 "nthread_per_process": 8,
@@ -162,7 +187,8 @@
                 "total_batch_size": 5120,
                 "tensor_parallel": 2,
                 "sequence_parallel": true,
-                "pipeline_parallel": 2
+                "pipeline_parallel": 2,
+                "virtual_pipeline_parallel": 2
             }
         }
     ]