Revert "reuse pipeline from torchtitan"

This reverts commit 09f0c94.

Author: 3outeille

torchtitan/distributed/pipeline_parallel.py

@@ -228,7 +228,6 @@ def generate_llm_fqn_per_model_part(
     num_layers: int,
     input_weight: int = 1,
     output_weight: int = 1,
-    include_rotary_emb: bool = False,
 ) -> list[list[str]]:
     """
     Programmatically generates module names model part, focused on LLMs models.
@@ -238,7 +237,6 @@ def generate_llm_fqn_per_model_part(
         num_layers: Total number of transformer layers in the model
         input_weight: Weight for input modules (tok_embeddings) in layer calculation
         output_weight: Weight for output modules (norm + output) in layer calculation
-        include_rotary_emb: Whether to include rotary_emb in each model part
 
     Returns:
         List of lists containing module names for each model part
@@ -253,10 +251,7 @@ def generate_llm_fqn_per_model_part(
     if num_stages == 1:
         # Single stage gets everything
         layer_names = [f"layers.{i}" for i in range(num_layers)]
-        result = [["tok_embeddings"] + layer_names + ["norm", "output"]]
-        if include_rotary_emb:
-            result[0].append("rotary_emb")
-        return result
+        return [["tok_embeddings"] + layer_names + ["norm", "output"]]
 
     # Calculate effective layers including weights
     num_effective_layers = num_layers + input_weight + output_weight
@@ -334,8 +329,6 @@ def generate_llm_fqn_per_model_part(
                     stage_modules.append(f"layers.{current_layer}")
                     current_layer += 1
 
-        if include_rotary_emb:
-            stage_modules.append("rotary_emb")
         module_names_per_stage.append(stage_modules)
 
     return module_names_per_stage
@@ -347,7 +340,6 @@ def pipeline_module_split(
     pp_schedule: str,
     device: torch.device,
     module_names_per_stage: list[list[str]],
-    use_identity_for_missing_modules: bool = False,
 ) -> tuple[list[PipelineStage], list[nn.Module]]:
     """
     This API creates pipeline stages based on specified module names for each stage.
@@ -369,8 +361,6 @@ def pipeline_module_split(
                                - "layers.0", "layers.1" for specific transformer layers
                                - "norm" for the final normalization layer
                                - "output" for the output projection layer
-        use_identity_for_missing_modules: If True, replace missing modules with nn.Identity(),
-                                         otherwise replace with None
 
     Returns:
         Tuple of (stages, models) where stages are PipelineStage objects and models are the
@@ -427,11 +417,8 @@ def _build_stage_from_modules(
                         setattr(model, module_name, nn.ModuleList())
             # Handle simple module attributes (e.g., "linear", "norm")
             elif module_name not in modules_to_keep:
-                # Replace with Identity or None based on configuration
-                replacement = (
-                    nn.Identity() if use_identity_for_missing_modules else None
-                )
-                setattr(model, module_name, replacement)
+                # Replace with None
+                setattr(model, module_name, None)
 
         stage = PipelineStage(
             model,

torchtitan/experiments/transformers_backend/__init__.py

@@ -12,7 +12,7 @@
 
 from .infra.parallelize import parallelize_hf_transformers
 
-from .infra.pipeline import pipeline_hf_transformers
+from .infra.pipeline_parallel import pipeline_hf_transformers
 from .model.args import HFTransformerModelArgs, TitanDenseModelArgs
 from .model.model import HFTransformerModel