PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
Copy file name to clipboard
+108-23Lines changed: 108 additions & 23 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
Copy file name to clipboard
+108-23Lines changed: 108 additions & 23 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
Copy file name to clipboard
+60-15Lines changed: 60 additions & 15 deletions b/‎python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
Copy file name to clipboard
+60-15Lines changed: 60 additions & 15 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/parallel_api.py
Copy file name to clipboard
+13-9Lines changed: 13 additions & 9 deletions b/‎test/auto_parallel/hybrid_strategy/parallel_api.py
Copy file name to clipboard
+13-9Lines changed: 13 additions & 9 deletions
@@ -118,7 +118,8 @@ def __init__(self, model):
         self.tp_parallelizer = None
         self.sharding_parallelizer = None
         self.model = None
-
+        self.share_param_list = {}
+        self.layer_param_placements = {}
         if isinstance(model, ParallelModel):
             self.pp_parallelizer = model.pp_parallelizer
             self.tp_parallelizer = model.tp_parallelizer
@@ -147,8 +148,9 @@ def parallelize_model(self):
 
         if self.tp_parallelizer is not None:
             assert callable(self.tp_parallelizer)
-            self.model = self.tp_parallelizer(self.model)
-
+            self.model, self.layer_param_placements = self.tp_parallelizer(
+                self.model
+            )
         if self.sharding_parallelizer is not None:
             assert callable(self.sharding_parallelizer)
             self.model = self.sharding_parallelizer(self.model)
@@ -157,36 +159,119 @@ def parallelize_model(self):
 
         return self.model
 
+    def _process_share_weight_layer(
+        self, layer, origin_weight, param_name, param_placements
+    ):
+        ipp = (
+            layer.pipeline_stage_index
+            if hasattr(layer, "pipeline_stage_index")
+            else 0
+        )
+
+        def create_pre_hook(origin_weight, param_name):
+            def forward_pre_hook(layer, input):
+                setattr(
+                    layer,
+                    param_name,
+                    None,
+                )
+                delattr(layer, param_name)
+                mesh = self.get_mesh(ipp)
+                share_weight = dist.reshard(
+                    origin_weight,
+                    mesh,
+                    param_placements,
+                )
+                setattr(
+                    layer,
+                    param_name,
+                    share_weight,
+                )
+
+            return forward_pre_hook
+
+        def create_post_hook(origin_weight, param_name):
+            def forward_post_hook(layer, input, output):
+                setattr(
+                    layer,
+                    param_name,
+                    origin_weight,
+                )
+
+            return forward_post_hook
+
+        layer.register_forward_pre_hook(
+            create_pre_hook(origin_weight, param_name)
+        )
+        layer.register_forward_post_hook(
+            create_post_hook(origin_weight, param_name)
+        )
+
     def _shard_all_param(self, model):
         param_name_to_shard_param = {}
+        param_name_to_pp_stage = {}
 
         def shard_layer_param(layer):
             if self.pp_parallelizer is not None:
                 assert hasattr(layer, "pipeline_stage_index")
             for param_name in list(layer._parameters.keys()):
                 param = getattr(layer, param_name)
-                if param is not None and not param.is_dist():
+                if param is not None:
                     param_full_name = param.name
-                    if param_full_name in param_name_to_shard_param:
-                        setattr(
-                            layer,
-                            param_name,
-                            param_name_to_shard_param[param_full_name],
-                        )
+                    ipp = (
+                        layer.pipeline_stage_index
+                        if hasattr(layer, "pipeline_stage_index")
+                        else 0
+                    )
+                    mesh = self.get_mesh(ipp)
+                    param_placements = [
+                        dist.Replicate() for _ in range(len(mesh._shape))
+                    ]
+                    if layer in self.layer_param_placements:
+                        if param_name in self.layer_param_placements[layer]:
+                            param_placements = (
+                                self.layer_param_placements[layer][param_name]
+                                if self.layer_param_placements[layer][
+                                    param_name
+                                ]
+                                is not None
+                                else param_placements
+                            )
+                    if not param.is_dist():
+                        if param_full_name in param_name_to_shard_param:
+                            setattr(
+                                layer,
+                                param_name,
+                                param_name_to_shard_param[param_full_name],
+                            )
+                            if ipp != param_name_to_pp_stage[param_full_name]:
+                                self._process_share_weight_layer(
+                                    layer,
+                                    param_name_to_shard_param[param_full_name],
+                                    param_name,
+                                    param_placements,
+                                )
+                        else:
+                            param = dist.shard_tensor(
+                                param, mesh, param_placements
+                            )
+                            param_name_to_shard_param[param_full_name] = param
+                            param_name_to_pp_stage[param_full_name] = ipp
+                            setattr(layer, param_name, param)
                     else:
-                        ipp = (
-                            layer.pipeline_stage_index
-                            if hasattr(layer, "pipeline_stage_index")
-                            else 0
-                        )
-                        mesh = self.get_mesh(ipp)
-                        param = dist.shard_tensor(
-                            param,
-                            mesh,
-                            [dist.Replicate() for _ in range(len(mesh._shape))],
-                        )
-                        param_name_to_shard_param[param_full_name] = param
-                        setattr(layer, param_name, param)
+                        if (
+                            param_full_name in param_name_to_shard_param
+                            and ipp != param_name_to_pp_stage[param_full_name]
+                        ):
+                            self._process_share_weight_layer(
+                                layer,
+                                param_name_to_shard_param[param_full_name],
+                                param_name,
+                                param_placements,
+                            )
+                        elif param_full_name not in param_name_to_shard_param:
+                            param_name_to_shard_param[param_full_name] = param
+                            param_name_to_pp_stage[param_full_name] = ipp
 
         for name, layer in model.named_sublayers():
             shard_layer_param(layer)
 
@@ -82,7 +82,7 @@ def c_concat(x, process_mesh, need_transpose):
 
 class PlanBase:
     def __init__(self):
-        pass
+        self.share_param_list = {}
 
     def apply(self, layer, process_mesh, shard_weight, shard_bias):
         raise NotImplementedError("Don't call the PlanBase directly.")
@@ -143,6 +143,7 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=True):
         index = process_mesh.dim_names.index('mp')  # get the axis for the split
         size = len(process_mesh.shape)
         placement = [dist.Replicate() for _ in range(size)]
+        param_placements = {}
         assert isinstance(layer, paddle.nn.Layer)
         if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
             logging.warning(
@@ -157,20 +158,39 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=True):
         ):
             placement[index] = dist.Shard(1)
             assert len(layer.weight.shape) == 2
-            layer.weight = dist.shard_tensor(
-                layer.weight,
-                process_mesh,
-                placement,
-            )
+            # NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
+            if (
+                self.share_param_list is not None
+                and layer.weight.name in self.share_param_list
+                and self.share_param_list[layer.weight.name] > 1
+            ):
+                param_placements.update({"weight": placement})
+            else:
+                layer.weight = dist.shard_tensor(
+                    layer.weight,
+                    process_mesh,
+                    placement,
+                )
         if hasattr(layer, "bias") and layer.bias is not None and shard_bias:
             placement[index] = dist.Shard(0)
             assert len(layer.bias.shape) == 1
-            layer.bias = dist.shard_tensor(layer.bias, process_mesh, placement)
+            # NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
+            if (
+                self.share_param_list is not None
+                and layer.bias.name in self.share_param_list
+                and self.share_param_list[layer.bias.name] > 1
+            ):
+                param_placements.update({"bias": placement})
+            else:
+                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)
             )
+        return param_placements
 
 
 class RowWiseParallel(PlanBase):
@@ -185,7 +205,7 @@ class RowWiseParallel(PlanBase):
 
     Args:
         is_input_parallel (bool): Whether the input is a local tensor or a global tensor. If the input is a
-            global tensor, an extra split will be called. The default value is `True`，
+            global tensor, an extra split will be called. The default value is `True`,
             which means the input is a local tensor.
 
     Examples:
@@ -225,6 +245,7 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=False):
         size = len(process_mesh.shape)
         placement = [dist.Replicate() for _ in range(size)]
         placement[index] = dist.Shard(0)
+        param_placements = {}
         assert isinstance(layer, paddle.nn.Layer)
         if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
             logging.warning(
@@ -238,13 +259,22 @@ def apply(self, layer, process_mesh, shard_weight=True, shard_bias=False):
             and shard_weight
         ):
             assert len(layer.weight.shape) == 2
-            layer.weight = dist.shard_tensor(
-                layer.weight,
-                process_mesh,
-                placement,
-            )
+            # NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
+            if (
+                self.share_param_list is not None
+                and layer.weight.name in self.share_param_list
+                and self.share_param_list[layer.weight.name] > 1
+            ):
+                param_placements.update({"weight": placement})
+            else:
+                layer.weight = dist.shard_tensor(
+                    layer.weight,
+                    process_mesh,
+                    placement,
+                )
         if not self.is_input_parallel:
             layer.register_forward_pre_hook(self.split_input_hook(process_mesh))
+        return param_placements
 
 
 class PrepareLayerInput(PlanBase):
@@ -626,20 +656,35 @@ def match_layer(self, name):
     def tensor_parallelizer_fn(self, model):
         if self.parallelize_plan is None:
             return
+        layer_param_placements = {}
+        share_param_list = {}
+        for name, layer in model.named_sublayers():
+            for param_name in list(layer._parameters.keys()):
+                param = getattr(layer, param_name)
+                if param.name not in share_param_list:
+                    share_param_list[param.name] = 1
+                    continue
+                share_param_list[param.name] += 1
         for name, layer in model.named_sublayers():
             plans = self.match_layer(name)
+            layer_param_placements[layer] = {}
             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(
+                        p.share_param_list = share_param_list
+                        param_placements = p.apply(
                             layer,
                             self.get_mesh(pp_idx),
                             shard_weight,
                             shard_bias,
                         )
-        return model
+                        if param_placements is not None and param_placements:
+                            layer_param_placements[layer].update(
+                                param_placements
+                            )
+        return model, layer_param_placements
 
 
 def tensor_parallel(model, optimizer=None, config=None):
 
@@ -135,23 +135,23 @@ if((WITH_GPU) AND (LINUX))
     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")
+                       PROPERTIES TIMEOUT "400" 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_2d
-                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
 endif()
 if((WITH_GPU) AND (LINUX))
   py_test_modules(
     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_3d
-                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
 endif()
 if((WITH_GPU) AND (LINUX))
   py_test_modules(
 
@@ -160,7 +160,7 @@ def init_dist_env(self):
         global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
         dist.auto_parallel.set_mesh(global_mesh)
 
-    def check_mp(self, layer):
+    def check_mp(self, layer, share_embedding):
         if self.mp == 1:
             return
         for name, sub_layer in layer.named_sublayers():
@@ -174,12 +174,14 @@ def check_mp(self, layer):
                         dist.Replicate(),
                         dist.Shard(0),
                     ]
+                if 'gate_proj' in name or 'up_proj' in name:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(1),
+                    ]
                 if (
-                    'gate_proj' in name
-                    or 'up_proj' in name
-                    or 'embed_tokens' in name
-                    or 'lm_head' in name
-                ):
+                    'embed_tokens' in name or 'lm_head' in name
+                ) and not share_embedding:
                     assert sub_layer.weight.placements == [
                         dist.Replicate(),
                         dist.Shard(1),
@@ -196,7 +198,7 @@ def check_mp(self, layer):
                         dist.Shard(0),
                     ]
 
-    def parallel_model(self, layer):
+    def parallel_model(self, layer, share_embedding=False):
         dp_config = None
         mp_config = None
         pp_config = None
@@ -306,7 +308,7 @@ def parallel_model(self, layer):
                 optimizer,
                 config=config,
             )
-        self.check_mp(layer)
+        self.check_mp(layer, share_embedding)
         return layer, optimizer, lr_scheduler
 
     def run_llama(
@@ -322,7 +324,9 @@ def run_llama(
                 self.config, share_embedding, position_embedding
             )
 
-        model, optimizer, lr_scheduler = self.parallel_model(model)
+        model, optimizer, lr_scheduler = self.parallel_model(
+            model, share_embedding
+        )
 
         criterion = LlamaPretrainingCriterion(self.config)