PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
Copy file name to clipboard
+25-6Lines changed: 25 additions & 6 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
Copy file name to clipboard
+25-6Lines changed: 25 additions & 6 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/parallelize.py
Copy file name to clipboard
+45Lines changed: 45 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/parallelize.py
Copy file name to clipboard
+45Lines changed: 45 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py
Copy file name to clipboard
+8-6Lines changed: 8 additions & 6 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py
Copy file name to clipboard
+8-6Lines changed: 8 additions & 6 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
Copy file name to clipboard
+39-2Lines changed: 39 additions & 2 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
Copy file name to clipboard
+39-2Lines changed: 39 additions & 2 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+13-4Lines changed: 13 additions & 4 deletions b/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+13-4Lines changed: 13 additions & 4 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/parallel_api.py
Copy file name to clipboard
+25-29Lines changed: 25 additions & 29 deletions b/‎test/auto_parallel/hybrid_strategy/parallel_api.py
Copy file name to clipboard
+25-29Lines changed: 25 additions & 29 deletions
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 
 import paddle
 import paddle.distributed as dist
@@ -42,11 +43,24 @@ def __init__(self, optimizer, level=None):
 
         if isinstance(optimizer, ParallelOptimizer):
             self.optimizer = optimizer.optimizer
-            self.level = optimizer.level
+            if level is None:
+                self.level = optimizer.level
+            else:
+                if isinstance(level, int):
+                    level = str(level)
+                assert level in ("0", "1", "2", "3", None)
+                if optimizer.level is not None:
+                    assert (
+                        level == optimizer.level
+                    ), f"The level passed in is not identical with previous level. Current level is {level}, previous level is {optimizer.level}"
+                self.level = level
         else:
             assert isinstance(optimizer, Optimizer)
             self.optimizer = optimizer
-            assert level in ("os", "os_g", "p_g_os", None)
+            if isinstance(level, int):
+                level = str(level)
+            assert level in ("0", "1", "2", "3", None)
+            # level=0 and level=None are all mean pure dp
             self.level = level
 
         self.is_initialized = False
@@ -65,15 +79,15 @@ def parallelize(self, parallelized_parameters):
             self.optimizer._param_groups = self.optimizer._parameter_list
         # 2.wrap with shard_optimizer
         mesh = fleet.auto.get_mesh()
-        if self.level == "os":
+        if self.level == "1":
             self.optimizer = dist.shard_optimizer(
                 self.optimizer, dist.ShardingStage1(mesh)
             )
-        elif self.level == "os_g":
+        elif self.level == "2":
             self.optimizer = dist.shard_optimizer(
                 self.optimizer, dist.ShardingStage2(mesh)
             )
-        elif self.level == "p_g_os":
+        elif self.level == "3":
             self.optimizer = dist.shard_optimizer(
                 self.optimizer, dist.ShardingStage3(mesh)
             )
@@ -132,7 +146,12 @@ def parallelize_model(self):
 
 
 def parallelize_model_and_optimizer(model, optimizer=None):
-    assert isinstance(model, ParallelModel)
+    if not isinstance(model, ParallelModel):
+        assert not isinstance(optimizer, ParallelOptimizer)
+        logging.warning(
+            "The method `parallelize_model_and_optimizer` won't do anything since the model is not parallelized."
+        )
+        return model, optimizer
     parallelized_model = model.parallelize_model()
     parallelized_optimizer = None
     if optimizer is not None:
 
@@ -0,0 +1,45 @@
+#   Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .parallel_base import parallelize_model_and_optimizer
+from .pipeline_parallel import pipeline_parallel
+from .sharded_data_parallel import sharded_data_parallel
+from .tensor_parallel import tensor_parallel
+
+
+def parallelize(
+    model, optimizer, mesh=None, dp_config=None, mp_config=None, pp_config=None
+):
+    # TODO(yaliu): global mesh and split axis support
+    if pp_config is not None:
+        assert isinstance(pp_config, dict)
+        model, optimizer = pipeline_parallel(
+            model, optimizer, pp_config.get('split_spec')
+        )
+    if mp_config is not None:
+        assert isinstance(mp_config, dict)
+        model, optimizer = tensor_parallel(
+            model, optimizer, mp_config.get('parallelize_plan')
+        )
+    if dp_config is not None:
+        assert isinstance(dp_config, dict)
+        model, optimizer = sharded_data_parallel(
+            model,
+            optimizer,
+            level=dp_config.get('sharding_level'),
+            offload=bool(dp_config.get('offload')),
+            exclude_layer=dp_config.get('exclude_layer'),
+        )
+    model, optimizer = parallelize_model_and_optimizer(model, optimizer)
+    return model, optimizer
@@ -26,9 +26,10 @@ class ShardedDataParallel(ParallelModel):
         model (paddle.nn.Layer): A single card model to be distributed.
         optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed.
         level (str): Zero stage, can be the following values:
-            os: Zero Stage1
-            os_g: Zero Stage2
-            p_g_os: Zero Stage3
+            0: no sharding (pure dp)
+            1: Zero Stage1
+            2: Zero Stage2
+            3: Zero Stage3
             Default: None, which means optimizer is replicated among all process.
         offload (bool): whether enable cpu offload strategy, not implemented currently.
         exclude_layer (list): Specify which layers do not use the zero stage strategy, not implemented currently.
@@ -60,9 +61,10 @@ def sharded_data_parallel(
         model (paddle.nn.Layer): A single card model to be distributed
         optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed
         level (str): Zero stage, can be the following values:
-            os: Zero Stage1
-            os_g: Zero Stage2
-            p_g_os: Zero Stage3
+            0: no sharding (pure dp)
+            1: Zero Stage1
+            2: Zero Stage2
+            3: Zero Stage3
             Default: None, which means optimizer is replicated among all process.
         offload (bool): whether enable cpu offload strategy, not implemented currently.
         exclude_layer (list): Specify which layers do not use the zero stage strategy, not implemented currently.
 
@@ -258,7 +258,7 @@ def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
         )
 
 
-class SequenceParallel(PlanBase):
+class SequenceParallelEnable(PlanBase):
     """
     Do sequence parallel on the layer. Note the input should be in [b, s, h] format.
     """
@@ -293,6 +293,43 @@ def apply(self, layer, process_mesh, shard_weight=None, shard_bias=None):
         )
 
 
+class SequenceParallelDisable(PlanBase):
+    """
+    Disable sequence parallel on the layer.
+    If the need_transpose is true:
+        - change the input from  [s/mp, b, h] to [b, s, h]
+        - change the output from [b, s, h] to [s/mp, b, h]
+    If the need_transpose is False:
+        - change the input from  [s/mp, b, h] to [s, b, h]
+        - change the output from [s, b, h] to [s/mp, b, h]
+    """
+
+    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 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)
+        )
+
+        layer.register_forward_post_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+
+
 class TensorParallel(ParallelModel):
     def __init__(self, model, parallelize_plan=None):
         super().__init__(model)
@@ -358,7 +395,7 @@ def tensor_parallelizer_fn(self, model):
         return model
 
 
-def tensor_parallel(model, parallelize_plan=None, optimizer=None):
+def tensor_parallel(model, optimizer=None, parallelize_plan=None):
     """
     Tensor parallel.
     :param model: paddle.nn.Layer, the model to be shard into tensor parallel.
 
@@ -131,16 +131,25 @@ if((WITH_GPU) AND (LINUX))
 endif()
 if((WITH_GPU) AND (LINUX))
   py_test_modules(
-    test_parallel_api_with_llama MODULES test_parallel_api_with_llama ENVS
+    test_parallel_api_with_llama_1d MODULES test_parallel_api_with_llama_1d
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_1d
+                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_2d MODULES test_parallel_api_with_llama_2d
+    ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
-  set_tests_properties(test_parallel_api_with_llama
+  set_tests_properties(test_parallel_api_with_llama_2d
                        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
+    test_parallel_api_with_llama_3d MODULES test_parallel_api_with_llama_3d
     ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
-  set_tests_properties(test_parallel_api_with_llama_mp
+  set_tests_properties(test_parallel_api_with_llama_3d
                        PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
 endif()
@@ -23,23 +23,19 @@
 import paddle
 import paddle.distributed as dist
 from paddle import LazyGuard
-from paddle.distributed.auto_parallel.intermediate.parallel_base import (
-    parallelize_model_and_optimizer,
+from paddle.distributed.auto_parallel.intermediate.parallelize import (
+    parallelize,
 )
 from paddle.distributed.auto_parallel.intermediate.pipeline_parallel import (
     SplitPoint,
-    pipeline_parallel,
-)
-from paddle.distributed.auto_parallel.intermediate.sharded_data_parallel import (
-    sharded_data_parallel,
 )
 from paddle.distributed.auto_parallel.intermediate.tensor_parallel import (
     ColWiseParallel,
     RowWiseParallel,
-    SequenceParallel,
     SequenceParallelBegin,
+    SequenceParallelDisable,
+    SequenceParallelEnable,
     SequenceParallelEnd,
-    tensor_parallel,
 )
 from paddle.io import BatchSampler, DataLoader, Dataset
 
@@ -141,10 +137,7 @@ def __init__(self):
             self.amp_level = os.getenv("amp_level")
         if os.getenv("amp_master_grad") == "true":
             self.amp_master_grad = True
-        self.sharding_stage_to_level = [None, "os", "os_g", "p_g_os"]
-        self.level = self.sharding_stage_to_level[
-            int(os.getenv("sharding_stage", 0))
-        ]
+        self.level = os.getenv("sharding_stage", "0")
         self.sequence_parallel = False
         if os.getenv("sequence_parallel") == "true":
             self.sequence_parallel = True
@@ -207,6 +200,9 @@ def check_mp(self, layer):
                     ]
 
     def parallel_model(self, layer, optimizer=None):
+        dp_config = None
+        mp_config = None
+        pp_config = None
         if self.pp > 1:
             decoders_per_rank = self.config.num_hidden_layers // self.pp
             split_spec = OrderedDict(
@@ -218,11 +214,9 @@ def parallel_model(self, layer, optimizer=None):
                     for i in range(1, self.pp)
                 ]
             )
-            layer, optimizer = pipeline_parallel(layer, optimizer, split_spec)
+            pp_config = {'split_spec': split_spec}
         if self.dp > 1:
-            layer, optimizer = sharded_data_parallel(
-                layer, optimizer, self.level
-            )
+            dp_config = {'sharding_level': self.level}
         if self.mp > 1:
             if not self.sequence_parallel:
                 plan = {
@@ -248,9 +242,9 @@ def parallel_model(self, layer, optimizer=None):
                         "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(),
+                        "llama.layers.*.input_layernorm": SequenceParallelEnable(),
+                        "llama.layers.*.post_attention_layernorm": SequenceParallelEnable(),
+                        "llama.norm": SequenceParallelEnable(),
                     }
                 else:
                     plan = {
@@ -262,22 +256,24 @@ def parallel_model(self, layer, optimizer=None):
                         "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.*.self_attn": SequenceParallelDisable(),
                         "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),
-                        ],
+                        "llama.layers.*.mlp": SequenceParallelDisable(
+                            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)
+            mp_config = {'parallelize_plan': plan}
+        layer, optimizer = parallelize(
+            layer,
+            optimizer,
+            dp_config=dp_config,
+            mp_config=mp_config,
+            pp_config=pp_config,
+        )
         self.check_mp(layer)
         if optimizer is None:
             return layer