Add microbatch_group_size_per_vp_stage as configurable.

README.md

@@ -72,6 +72,8 @@ uv run python main.py strategy=interleave num_devices=4 num_stages=8 num_batches
 ```
 ![interleave](assets/interleave_1f1b.png)
 
+You can optionally setting `microbatch_group_size_per_vp_stage`.
+
 ### Running for ZB-1P strategy:
 ```bash
 uv run python main.py strategy=zb1p num_devices=4 num_stages=4 num_batches=8

conf/config.yaml

@@ -5,6 +5,7 @@ num_batches: 8
 visualization_port: 8050
 strategy: "1f1b"  # Options: "1f1b", "interleave"
 p2p_latency: 0.0
+microbatch_group_size_per_vp_stage: null
 
 # Operation time configurations
 op_times:

main.py

@@ -71,6 +71,7 @@ def run_interleave(cfg: DictConfig) -> None:
         p2p_latency=cfg.p2p_latency,
         placement_strategy="interleave",
         op_times=op_times,
+        microbatch_group_size_per_vp_stage=cfg.microbatch_group_size_per_vp_stage,
     )
     schedule = generate_1f1b_interleave_schedule(schedule_config)
     schedule.execute()

src/execution_model.py

@@ -83,6 +83,7 @@ class ScheduleConfig:
         placement_strategy: str = "standard",
         split_backward: bool = False,
         op_times: Optional[Dict[str, Union[float, Dict[int, float]]]] = None,
+        microbatch_group_size_per_vp_stage: Optional[int] = None,
     ):
         self.num_devices = num_devices
         self.num_stages = num_stages
@@ -90,6 +91,10 @@ class ScheduleConfig:
         self.p2p_latency = p2p_latency
         self.placement_strategy = placement_strategy
         self.split_backward = split_backward
+        if microbatch_group_size_per_vp_stage is None:
+            self.microbatch_group_size_per_vp_stage = num_devices
+        else:
+            self.microbatch_group_size_per_vp_stage = microbatch_group_size_per_vp_stage
 
         # Initialize default operation times
         if self.split_backward:

src/strategies.py

@@ -244,7 +244,7 @@ def generate_1f1b_interleave_schedule(config: ScheduleConfig):
     schedule = Schedule(config)
 
     for device_id in range(config.num_devices):
-        microbatch_group_size_per_vp_stage = config.num_devices
+        microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage
         num_warmup_microbatches = _get_pp_rank_microbatches(
             config.num_batches,
             config.num_devices,