fix bugs in PP

Author: junjzhang

torchtitan/experiments/train_llama_hf/model/__init__.py

@@ -6,13 +6,13 @@
 #
 # Copyright (c) Meta Platforms, Inc. All Rights Reserved.
 
-from loss import cross_entropy_loss_hf
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 from torchtitan.components.optimizer import build_lr_schedulers, build_optimizers
 from torchtitan.experiments.train_llama_hf.dataset import (
     build_pos_included_hf_dataloader,
 )
+from torchtitan.experiments.train_llama_hf.loss import cross_entropy_loss_hf
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
 
 from .parallelize_llama import parallelize_llama

torchtitan/experiments/train_llama_hf/model/parallelize_llama.py

@@ -167,7 +167,11 @@ def apply_tp(
     # NOTE: At the cost of model code change, we can accelerate Sequence Parallel
     #       by folding (and unfolding) the batch dimension and the sequence dimension.
     #       Examples can be found at https://github.com/pytorch/torchtitan/pull/437
-    for transformer_block in model.model.layers:
+    if isinstance(model.model.layers, nn.ModuleDict):
+        transformer_blocks = model.model.layers.values()
+    else:
+        transformer_blocks = model.model.layers
+    for transformer_block in transformer_blocks:
         layer_plan = {
             "input_layernorm": SequenceParallel(),
             "self_attn": prepare_module_input(
@@ -260,8 +264,12 @@ def apply_fsdp(
     fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy}
     if cpu_offload:
         fsdp_config["offload_policy"] = CPUOffloadPolicy()
+    if isinstance(model.model.layers, nn.ModuleDict):
+        layer_items = [(int(k), v) for (k, v) in model.model.layers.items()]
+    else:
+        layer_items = list(enumerate(model.model.layers))
 
-    for layer_id, transformer_block in enumerate(model.model.layers):
+    for layer_id, transformer_block in layer_items:
         if reshard_after_forward_policy == "always":
             reshard_after_forward = True
         elif reshard_after_forward_policy == "never":
@@ -274,7 +282,7 @@ def apply_fsdp(
             else:
                 # As an optimization, do not reshard after forward for the last
                 # transformer block since FSDP would prefetch it immediately
-                reshard_after_forward = int(layer_id) < len(model.model.layers) - 1
+                reshard_after_forward = layer_id < len(layer_items) - 1
         else:
             raise ValueError(
                 f"Invalid reshard_after_forward_policy: {reshard_after_forward_policy}."

torchtitan/experiments/train_llama_hf/model/pipeline_llama.py

@@ -29,7 +29,6 @@
 )
 from torchtitan.tools.logging import logger
 
-
 DeviceType = Union[int, str, torch.device]
 
 
@@ -87,8 +86,10 @@ def forward(
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
-        # decoder layers
-        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+        # decoder layers, ok since ModuleDict is ordered
+        for decoder_layer in list(self.layers.values())[
+            : self.config.num_hidden_layers
+        ]:
 
             if self.gradient_checkpointing and self.training:
                 layer_outputs = self._gradient_checkpointing_func(
@@ -217,6 +218,10 @@ def pipeline_llama(
     model_config: PretrainedConfig,
     loss_fn: Callable[..., torch.Tensor],
 ) -> tuple[_PipelineSchedule, list[nn.Module], bool, bool]:
+    logger.info("Changing model.model.layers to nn.ModuleDict")
+    model.model.layers = nn.ModuleDict(
+        {str(i): layer for i, layer in enumerate(model.model.layers)}
+    )
     logger.info(
         "Patching Llama forward method for pipeline parallelism, it will disable some features of orignal HF model"
     )
@@ -277,20 +282,14 @@ def _build_stage(
             model.model.embed_tokens = None
 
         drop_layers = start_layer is not None
-        del_indexes = []
-        for i in range(len(model.model.layers)):
+        for name in list(model.model.layers.keys()):
             # we keep layers in a contiguous region between start (inclusive) and stop (exclusive)
-            if f"layers.{i}" == start_layer:
+            if f"layers.{name}" == start_layer:
                 drop_layers = False
-            if f"layers.{i}" == stop_layer:
+            if f"layers.{name}" == stop_layer:
                 drop_layers = True
             if drop_layers:
-                del_indexes.append(i)
-
-        # delete layers in reverse order to avoid index shifting
-        del_indexes.reverse()
-        for i in del_indexes:
-            del model.model.layers[i]
+                del model.model.layers[name]
 
         if not is_last:
             model.model.norm = None

torchtitan/experiments/train_llama_hf/test_loading_hf_weights_helper.py

@@ -16,13 +16,13 @@
     normalize_state_dict_key,
 )
 
+from torchtitan.experiments.train_llama_hf.loss import cross_entropy_loss_hf
+
 from torchtitan.experiments.train_llama_hf.model.parallelize_llama import (
     apply_fsdp,
     apply_tp,
 )
-from torchtitan.experiments.train_llama_hf.model.pipeline_llama import (
-    pipeline_llama_manual_split,
-)
+from torchtitan.experiments.train_llama_hf.model.pipeline_llama import pipeline_llama
 
 
 def main(job_config: JobConfig):
@@ -52,13 +52,14 @@ def main(job_config: JobConfig):
     # apply parallelisms
     if parallel_dims.pp_enabled:
         # apply PT-D Pipeline Parallel
-        _, model_parts = pipeline_llama_manual_split(
+        _, model_parts, _, _ = pipeline_llama(
             model,
             world_mesh["pp"],
             parallel_dims,
             job_config,
             device,
             model_config,
+            loss_fn=cross_entropy_loss_hf,
         )
     else:
         model_parts = [model]