add very dumb layers splitting

Author: 3outeille

src/transformers/distributed/mixin.py

@@ -28,13 +28,13 @@
 from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
 from .fsdp import is_fsdp_managed_module
-from .pipeline_parallel import initialize_pipeline_parallelism
 from .utils import (
     _get_torch_distributed_rank,
     _is_torch_distributed_initialized,
     distribute_model,
     gather_full_state_dict,
     initialize_fully_sharded_data_parallelism,
+    initialize_pipeline_parallelism,
     save_model_checkpoint_distributed,
 )
 

src/transformers/distributed/pipeline_parallel.py

@@ -13,45 +13,61 @@
 # limitations under the License.
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING
 
-from ..utils import is_torch_available, is_torch_greater_or_equal
-from .utils import _ensure_torch_distributed
-
+from ..utils import is_torch_available
 
 if TYPE_CHECKING:
-    from .configuration_utils import DistributedConfig
+    import torch.nn as nn
 
 if is_torch_available():
     import torch
+    import torch.nn as nn
+
+
+class PPMissingLayer(nn.Identity):
+    """A placeholder layer for missing layers in a pipeline parallel model."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        """Return the first arg from args or the first value from kwargs."""
+        return args[0] if args else next(iter(kwargs.values()))
+
+
+def apply_pipeline_parallelism(model: nn.Module, device_mesh: torch.distributed.device_mesh.DeviceMesh) -> nn.Module:
+    """Naive even split of ``base_model.layers`` across PP ranks."""
+    pp_size = device_mesh.size()
+    if pp_size <= 1:
+        return model
+
+    pp_rank = device_mesh.get_local_rank()
+    is_first_rank = pp_rank == 0
+    is_last_rank = pp_rank == pp_size - 1
 
+    base_model = getattr(model, model.base_model_prefix)
+    layers = base_model.layers
+    num_layers = len(layers)
 
-def initialize_pipeline_parallelism(
-    distributed_config: DistributedConfig,
-):
-    if not is_torch_greater_or_equal("2.5"):
-        raise OSError("Pipeline parallelism with DistributedConfig requires `torch>=2.5`.")
+    layers_per_rank = num_layers // pp_size
+    start_layer = pp_rank * layers_per_rank
+    end_layer = num_layers if is_last_rank else start_layer + layers_per_rank
 
-    device_type = torch._C._get_accelerator().type
-    _ensure_torch_distributed(device_type)
+    if not is_first_rank:
+        base_model.embed_tokens = PPMissingLayer()
 
-    world_size = torch.distributed.get_world_size()
-    pp_size = distributed_config.pp_size
-    if world_size != pp_size:
-        raise RuntimeError(f"world_size ({world_size}) must be equal to pp_size ({pp_size})")
+    for i in range(num_layers):
+        if i < start_layer or i >= end_layer:
+            layers[i] = PPMissingLayer()
 
-    if device_type != "cpu":
-        local_rank = int(os.environ.get("LOCAL_RANK", 0))
-        getattr(torch, device_type).set_device(local_rank)
-        device_map = torch.device(device_type, local_rank)
-    else:
-        device_map = torch.device(device_type)
-    
+    if not is_last_rank:
+        base_model.norm = PPMissingLayer()
+        model.lm_head = PPMissingLayer()
 
-    assert world_size == pp_size, f"world_size ({world_size}) must be equal to pp_size ({pp_size})"
-    mesh = torch.distributed.init_device_mesh(device_type, (pp_size,), mesh_dim_names=("pp",))
+    model._pp_rank = pp_rank
+    model._pp_size = pp_size
+    return model
 
-    return device_map, mesh
 
 #TODO(3outeille): probably have to introduce pipeline_communicate here ?

src/transformers/distributed/utils.py

@@ -19,6 +19,7 @@
 from ..integrations.tensor_parallel import apply_tensor_parallelism
 from ..utils import is_torch_available, is_torch_greater_or_equal
 from .fsdp import apply_fully_sharded_data_parallelism
+from .pipeline_parallel import apply_pipeline_parallelism
 
 
 if TYPE_CHECKING:
@@ -153,6 +154,34 @@ def initialize_fully_sharded_data_parallelism(distributed_config: DistributedCon
     return device_map, mesh
 
 
+def initialize_pipeline_parallelism(
+    distributed_config: DistributedConfig,
+):
+    if not is_torch_greater_or_equal("2.5"):
+        raise OSError("Pipeline parallelism with DistributedConfig requires `torch>=2.5`.")
+
+    device_type = torch._C._get_accelerator().type
+    _ensure_torch_distributed(device_type)
+
+    world_size = torch.distributed.get_world_size()
+    pp_size = distributed_config.pp_size
+    if world_size != pp_size:
+        raise RuntimeError(f"world_size ({world_size}) must be equal to pp_size ({pp_size})")
+
+    if device_type != "cpu":
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        getattr(torch, device_type).set_device(local_rank)
+        device_map = torch.device(device_type, local_rank)
+    else:
+        device_map = torch.device(device_type)
+    
+
+    assert world_size == pp_size, f"world_size ({world_size}) must be equal to pp_size ({pp_size})"
+    mesh = torch.distributed.init_device_mesh(device_type, (pp_size,), mesh_dim_names=("pp",))
+
+    return device_map, mesh
+
+
 def distribute_model(
     model,
     distributed_config: DistributedConfig,
@@ -172,6 +201,9 @@ def distribute_model(
     elif distributed_config.fsdp_size > 1:
         fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
         model = apply_fully_sharded_data_parallelism(model, fsdp_mesh)
+    elif distributed_config.pp_size > 1:
+        pp_mesh = device_mesh["pp"] if device_mesh.ndim > 1 else device_mesh
+        model = apply_pipeline_parallelism(model, pp_mesh)
 
     return model