[PP] Add microbatch_size and remove num_microbatches config (#1073)

landing #1050

Author: H-Huang

tests/integration_tests.py

@@ -156,7 +156,6 @@ def build_test_list():
                 [
                     "--parallelism.pipeline_parallel_degree 2",
                     "--parallelism.pipeline_parallel_schedule ZBVZeroBubble",
-                    "--parallelism.pipeline_parallel_microbatches 8",
                 ],
             ],
             "PP zero bubble test (v shaped)",
@@ -280,7 +279,6 @@ def build_test_list():
                     "--parallelism.pipeline_parallel_degree 2",
                     "--parallelism.pipeline_parallel_schedule PipelineScheduleMulti",
                     "--parallelism.pipeline_parallel_schedule_csv ./tests/assets/custom_schedule.csv",
-                    "--parallelism.pipeline_parallel_microbatches 8",
                 ],
             ],
             "PP with custom pipeline schedule loaded from CSV file",

torchtitan/config_manager.py

@@ -520,15 +520,16 @@ def __init__(self):
             """,
         )
         self.parser.add_argument(
-            "--parallelism.pipeline_parallel_microbatches",
+            "--parallelism.pipeline_parallel_microbatch_size",
             type=int,
-            default=None,
+            default=1,
             help="""
-                How many microbatches to split the global training batch into when using pipeline parallelism.
+                The size of each pipeline parallel microbatch (default 1).
 
-                The global training batch size must be evenly divisible by the number of microbatches.
+                This value is used to compute the total number of microbatches by dividing batch_size with
+                pipeline_parallel_microbatch_size.
 
-                The default value will be the number of pipeline stages, if unspecified.
+                The global training batch size must be evenly divisible by pipeline_parallel_microbatch_size.
             """,
         )
         self.parser.add_argument(

torchtitan/distributed/pipeline.py

@@ -139,24 +139,23 @@ def build_pipeline_schedule(
         )
 
     looped_schedule = issubclass(schedule_class, PipelineScheduleMulti)
-    n_microbatches = job_config.parallelism.pipeline_parallel_microbatches
+    microbatch_size = job_config.parallelism.pipeline_parallel_microbatch_size
+    batch_size = job_config.training.batch_size
+    # validate that the batch size is divisible by the microbatch_size otherwise we'll hang or error during training
+    if batch_size % microbatch_size != 0:
+        raise ValueError(
+            f"Batch size {job_config.training.batch_size} must be divisible by number of microbatches {n_microbatches}. "
+            "Update the config arguments for either batch_size or pipeline_parallel_microbatch_size."
+        )
+    n_microbatches = batch_size // microbatch_size
     # We expect that the number of local stages (`len(stages)`) is the same across all ranks
     num_total_stages = job_config.parallelism.pipeline_parallel_degree * len(stages)
-    if n_microbatches is None:
-        n_microbatches = num_total_stages
-    elif n_microbatches < num_total_stages:
+    if n_microbatches < num_total_stages:
         logger.warning(
             f"Number of microbatches ({n_microbatches}) is less than the total number "
             f"of stages ({num_total_stages}) which may result in a bubble in the pipeline."
         )
 
-    # validate that the batch size is divisible by the number of microbatches otherwise we'll hang or error during training
-    if job_config.training.batch_size % n_microbatches != 0:
-        raise ValueError(
-            f"Batch size {job_config.training.batch_size} must be divisible by number of microbatches {n_microbatches}. "
-            "Update the config arguments for either batch_size or pipeline_parallel_microbatches."
-        )
-
     schedule = schedule_class(
         stages if looped_schedule else stages[0],
         n_microbatches=n_microbatches,