Support fused_attention_qkv for auto_parallel llama (#8432)

* add

* add

* add

* add

* add

Author: zhangbo9674

llm/llama/auto_parallel/run_pretrain_auto.py

@@ -86,6 +86,10 @@ class PreTrainingArguments(TrainingArguments):
             "help": "Enable fused_linear_param_grad pass, which should replace add_n_op with add_op for gradients accumulation."
         },
     )
+    fuse_allreduce_split_to_reducescatter: bool = field(
+        default=False,
+        metadata={"help": "Enable fuse_allreduce_split_to_reducescatter pass."},
+    )
     eliminate_transpose: bool = field(
         default=False,
         metadata={
@@ -138,6 +142,11 @@ def __post_init__(self):
             fused_passes.enable = True
             fused_passes.fused_passes_list.append("fused_linear_param_grad_add_pass")
 
+        if self.fuse_allreduce_split_to_reducescatter:
+            fused_passes = self.strategy.fused_passes
+            fused_passes.enable = True
+            fused_passes.fused_passes_list.append("fuse_allreduce_split_to_reducescatter_pass")
+
         if self.eliminate_transpose:
             fused_passes = self.strategy.fused_passes
             fused_passes.enable = True

paddlenlp/transformers/llama/modeling_auto.py

@@ -272,16 +272,14 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
         self.head_dim = self.hidden_size // config.num_attention_heads
 
         self.num_key_value_heads = config.num_key_value_heads
+        assert config.num_attention_heads // config.num_key_value_heads
         self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.gqa_or_mqa = config.num_attention_heads != config.num_key_value_heads
 
         self.max_position_embeddings = config.max_position_embeddings
         self.seq_length = config.seq_length
 
         self.fuse_attention_qkv = config.fuse_attention_qkv
-        if self.fuse_attention_qkv and config.num_attention_heads != config.num_key_value_heads:
-            raise ValueError(
-                f"fuse_attention_qkv can't be True when num_attention_heads {config.num_attention_heads}!= num_key_value_heads {config.num_key_value_heads}"
-            )
 
         self.kv_indices = None
         # Note that we will actually perform a recompute only if both enable_recompute and layerwise_recompute are set to True
@@ -303,7 +301,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False, ipp:
         if self.fuse_attention_qkv:
             self.qkv_proj = nn.Linear(
                 self.hidden_size,
-                3 * self.hidden_size,
+                self.hidden_size + 2 * self.config.num_key_value_heads * self.head_dim,
                 bias_attr=False,
             )
             self.qkv_proj.weight = dist.shard_tensor(
@@ -415,10 +413,16 @@ def forward(
             )
 
         if self.fuse_attention_qkv:
-            target_shape = [0, 0, self.num_heads, 3 * self.head_dim]
+            target_shape = [0, 0, self.num_key_value_heads, (self.num_key_value_groups + 2) * self.head_dim]
             mix_layer = self.qkv_proj(hidden_states)
             mix_layer = paddle.reshape_(mix_layer, target_shape)
-            query_states, key_states, value_states = paddle.split(mix_layer, num_or_sections=3, axis=-1)
+            query_states, key_states, value_states = paddle.split(
+                mix_layer,
+                num_or_sections=[self.num_key_value_groups * self.head_dim, self.head_dim, self.head_dim],
+                axis=-1,
+            )
+            if self.gqa_or_mqa:
+                query_states = paddle.reshape(query_states, [0, 0, self.num_heads, self.head_dim])
         else:
             target_query_shape = [0, 0, self.num_heads, self.head_dim]
             target_key_value_shape = [0, 0, self.num_key_value_heads, self.head_dim]