[Auto Parallel] Add two flags to mp plan. (#69776)

FeixLiu · web-flow · commit 109965ced94f · 2024-11-29T07:07:13.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -20,11 +20,61 @@
 from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
 
 
+def c_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 is_tensor(target_x)
+    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 c_concat(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 is_tensor(target_x)
+    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 PlanBase:
     def __init__(self):
         pass
 
-    def apply(self, param, process_mesh, shard_weight, shard_bias):
+    def apply(self, layer, process_mesh, shard_weight, shard_bias):
         raise NotImplementedError("Don't call the PlanBase directly.")
 
 
@@ -40,8 +90,16 @@ class ColWiseParallel(PlanBase):
     Note: `layer.bias` should have one dim.
     """
 
-    def __init__(self):
+    def __init__(self, gather_output=False):
         super().__init__()
+        self.gather_output = gather_output
+
+    def gather_output_hook(self, process_mesh):
+        def gather_hook(layer, input, output):
+            assert output is not None
+            return c_concat(output, process_mesh, False)
+
+        return gather_hook
 
     def apply(self, layer, process_mesh, shard_weight=True, shard_bias=True):
         """
@@ -79,6 +137,11 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=True):
             assert len(layer.bias.shape) == 1
             layer.bias = dist.shard_tensor(layer.bias, process_mesh, placement)
 
+        if self.gather_output:
+            layer.register_forward_post_hook(
+                self.gather_output_hook(process_mesh)
+            )
+
 
 class RowWiseParallel(PlanBase):
     """
@@ -90,8 +153,15 @@ class RowWiseParallel(PlanBase):
     Note: `layer.weight` should have two dims.
     """
 
-    def __init__(self):
+    def __init__(self, is_input_parallel=True):
         super().__init__()
+        self.is_input_parallel = is_input_parallel
+
+    def split_input_hook(self, process_mesh):
+        def split_hook(layer, input, output):
+            return c_split(input, process_mesh, False)
+
+        return split_hook
 
     def apply(self, layer, process_mesh, shard_weight=True, shard_bias=False):
         """
@@ -124,6 +194,8 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=False):
                 process_mesh,
                 placement,
             )
+        if not self.is_input_parallel:
+            layer.register_forward_pre_hook(self.split_input_hook(process_mesh))
 
 
 class PrepareLayerInput(PlanBase):
@@ -156,56 +228,6 @@ 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 is_tensor(target_x)
-    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 is_tensor(target_x)
-    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].
@@ -222,7 +244,7 @@ def __init__(self, need_transpose=True):
     def sequence_parallel_begin(self, process_mesh):
         def begin(layer, input, output):
             assert output is not None
-            return sp_split(output, process_mesh, self.need_transpose)
+            return c_split(output, process_mesh, self.need_transpose)
 
         return begin
 
@@ -248,7 +270,7 @@ def __init__(self, need_transpose=True):
     def sequence_parallel_end(self, process_mesh):
         def end(layer, input, output=None):
             assert input is not None
-            return sp_reduce_scatter(input, process_mesh, self.need_transpose)
+            return c_concat(input, process_mesh, self.need_transpose)
 
         return end
 
@@ -269,14 +291,14 @@ def __init__(self):
     def sequence_parallel_begin(self, process_mesh):
         def begin(layer, input, output=None):
             assert input is not None
-            return sp_split(input, process_mesh, True)
+            return c_split(input, process_mesh, True)
 
         return begin
 
     def sequence_parallel_end(self, process_mesh):
         def end(layer, input, output):
             assert output is not None
-            return sp_reduce_scatter(output, process_mesh, True)
+            return c_concat(output, process_mesh, True)
 
         return end
 
@@ -310,13 +332,13 @@ def __init__(self, need_transpose=True):
 
     def sequence_parallel_begin(self, process_mesh):
         def begin(layer, input, output=None):
-            return sp_split(output, process_mesh, self.need_transpose)
+            return c_split(output, process_mesh, self.need_transpose)
 
         return begin
 
     def sequence_parallel_end(self, process_mesh):
         def end(layer, input, output=None):
-            return sp_reduce_scatter(input, process_mesh, self.need_transpose)
+            return c_concat(input, process_mesh, self.need_transpose)
 
         return end
 
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -23,6 +23,7 @@
 import paddle.distributed as dist
 from paddle import LazyGuard
 from paddle.distributed.auto_parallel.intermediate.parallelize import (
+    parallelize,
     parallelize_model,
     parallelize_optimizer,
 )
@@ -55,12 +56,12 @@ def get_mesh(pp_idx=None):
 
 
 class Config:
-    vocab_size = 32000
-    hidden_size = 4096
-    intermediate_size = 11008
-    seq_length = 2048
+    vocab_size = 8192
+    hidden_size = 512
+    intermediate_size = 2048
+    seq_length = 512
     num_hidden_layers = 2
-    num_attention_heads = 32
+    num_attention_heads = 8
     rms_norm_eps = 1e-6
     use_lazy_init = False
 
@@ -146,6 +147,10 @@ def __init__(self):
         if num_hidden_layers:
             self.config.num_hidden_layers = int(num_hidden_layers)
 
+        self.one_api = False
+        if os.getenv("one_api") == "true":
+            self.one_api = True
+
         seed = int(os.getenv("seed", 2024))
         np.random.seed(seed)
         random.seed(seed)
@@ -219,12 +224,20 @@ def parallel_model(self, layer):
         if self.mp > 1:
             if not self.sequence_parallel:
                 plan = {
-                    "llama.embed_tokens": ColWiseParallel(),
+                    "llama.embed_tokens": ColWiseParallel(gather_output=True),
                     "llama.position_embedding": 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.*.self_attn.q_proj": ColWiseParallel(
+                        gather_output=True
+                    ),
+                    "llama.layers.*.self_attn.k_proj": ColWiseParallel(
+                        gather_output=True
+                    ),
+                    "llama.layers.*.self_attn.v_proj": ColWiseParallel(
+                        gather_output=True
+                    ),
+                    "llama.layers.*.self_attn.o_proj": RowWiseParallel(
+                        is_input_parallel=False
+                    ),
                     "llama.layers.*.mlp.gate_proj": ColWiseParallel(),
                     "llama.layers.*.mlp.up_proj": ColWiseParallel(),
                     "llama.layers.*.mlp.down_proj": RowWiseParallel(),
@@ -272,14 +285,36 @@ def parallel_model(self, layer):
                         "lm_head": SequenceParallelEnd(),
                     }
             mp_config = {'parallelize_plan': plan}
-        layer = parallelize_model(
-            layer,
-            dp_config=dp_config,
-            mp_config=mp_config,
-            pp_config=pp_config,
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
         )
+
+        if self.one_api:
+            optimizer = create_optimizer(layer, lr_scheduler)
+            model, optimizer = parallelize(
+                layer,
+                optimizer,
+                dp_config=dp_config,
+                mp_config=mp_config,
+                pp_config=pp_config,
+            )
+        else:
+            layer = parallelize_model(
+                layer,
+                dp_config=dp_config,
+                mp_config=mp_config,
+                pp_config=pp_config,
+            )
+            optimizer = create_optimizer(layer, lr_scheduler)
+            optimizer = parallelize_optimizer(
+                optimizer,
+                dp_config=dp_config,
+                mp_config=mp_config,
+                pp_config=pp_config,
+            )
         self.check_mp(layer)
-        return layer, dp_config, mp_config, pp_config
+        return layer, optimizer, lr_scheduler
 
     def run_llama(
         self, share_embedding=False, position_embedding=False, to_static=0
@@ -294,19 +329,7 @@ def run_llama(
                 self.config, share_embedding, position_embedding
             )
 
-        model, dp_config, mp_config, pp_config = self.parallel_model(model)
-
-        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
-            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
-        )
-        optimizer = create_optimizer(model, lr_scheduler)
-
-        optimizer = parallelize_optimizer(
-            optimizer,
-            dp_config=dp_config,
-            mp_config=mp_config,
-            pp_config=pp_config,
-        )
+        model, optimizer, lr_scheduler = self.parallel_model(model)
 
         criterion = LlamaPretrainingCriterion(self.config)
 
@@ -403,7 +426,7 @@ def run_llama(
                     tr_loss = 0
 
                 global_step += 1
-                if global_step // self.gradient_accumulation_steps >= 10:
+                if global_step // self.gradient_accumulation_steps >= 3:
                     break
         else:
             strategy = dist.Strategy()
@@ -433,7 +456,7 @@ def run_llama(
                 input_ids, labels = inputs
                 loss = dist_model(input_ids, labels)
                 logging.info(f"step: {step}  loss: {loss}")
-                if step >= 10:
+                if step >= 3:
                     break
 
     def run_test_cases(self, share_embedding=False, position_embedding=False):
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d.py
@@ -44,6 +44,7 @@ def setUp(self):
             "test_position_embedding": [
                 "1",
             ],
+            "one_api": ["true", "false"],
         }
 
     def test_simple_net_mp2_pp2(self):
@@ -86,6 +87,7 @@ def setUp(self):
             "test_position_embedding": [
                 "1",
             ],
+            "one_api": ["true", "false"],
         }
 
     def test_simple_net_dp2_pp2(self):
@@ -129,6 +131,7 @@ def setUp(self):
             "test_position_embedding": [
                 "1",
             ],
+            "one_api": ["true", "false"],
         }
 
     def test_simple_net_mp2_pp2(self):
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_3d.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_3d.py

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ def setUp(self):`
`44`	`44`	`"test_position_embedding": [`
`45`	`45`	`"1",`
`46`	`46`	`],`
	`47`	`+ "one_api": ["true", "false"],`
`47`	`48`	`}`
`48`	`49`
`49`	`50`	`def test_simple_net_mp2_pp2(self):`
`@@ -86,6 +87,7 @@ def setUp(self):`
`86`	`87`	`"test_position_embedding": [`
`87`	`88`	`"1",`
`88`	`89`	`],`
	`90`	`+ "one_api": ["true", "false"],`
`89`	`91`	`}`
`90`	`92`
`91`	`93`	`def test_simple_net_dp2_pp2(self):`
`@@ -129,6 +131,7 @@ def setUp(self):`
`129`	`131`	`"test_position_embedding": [`
`130`	`132`	`"1",`
`131`	`133`	`],`
	`134`	`+ "one_api": ["true", "false"],`
`132`	`135`	`}`
`133`	`136`
`134`	`137`	`def test_simple_net_mp2_pp2(self):`