Refactor pipeline_parallel.py for graph PP reuse

Author: Sanket Jayant Purandare

torchtitan/distributed/pipeline_parallel.py

@@ -43,6 +43,8 @@
     "build_pipeline_schedule",
     "generate_llm_fqn_per_model_part",
     "pipeline_module_split",
+    "split_module",
+    "get_pp_rank_to_stage_indices_mapping",
 ]
 
 
@@ -64,10 +66,10 @@ def pipeline_llm(
     pp_mesh = parallel_dims.get_mesh("pp")
 
     num_virtual_stages, num_layers, input_weight, output_weight = get_pipeline_metadata(
-        parallel_dims, job_config, model_args
+        parallel_dims, parallelism, model_config
     )
 
-    module_names_per_stage = job_config.parallelism.module_fqns_per_model_part
+    module_names_per_stage = parallelism.module_fqns_per_model_part
     if module_names_per_stage is None:
         module_names_per_stage = generate_llm_fqn_per_model_part(
             num_virtual_stages, num_layers, input_weight, output_weight
@@ -78,7 +80,7 @@ def pipeline_llm(
     stages, model_parts = pipeline_module_split(
         model,
         pp_mesh,
-        job_config.parallelism.pipeline_parallel_schedule,
+        parallelism.pipeline_parallel_schedule,
         device,
         module_names_per_stage,
     )
@@ -88,13 +90,27 @@ def pipeline_llm(
     # optimizer, and checkpointing
     for i, m in enumerate(model_parts):
         # apply SPMD-style PT-D techniques
-        m = parallelize_fn(m, parallel_dims, job_config)
+        m = parallelize_fn(
+            m,
+            parallel_dims=parallel_dims,
+            training=training,
+            model_converters=model_converters,
+            parallelism=parallelism,
+            compile_config=compile_config,
+            ac_config=ac_config,
+            dump_folder=dump_folder,
+        )
         model_parts[i] = m
         # NOTE: this is to update the model in the stage
         #       in case the model is modified e.g. by torch.compile
         stages[i].submod = m
 
-    pp_schedule = build_pipeline_schedule(job_config, stages, loss_fn)
+    pp_schedule = build_pipeline_schedule(
+        parallelism=parallelism,
+        local_batch_size=training.local_batch_size,
+        stages=stages,
+        loss_fn=loss_fn,
+    )
 
     # This is used in the train loop to determine whether to pass in the input_ids and labels
     has_first_stage = False
@@ -110,16 +126,16 @@ def pipeline_llm(
 
 def get_pipeline_metadata(
     parallel_dims: ParallelDims,
-    job_config: JobConfig,
-    model_args: BaseModelArgs,
+    parallelism: ParallelismConfig,
+    model_config: BaseModel.Config,
 ) -> tuple[int, int, int, int]:
     """
     Determine the number of virtual stages and the number of layers in the model.
 
     Args:
         parallel_dims (ParallelDims): Parallel dimensions.
-        job_config (JobConfig): Job configuration.
-        model_args (BaseModelArgs): Model arguments.
+        parallelism (ParallelismConfig): Parallelism configuration.
+        model_config (BaseModel.Config): Model configuration.
 
     Returns:
         tuple: A tuple containing the number of virtual stages, the number of layers in the model,
@@ -194,6 +210,8 @@ def build_pipeline_schedule(
     local_batch_size: int,
     stages: list[PipelineStage],
     loss_fn: Callable,
+    backward_requires_autograd: bool = True,
+    scale_grads: bool = False,
 ) -> _PipelineSchedule:
     """Builds a pipeline schedule for the given job configuration and stages.
 
@@ -242,7 +260,8 @@ def build_pipeline_schedule(
         stages if looped_schedule else stages[0],
         n_microbatches=n_microbatches,
         loss_fn=loss_fn,
-        scale_grads=False,
+        backward_requires_autograd=backward_requires_autograd,
+        scale_grads=scale_grads,
     )
     logger.info(
         f"Using pipeline schedule {parallelism.pipeline_parallel_schedule} "
@@ -403,7 +422,9 @@ def split_module(
     modules_to_keep = set(module_names)
     for module_name, module_value in model.named_children():
         # Handle layer-like structures (e.g., "layers.0", "layers.1")
-        if isinstance(module_value, (nn.ModuleDict, nn.ModuleList)):
+        if isinstance(
+            module_value, (nn.ModuleDict, nn.ModuleList, ModuleDict, ModuleList)
+        ):
             layers_to_keep = {
                 name.split(".", 1)[1]
                 for name in modules_to_keep
@@ -419,7 +440,7 @@ def split_module(
                     indices_to_keep = {
                         int(idx) for idx in layers_to_keep if idx.isdigit()
                     }
-                    new_layers = nn.ModuleList(
+                    new_layers = ModuleList(
                         [
                             layer
                             for i, layer in enumerate(module_value)
@@ -429,10 +450,10 @@ def split_module(
                     setattr(model, module_name, new_layers)
             else:
                 # No layers from this structure needed, set to empty structure
-                if isinstance(module_value, nn.ModuleDict):
-                    setattr(model, module_name, nn.ModuleDict())
-                elif isinstance(module_value, nn.ModuleList):
-                    setattr(model, module_name, nn.ModuleList())
+                if isinstance(module_value, (nn.ModuleDict, ModuleDict)):
+                    setattr(model, module_name, ModuleDict())
+                elif isinstance(module_value, (nn.ModuleList, ModuleList)):
+                    setattr(model, module_name, ModuleList())
         # Handle simple module attributes (e.g., "linear", "norm")
         elif module_name not in modules_to_keep:
             # Replace with None