Add DualPipeV

Author: H-Huang

torchtitan/distributed/pipeline_parallel.py

@@ -19,6 +19,7 @@
     PipelineScheduleMulti,
     PipelineScheduleSingle,
     ScheduleZBVZeroBubble,
+    ScheduleDualPipeV,
 )
 
 from torchtitan.config import JobConfig
@@ -335,7 +336,7 @@ def _build_stage_from_modules(
     models = []
 
     schedule_class = get_schedule_class(pp_schedule)
-    style = "v" if schedule_class == ScheduleZBVZeroBubble else "loop"
+    style = "v" if schedule_class in (ScheduleZBVZeroBubble, ScheduleDualPipeV) else "loop"
 
     for stage_idx in stage_ids_this_rank(pp_rank, pp_size, num_stages, style=style):
         module_names = module_names_per_stage[stage_idx]