[Auto Parallel] Support 3D parallel with intermediate API (#69365)

deepllz · web-flow · commit 1eed387a6fbb · 2024-11-14T09:35:10.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 from collections import OrderedDict
 from enum import Enum
 
@@ -85,11 +86,38 @@ def forward_pre_hook(layer, input):
             # TODO(deepllz): support in the future
             return input
 
+        # step1: set every layer's own pipeline_stage_index
         split_layer_names = list(self.split_spec.keys())
-        for index, name in enumerate(split_layer_names):
+        sublayer_names = [name for name, _ in model.named_sublayers()]
+        # Mark which layer is the next pipeline stage
+        pipline_layer_mark = [0 for _ in range(len(sublayer_names))]
+        for split_layer_name in split_layer_names:
+            split_point = self.split_spec[split_layer_name]
+            index = sublayer_names.index(split_layer_name)
+            if split_point == SplitPoint.END:
+                is_valid = False
+                for i in range(index + 1, len(sublayer_names)):
+                    if not sublayer_names[i].startswith(split_layer_name):
+                        pipline_layer_mark[i] = 1
+                        is_valid = True
+                        break
+                assert (
+                    is_valid
+                ), f"the last layer:{split_layer_name} must not be SplitPoint.END, please check the split_spec"
+            else:
+                raise NotImplementedError(
+                    "SplitPoint.BEGINNING is not supported currently"
+                )
+                pipline_layer_mark[index] = 1
+        # the inclusiveSum of pipline_layer_mark is the pipeline stage index
+        pipline_stage_index = list(itertools.accumulate(pipline_layer_mark))
+        for index, (name, layer) in enumerate(model.named_sublayers()):
+            layer.pipeline_stage_index = pipline_stage_index[index]
+
+        # step2: insert reshard
+        for name in split_layer_names:
             layer = get_layer_by_name(name)
             split_point = self.split_spec[name]
-            layer.pipeline_stage_index = index
             layer.split_point = split_point
             if split_point == SplitPoint.END:
                 layer.register_forward_post_hook(forward_post_hook)
@@ -98,6 +126,7 @@ def forward_pre_hook(layer, input):
                     "SplitPoint.BEGINNING is not supported currently"
                 )
                 layer.register_forward_pre_hook(forward_pre_hook)
+
         return model
 
 
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -220,7 +220,8 @@ def __init__(self, need_transpose=True):
         self.need_transpose = need_transpose
 
     def sequence_parallel_begin(self, process_mesh):
-        def begin(layer, input, output=None):
+        def begin(layer, input, output):
+            assert output is not None
             return sp_split(output, process_mesh, self.need_transpose)
 
         return begin
@@ -246,6 +247,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 end
@@ -266,12 +268,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 begin
 
     def sequence_parallel_end(self, process_mesh):
-        def end(layer, input, output=None):
+        def end(layer, input, output):
+            assert output is not None
             return sp_reduce_scatter(output, process_mesh, True)
 
         return end
@@ -309,20 +313,6 @@ def __init__(self, model, parallelize_plan=None):
             self.parallelize_plan = parallelize_plan
             self.tp_parallelizer = self.tensor_parallelizer_fn
 
-    def get_mesh(self):
-        # TODO(yaliu): fit pp
-        # Get local mesh for current pp.
-        assert "mp" in self.global_mesh.dim_names
-        if "pp" in self.global_mesh.dim_names:
-            assert (
-                self.global_mesh.get_dim_size("pp") == 1
-            ), "Not support pp with mp for now."
-            mesh = self.global_mesh.get_mesh_with_dim("pp")[0]
-        else:
-            mesh = self.global_mesh
-        assert len(mesh.shape) in [1, 2]
-        return mesh
-
     def match_layer(self, name):
         # Match the layer to a plan.
         # Will return the plan if the layer hits one, otherwise return None.
@@ -355,11 +345,15 @@ def tensor_parallelizer_fn(self, model):
         for name, layer in model.named_sublayers():
             plans = self.match_layer(name)
             if len(plans) > 0:
+                pp_idx = getattr(layer, "pipeline_stage_index", 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
+                            layer,
+                            self.get_mesh(pp_idx),
+                            shard_weight,
+                            shard_bias,
                         )
         return model
 
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
@@ -64,5 +64,77 @@ def test_simple_net_mp2(self):
             ckpt_path.cleanup()
 
 
+class TestMPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=120, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["true"],
+            "prepare_input_output": ["false"],
+        }
+
+    def test_simple_net_mp2_pp2(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 TestDPMPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=120, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["true"],
+            "prepare_input_output": ["false"],
+        }
+
+    def test_simple_net_dp2_mp2_pp2(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()
+
+
 if __name__ == "__main__":
     unittest.main()
