[Auto Parallel] support global mesh output with pipeline (#69628)

deepllz · web-flow · commit 9e405f7d0456 · 2024-11-28T10:58:58.000+08:00
* [Auto Parallel] add global mesh plan

* [Auto Parallel] support global mesh output with pipeline

* [Auto Parallel] support global mesh output with pipeline

* [Auto Parallel] support global mesh output with pipeline
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallelize.py b/python/paddle/distributed/auto_parallel/intermediate/parallelize.py
@@ -26,13 +26,13 @@ def parallelize(
     if pp_config is not None:
         assert isinstance(pp_config, dict)
         model, optimizer = pipeline_parallel(
-            model, optimizer, pp_config.get('split_spec')
+            model,
+            optimizer,
+            pp_config,
         )
     if mp_config is not None:
         assert isinstance(mp_config, dict)
-        model, optimizer = tensor_parallel(
-            model, optimizer, mp_config.get('parallelize_plan')
-        )
+        model, optimizer = tensor_parallel(model, optimizer, mp_config)
     if dp_config is not None:
         assert isinstance(dp_config, dict)
         if 'sharding_level' not in dp_config.keys():
diff --git a/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import itertools
+import logging
 import re
 from collections import OrderedDict
 from enum import Enum
@@ -32,10 +33,21 @@ class SplitPoint(Enum):
 
 
 class PipelineParallel(ParallelModel):
-    def __init__(self, model, split_spec):
+    def __init__(self, model, split_spec, global_spec, pipeline_layers):
         super().__init__(model)
         self.split_spec = split_spec
+        self.global_spec = global_spec
+        self.pipeline_layers = pipeline_layers
         self.pp_parallelizer = self.pipeline_parallel_fn
+        self.name_to_layer = {}
+        for layer_name, layer in model.named_sublayers():
+            self.name_to_layer[layer_name] = layer
+
+    def get_layer_by_name(self, name):
+        assert (
+            name in self.name_to_layer
+        ), f"layer name:{name} not in the model, please check the split_spec"
+        return self.name_to_layer[name]
 
     def pipeline_parallel_fn(self, model):
         mesh = fleet.auto.get_mesh()
@@ -46,12 +58,6 @@ def pipeline_parallel_fn(self, model):
         for layer_name, layer in model.named_sublayers():
             name_to_layer[layer_name] = layer
 
-        def get_layer_by_name(name):
-            assert (
-                name in name_to_layer
-            ), f"layer name:{name} not in the model, please check the split_spec"
-            return name_to_layer[name]
-
         def forward_post_hook(layer, input, output):
             pipeline_stage_index = layer.pipeline_stage_index
             split_point = layer.split_point
@@ -86,6 +92,7 @@ def forward_post_hook(layer, input, output):
             return output
 
         def forward_pre_hook(layer, input):
+            split_point = layer.split_point
             assert split_point == SplitPoint.BEGINNING
             # TODO(deepllz): support in the future
             return input
@@ -120,7 +127,7 @@ def forward_pre_hook(layer, input):
 
         # step2: insert reshard
         for name in split_layer_names:
-            layer = get_layer_by_name(name)
+            layer = self.get_layer_by_name(name)
             split_point = self.split_spec[name]
             layer.split_point = split_point
             if split_point == SplitPoint.END:
@@ -130,45 +137,104 @@ def forward_pre_hook(layer, input):
                     "SplitPoint.BEGINNING is not supported currently"
                 )
                 layer.register_forward_pre_hook(forward_pre_hook)
-
+        if self.global_spec:
+            self.process_global_mesh_layers()
         return model
 
+    def process_global_mesh_layers(self):
+        g_mesh = fleet.auto.get_mesh()
+        g_mesh = g_mesh.get_mesh_with_dim("pp")
+
+        def forward_post_hook(layer, input, output):
+            if isinstance(output, (list, tuple)):
+                global_output = list(output)
+                for ind in range(len(global_output)):
+                    if is_tensor(global_output[ind]):
+                        global_output[ind] = dist.shard_tensor(
+                            global_output[ind],
+                            g_mesh,
+                            [
+                                dist.Replicate()
+                                for _ in range(len(g_mesh._shape))
+                            ],
+                        )
+                if isinstance(output, tuple):
+                    global_output = tuple(global_output)
+                return global_output
+            elif is_tensor(output):
+                return dist.shard_tensor(
+                    output,
+                    g_mesh,
+                    [dist.Replicate() for _ in range(len(g_mesh._shape))],
+                )
+            else:
+                raise TypeError(
+                    "layer output can only be tensor or list/tuple of tensor"
+                )
+
+        def forward_pre_hook(layer, input):
+            pp_idx = getattr(layer, "pipeline_stage_index", 0)
+            new_input = []
+            for t in input:
+                if is_tensor(t) and t.is_dist() and t.process_mesh == g_mesh:
+                    new_input.append(
+                        dist.reshard(
+                            t,
+                            self.get_mesh(pp_idx),
+                            [dist.Replicate(), dist.Replicate()],
+                        )
+                    )
+                else:
+                    new_input.append(t)
+            return tuple(new_input)
+
+        for layer_name in self.global_spec:
+            layer = self.get_layer_by_name(layer_name)
+            layer.register_forward_post_hook(forward_post_hook)
+
+        for layer_name in self.pipeline_layers:
+            layer = self.get_layer_by_name(layer_name)
+            layer.register_forward_pre_hook(forward_pre_hook)
 
-def pipeline_parallel(model, optimizer, split_spec, mesh=None, dimension=None):
+
+def pipeline_parallel(model, optimizer=None, config=None):
     """
     pipeline_parallel converts model and optimizer to pipelined distributed model
 
     Args:
         model (paddle.nn.Layer): A single card model to be distributed
         optimizer (paddle.optimizer.Optimizer): An optimizer to be distributed
-        split_spec (OrderedDict|dict|str|list(str)): The pipeline parallel split point.
-            if split_spec is a string or list, such as "llama.layer" or ["llama.layerA", "llama.layerB"], Then the layer with same prefix a will be divided equally according to the size of pipeline degree.
-            if split_spec is a OrderedDict|dict, key is the layer name, and the value is the split position that can be SplitPoint.BEGINNING or SplitPoint.END, the order of the keys is the order of the pipeline stage.
-            NOTE: dict is also ordered after python3.7, so use dict at this time.
-        the order of the keys is the order of the pipeline stage
-        mesh (ProcessMesh): A ProcessMesh Object.
-        dimension (int|str): The mesh dimension to pipeline the model.
+        config (dict): {
+            "split_spec": OrderedDict|dict|str|list(str), The pipeline parallel split point.
+                if split_spec is a string or list, such as "llama.layer" or ["llama.layerA", "llama.layerB"], Then the layer with same prefix a will be divided equally according to the size of pipeline degree.
+                if split_spec is a OrderedDict|dict, key is the layer name, and the value is the split position that can be SplitPoint.BEGINNING or SplitPoint.END, the order of the keys is the order of the pipeline stage.
+                NOTE: dict is also ordered after python3.7, so use dict at this time.
+            "global_spec": str|list(str), make the output tensor of specific layers on global mesh.
+        }
 
     Returns:
         PipelineParallel: a distributed model
         ParallelOptimizer: a distributed optimizer
     """
-    if mesh is None:
-        mesh = fleet.auto.get_mesh()
-        assert (
-            mesh is not None
-        ), "global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
-        assert (
-            "pp" in mesh.dim_names
-        ), "pp must in the mesh dim_names when use pipeline_parallel"
-    else:
-        assert NotImplementedError(
-            "Specifying a custom mesh is not supported currently"
-        )
 
+    split_spec = config.get("split_spec")
+    if split_spec is None:
+        logging.warning("No split_spec, pipeline parallel won't do anything.")
+        return model, optimizer
+
+    mesh = fleet.auto.get_mesh()
+    assert (
+        mesh is not None
+    ), "global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
+    assert (
+        "pp" in mesh.dim_names
+    ), "pp must in the mesh dim_names when use pipeline_parallel"
+
+    global_spec = config.get("global_spec")
     if isinstance(split_spec, str):
         split_spec = [split_spec]
 
+    matched_layer_name = None
     if isinstance(split_spec, (list, tuple)):
         # match layer_name with split_spec following by a dot and numbers and no other characters
         # such as split_spec = ["llama.layer"], then llama.layer.0 is matched, llama.layer.0.mlp is not matched
@@ -204,10 +270,25 @@ def is_match(layer_name):
         )
     else:
         split_spec_dict = split_spec
+        if global_spec:
+            raise NotImplementedError(
+                "global_spec should be None if split_spec is a dict"
+            )
+
+    if isinstance(global_spec, str):
+        global_spec = [global_spec]
+    else:
+        assert isinstance(
+            global_spec, (list, tuple)
+        ), f"global_spec can only be list or list(str), but got:{type(global_spec)}"
 
-    logger.info(f"split_spec_dict: {split_spec_dict}")
+    logger.info(
+        f"split_spec_dict: {split_spec_dict}, global_spec: {global_spec}"
+    )
 
-    model = PipelineParallel(model, split_spec_dict)
+    model = PipelineParallel(
+        model, split_spec_dict, global_spec, matched_layer_name
+    )
     if optimizer is not None:
         optimizer = ParallelOptimizer(optimizer)
 
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -395,15 +395,18 @@ def tensor_parallelizer_fn(self, model):
         return model
 
 
-def tensor_parallel(model, optimizer=None, parallelize_plan=None):
+def tensor_parallel(model, optimizer=None, config=None):
     """
     Tensor parallel.
-    :param model: paddle.nn.Layer, the model to be shard into tensor parallel.
-    :param parallelize_plan: Dict, the plan to shard the layer.
-    :param optimizer: paddle.optimizer.Optimizer, the optimizer.
-    :return:
-        model: model after sharding
-        optimizer: optimizer after sharding
+    Args:
+        model (paddle.nn.Layer): the model to be shard into tensor parallel.
+        optimizer (paddle.optimizer.Optimizer): the optimizer.
+        config (dict): {
+            "parallelize_plan": dict, the plan to shard the layer.
+        }
+    Returns:
+        model: model after tp
+        optimizer: optimizer after tp
 
     NOTE: the plan should be a dict maps layer name or parameter name to a split_plan,
     which will be used to split the layer or the parameter. The name can be written in regular format.
@@ -423,6 +426,7 @@ def tensor_parallel(model, optimizer=None, parallelize_plan=None):
     }
     ```
     """
+    parallelize_plan = config.get("parallelize_plan")
     if parallelize_plan is None:
         # Do nothing if no plan.
         logging.warning(
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -210,7 +210,10 @@ def parallel_model(self, layer):
             #     f"llama.layers.{i * decoders_per_rank - 1}": SplitPoint.END
             #     for i in range(1, self.pp)
             # }
-            pp_config = {'split_spec': "llama.layers"}
+            pp_config = {
+                'split_spec': "llama.layers",
+                "global_spec": "llama.global_layer",
+            }
         if self.dp > 1:
             dp_config = {'sharding_level': self.level}
         if self.mp > 1:
diff --git a/test/auto_parallel/hybrid_strategy/single_llama_model.py b/test/auto_parallel/hybrid_strategy/single_llama_model.py
@@ -136,8 +136,8 @@ def __init__(self, config):
         self.input_layernorm = LlamaRMSNorm(self.config)
         self.post_attention_layernorm = LlamaRMSNorm(self.config)
 
-    def forward(self, hidden_states):
-        residual = hidden_states
+    def forward(self, hidden_states, global_tensor):
+        residual = hidden_states + global_tensor
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states)
         hidden_states = residual + hidden_states
@@ -150,6 +150,19 @@ def forward(self, hidden_states):
         return hidden_states
 
 
+class GlobalOutputNet(nn.Layer):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+
+    def forward(self, input):
+        return (
+            input
+            if input is not None
+            else paddle.rand([self.config.hidden_size], dtype="float32")
+        )
+
+
 class LlamaModel(nn.Layer):
     def __init__(self, config, position_embedding=False):
         super().__init__()
@@ -171,6 +184,8 @@ def __init__(self, config, position_embedding=False):
             else None
         )
 
+        self.global_layer = GlobalOutputNet(self.config)
+
         decoder_layers = []
         for i in range(self.config.num_hidden_layers):
             decoder_layers.append(LlamaDecoderLayer(self.config))
@@ -187,8 +202,10 @@ def forward(self, input_ids):
             position_embeddings = self.position_embedding(position_ids)
             hidden_states = hidden_states + position_embeddings
 
+        global_tensor = self.global_layer(None)
+
         for idx, (decoder_layer) in enumerate(self.layers):
-            hidden_states = decoder_layer(hidden_states)
+            hidden_states = decoder_layer(hidden_states, global_tensor)
 
         hidden_states = self.norm(hidden_states)
 
