Revert "[Auto-Parallel] optimize llama-7b benchmark in temporary solution to …" (#10682)

This reverts commit e0921f0.

Author: liym27

paddlenlp/transformers/llama/modeling_auto.py

@@ -191,7 +191,6 @@ def scaled_dot_product_attention(
 
 colwise_placements = [dist.Replicate(), dist.Shard(1)]
 rowise_placement = [dist.Replicate(), dist.Shard(0)]
-replicate_placements = [dist.Replicate(), dist.Replicate()]
 
 
 class LlamaRMSNormAuto(nn.Layer):
@@ -242,28 +241,28 @@ def __init__(self, config, ipp: Optional[int] = None):
             self.gate_up_fused_proj.weight = dist.shard_tensor(
                 self.gate_up_fused_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
         else:
             self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias_attr=False)
             self.gate_proj.weight = dist.shard_tensor(
                 self.gate_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
             self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias_attr=False)
             self.up_proj.weight = dist.shard_tensor(
                 self.up_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias_attr=False)
         self.down_proj.weight = dist.shard_tensor(
             self.down_proj.weight,
             get_mesh(self.ipp),
-            rowise_placement if self.config.tensor_parallel_degree > 1 else replicate_placements,
+            rowise_placement,
         )
 
     def forward(self, x):
@@ -323,7 +322,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
             self.qkv_proj.weight = dist.shard_tensor(
                 self.qkv_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
         else:
@@ -335,7 +334,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
             self.q_proj.weight = dist.shard_tensor(
                 self.q_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
             self.k_proj = nn.Linear(
@@ -346,7 +345,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
             self.k_proj.weight = dist.shard_tensor(
                 self.k_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
             self.v_proj = nn.Linear(
@@ -357,7 +356,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
             self.v_proj.weight = dist.shard_tensor(
                 self.v_proj.weight,
                 get_mesh(self.ipp),
-                colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+                colwise_placements,
             )
 
         self.o_proj = nn.Linear(
@@ -368,7 +367,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
         self.o_proj.weight = dist.shard_tensor(
             self.o_proj.weight,
             get_mesh(self.ipp),
-            rowise_placement if self.config.tensor_parallel_degree > 1 else replicate_placements,
+            rowise_placement,
         )
 
         if config.rope:
@@ -1220,7 +1219,7 @@ def __init__(self, config: LlamaConfig):
         self.weight = dist.shard_tensor(
             self.weight,
             get_mesh(-1),
-            colwise_placements if self.config.tensor_parallel_degree > 1 else replicate_placements,
+            colwise_placements,
         )
 
     def forward(self, hidden_states, tensor_parallel_output=None):