[Auto Parallel] Sequence parallel api (#69288)

FeixLiu · web-flow · commit e9497649c0a1 · 2024-11-12T11:20:29.000+08:00
* Sequence parallel api

* bug fix

* reformat the code
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -156,6 +156,139 @@ def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
         layer.register_forward_post_hook(self.fn(process_mesh=process_mesh))
 
 
+def sp_split(x, process_mesh, need_transpose):
+    index = process_mesh.dim_names.index('mp')  # get the axis for the split
+    if isinstance(x, tuple):
+        target_x = x[0]
+    else:
+        target_x = x
+    assert isinstance(target_x, paddle.Tensor)
+    assert len(target_x.shape) == 3
+    if need_transpose:
+        target_x = paddle.transpose(target_x, perm=[1, 0, 2])
+    placements = target_x.placements
+    if placements is None:
+        placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
+    placements[index] = dist.Shard(0)
+    target_x = dist.reshard(target_x, process_mesh, placements)
+    if isinstance(x, tuple):
+        x = list(x)
+        x[0] = target_x
+        x = tuple(x)
+    else:
+        x = target_x
+
+    return x
+
+
+def sp_reduce_scatter(x, process_mesh, need_transpose):
+    index = process_mesh.dim_names.index('mp')  # get the axis for the split
+    if isinstance(x, tuple):
+        target_x = x[0]
+    else:
+        target_x = x
+    assert isinstance(target_x, paddle.Tensor)
+    assert len(target_x.shape) == 3
+    placements = target_x.placements
+    if placements is None:
+        placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
+    placements[index] = dist.Replicate()
+    target_x = dist.reshard(target_x, process_mesh, placements)
+    if need_transpose:
+        target_x = paddle.transpose(target_x, perm=[1, 0, 2])
+    if isinstance(x, tuple):
+        x = list(x)
+        x[0] = target_x
+        x = tuple(x)
+    else:
+        x = target_x
+
+    return x
+
+
+class SequenceParallelBegin(PlanBase):
+    """
+    With need_transpose=True, this plan will transpose and reshard the output from [b, s, h] to [s/mp, b, h].
+    With need_transpose=False, this plan will reshard the output from [s, b, h] to [s/mp, b, h].
+
+    This plan marks the beginning of the sp and should be added to the LAST layer before the sp range.
+    DON'T mark any layer in the sp range.
+    """
+
+    def __init__(self, need_transpose=True):
+        super().__init__()
+        self.need_transpose = need_transpose
+
+    def sequence_parallel_begin(self, process_mesh):
+        def begin(layer, input, output=None):
+            return sp_split(output, process_mesh, self.need_transpose)
+
+        return begin
+
+    def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
+        layer.register_forward_post_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+
+
+class SequenceParallelEnd(PlanBase):
+    """
+    With need_transpose=True, this plan will reshard and transpose the input from [s/mp, b, h] to [b, s, h].
+    With need_transpose=False, this plan will reshard the input from [s/mp, b, h] to [s, b, h].
+
+    This plan marks the ending of the sp and should be added to the FIRST layer after the sp range.
+    DON'T mark any layer in the sp range.
+    """
+
+    def __init__(self, need_transpose=True):
+        super().__init__()
+        self.need_transpose = need_transpose
+
+    def sequence_parallel_end(self, process_mesh):
+        def end(layer, input, output=None):
+            return sp_reduce_scatter(input, process_mesh, self.need_transpose)
+
+        return end
+
+    def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
+        layer.register_forward_pre_hook(
+            self.sequence_parallel_end(process_mesh)
+        )
+
+
+class SequenceParallel(PlanBase):
+    """
+    Do sequence parallel on the layer. Note the input should be in [b, s, h] format.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def sequence_parallel_begin(self, process_mesh):
+        def begin(layer, input, output=None):
+            return sp_split(input, process_mesh, True)
+
+        return begin
+
+    def sequence_parallel_end(self, process_mesh):
+        def end(layer, input, output=None):
+            return sp_reduce_scatter(output, process_mesh, True)
+
+        return end
+
+    def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
+        logging.warning(
+            "Sequence parallel with the usage of SequenceParallel may not reach the best throughput. "
+            "Try to use SequenceParallelBegin/End to achieve better performance"
+        )
+        layer.register_forward_pre_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+        layer.register_forward_post_hook(
+            self.sequence_parallel_end(process_mesh)
+        )
+
+
 class TensorParallel(ParallelModel):
     def __init__(self, model, parallelize_plan=None):
         super().__init__(model)
@@ -207,7 +340,10 @@ def match_layer(self, name):
                 key = key.replace(".bias", "")
                 shard_weight = False
             re_find = re.match(key, name)
-            if key == name or (re_find is not None and re_find.string == name):
+            if key == name or (
+                re_find is not None
+                and int(re_find.end()) - int(re_find.start()) == len(name)
+            ):
                 if isinstance(plan, PlanBase):
                     plan = [plan]
                 plans.append([plan, shard_weight, shard_bias])
@@ -217,15 +353,14 @@ def tensor_parallelizer_fn(self, model):
         if self.parallelize_plan is None:
             return
         for name, layer in model.named_sublayers():
-            if len(layer.sublayers()) == 0:
-                plans = self.match_layer(name)
-                if len(plans) > 0:
-                    for plan in plans:
-                        real_plan, shard_weight, shard_bias = plan
-                        for p in real_plan:
-                            p.apply(
-                                layer, self.get_mesh(), shard_weight, shard_bias
-                            )
+            plans = self.match_layer(name)
+            if len(plans) > 0:
+                for plan in plans:
+                    real_plan, shard_weight, shard_bias = plan
+                    for p in real_plan:
+                        p.apply(
+                            layer, self.get_mesh(), shard_weight, shard_bias
+                        )
         return model
 
 
diff --git a/test/auto_parallel/hybrid_strategy/CMakeLists.txt b/test/auto_parallel/hybrid_strategy/CMakeLists.txt
@@ -136,3 +136,11 @@ if((WITH_GPU) AND (LINUX))
   set_tests_properties(test_parallel_api_with_llama
                        PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
 endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_mp MODULES test_parallel_api_with_llama_mp
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_mp
+                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+endif()
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -35,9 +35,10 @@
 )
 from paddle.distributed.auto_parallel.intermediate.tensor_parallel import (
     ColWiseParallel,
-    PrepareLayerInput,
-    PrepareLayerOutput,
     RowWiseParallel,
+    SequenceParallel,
+    SequenceParallelBegin,
+    SequenceParallelEnd,
     tensor_parallel,
 )
 from paddle.io import BatchSampler, DataLoader, Dataset
@@ -147,6 +148,9 @@ def __init__(self):
         self.sequence_parallel = False
         if os.getenv("sequence_parallel") == "true":
             self.sequence_parallel = True
+        self.prepare_input_output = False
+        if os.getenv("prepare_input_output") == "true":
+            self.sequence_parallel = True
 
         seed = int(os.getenv("seed", 2024))
         np.random.seed(seed)
@@ -166,18 +170,6 @@ def init_dist_env(self):
         global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
         dist.auto_parallel.set_mesh(global_mesh)
 
-    def output_transpose(self, process_mesh):
-        def transpose(layer, input, output=None):
-            return paddle.transpose(output, perm=[1, 0, 2])
-
-        return transpose
-
-    def input_transpose(self, process_mesh):
-        def transpose(layer, input, output=None):
-            return paddle.transpose(input, perm=[1, 0, 2])
-
-        return transpose
-
     def check_mp(self, layer):
         if self.mp == 1:
             return
@@ -244,34 +236,46 @@ def parallel_model(self, layer, optimizer=None):
                     "llama.layers.*.mlp.down_proj": RowWiseParallel(),
                     "lm_head.weight": ColWiseParallel(),
                 }
-            if self.sequence_parallel:
-                plan = {
-                    "llama.embed_tokens": [
-                        ColWiseParallel(),
-                        PrepareLayerOutput(self.output_transpose),
-                    ],
-                    "llama.layers.*.self_attn.q_proj": [
-                        ColWiseParallel(),
-                        PrepareLayerOutput(self.output_transpose),
-                    ],
-                    "llama.layers.*.self_attn.k_proj": [
-                        ColWiseParallel(),
-                        PrepareLayerOutput(self.output_transpose),
-                    ],
-                    "llama.layers.*.self_attn.v_proj": [
-                        ColWiseParallel(),
-                        PrepareLayerOutput(self.output_transpose),
-                    ],
-                    "llama.layers.*.self_attn.o_proj": [
-                        RowWiseParallel(),
-                        PrepareLayerInput(self.input_transpose),
-                    ],
-                    "llama.layers.*.mlp.gate_proj": ColWiseParallel(),
-                    "llama.layers.*.mlp.up_proj": ColWiseParallel(),
-                    "llama.layers.*.mlp.down_proj": RowWiseParallel(),
-                    "lm_head.weight": ColWiseParallel(),
-                    "lm_head": PrepareLayerInput(self.input_transpose),
-                }
+            else:
+                if self.prepare_input_output:
+                    plan = {
+                        "llama.embed_tokens": ColWiseParallel(),
+                        "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.o_proj": RowWiseParallel(),
+                        "llama.layers.*.mlp.gate_proj": ColWiseParallel(),
+                        "llama.layers.*.mlp.up_proj": ColWiseParallel(),
+                        "llama.layers.*.mlp.down_proj": RowWiseParallel(),
+                        "lm_head.weight": ColWiseParallel(),
+                        "llama.layers.*.input_layernorm": SequenceParallel(),
+                        "llama.layers.*.post_attention_layernorm": SequenceParallel(),
+                        "llama.norm": SequenceParallel(),
+                    }
+                else:
+                    plan = {
+                        "llama.embed_tokens": [
+                            ColWiseParallel(),
+                            SequenceParallelBegin(),
+                        ],
+                        "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
+                        "llama.layers.*.self_attn.o_proj": RowWiseParallel(),
+                        "llama.layers.*.self_attn": [
+                            SequenceParallelEnd(),
+                            SequenceParallelBegin(),
+                        ],
+                        "llama.layers.*.mlp.gate_proj": ColWiseParallel(),
+                        "llama.layers.*.mlp.up_proj": ColWiseParallel(),
+                        "llama.layers.*.mlp.down_proj": RowWiseParallel(),
+                        "llama.layers.*.mlp": [
+                            SequenceParallelEnd(need_transpose=False),
+                            SequenceParallelBegin(need_transpose=False),
+                        ],
+                        "lm_head.weight": ColWiseParallel(),
+                        "lm_head": SequenceParallelEnd(),
+                    }
             layer, optimizer = tensor_parallel(layer, plan, optimizer)
         layer, optimizer = parallelize_model_and_optimizer(layer, optimizer)
         self.check_mp(layer)
diff --git a/test/auto_parallel/hybrid_strategy/single_llama_model.py b/test/auto_parallel/hybrid_strategy/single_llama_model.py
@@ -98,9 +98,9 @@ def __init__(self, config):
             self.intermediate_size, self.hidden_size, bias_attr=False
         )
 
-    def forward(self, x):
+    def forward(self, x, test_for_list_input_output):
         out = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-        return out
+        return out, test_for_list_input_output
 
 
 class LlamaRMSNorm(nn.Layer):
@@ -144,7 +144,7 @@ def forward(self, hidden_states):
 
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
+        hidden_states, _ = self.mlp(hidden_states, "ONLY_FOR_TEST")
         hidden_states = residual + hidden_states
 
         return hidden_states
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama.py
@@ -21,41 +21,6 @@
 os.environ['FLAGS_enable_pir_api'] = '0'
 
 
-class TestTensorParallelAPI(test_base.CommunicationTestDistBase):
-    def setUp(self):
-        super().setUp(num_of_devices=2, timeout=120, nnode=1)
-        self._default_envs = {
-            "dtype": "float32",
-            "seed": "2023",
-            "dp": "1",
-            "mp": "2",
-            "pp": "1",
-            "acc_step": "2",
-        }
-        self._changeable_envs = {
-            "backend": ["gpu"],
-            "amp": ["true"],
-            "amp_level": ["O2"],
-            "amp_dtype": ["bfloat16"],
-            "amp_master_grad": ["true"],
-            "use_lazy_init": ["true", "false"],
-            "sequence_parallel": ["true", "false"],
-        }
-
-    def test_simple_net_mp2(self):
-        envs_list = test_base.gen_product_envs_list(
-            self._default_envs, self._changeable_envs
-        )
-        for envs in envs_list:
-            ckpt_path = tempfile.TemporaryDirectory()
-            envs["ckpt_path"] = ckpt_path.name
-            self.run_test_case(
-                "parallel_api.py",
-                user_defined_envs=envs,
-            )
-            ckpt_path.cleanup()
-
-
 class TestShardingParallelAPI(test_base.CommunicationTestDistBase):
     def setUp(self):
         super().setUp(num_of_devices=2, timeout=120, nnode=1)
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_mp.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_mp.py
diff --git a/test/auto_parallel/hybrid_strategy/testslist.csv b/test/auto_parallel/hybrid_strategy/testslist.csv
