Merge pull request #801 from zhaoyinglia/AutoParallel/add_fuse_qkv

[AutoParallel] add fuse_attn_qkv for gpt

ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_1.3B_dp8.yaml

@@ -18,6 +18,7 @@ Model:
   type_vocab_size: 16
   initializer_range: 0.02
   use_recompute: True
+  fuse_attn_qkv: True
 
 
 Distributed:

ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_1.3B_single_card.yaml

@@ -18,6 +18,7 @@ Model:
   type_vocab_size: 16
   initializer_range: 0.02
   use_recompute: True
+  fuse_attn_qkv: True
 
 
 Distributed:

ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_345M_single_card.yaml

@@ -18,6 +18,7 @@ Model:
   type_vocab_size: 16
   initializer_range: 0.02
   use_recompute: False
+  fuse_attn_qkv: True
 
 
 Distributed:

ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_6.7B_sharding16.yaml

@@ -18,6 +18,7 @@ Model:
   type_vocab_size: 16
   initializer_range: 0.02
   use_recompute: True
+  fuse_attn_qkv: True
 
 
 Distributed:

ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_base.yaml

@@ -27,6 +27,7 @@ Engine:
 Model:
   module: "GPTModuleAuto"
   name: "GPT"
+  fuse_attn_qkv: False
 
 
 Data:

ppfleetx/models/language_model/gpt/auto/auto_model.py

@@ -47,7 +47,7 @@ def __init__(self,
                  need_weights=False,
                  weight_attr=None,
                  bias_attr=None,
-                 fuse=False,
+                 fuse_attn_qkv=False,
                  mesh=None,
                  mesh_idx=None):
         super(MultiHeadAttention, self).__init__()
@@ -57,14 +57,14 @@ def __init__(self,
         self.num_heads = num_heads
         self.dropout = dropout
         self.need_weights = need_weights
-        self.fuse = fuse
+        self.fuse_attn_qkv = fuse_attn_qkv
         self.mesh = mesh
         self.mesh_idx = mesh_idx
 
         self.head_dim = embed_dim // num_heads
         assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
 
-        if self.fuse:
+        if self.fuse_attn_qkv:
             assert self.kdim == embed_dim
             assert self.vdim == embed_dim
             self.qkv_proj = nn.Linear(
@@ -80,6 +80,9 @@ def __init__(self,
             embed_dim, embed_dim, weight_attr, bias_attr=bias_attr)
 
     def _fuse_prepare_qkv(self, query, use_cache=False, cache=None):
+        auto.shard_tensor(self.qkv_proj.weight, self.mesh[self.mesh_idx],
+                          [None, self.mesh.mp])
+
         mix_layer = self.qkv_proj(query)
         mix_layer = paddle.reshape_(mix_layer,
                                     [0, 0, self.num_heads, 3 * self.head_dim])
@@ -220,13 +223,13 @@ def forward(self,
         value = query if value is None else value
         # compute q ,k ,v
         if use_cache is False:
-            if self.fuse:
+            if self.fuse_attn_qkv:
                 q, k, v = self._fuse_prepare_qkv(query, use_cache, cache)
             else:
                 q, k, v = self._prepare_qkv(query, key, value, use_cache,
                                             cache)
         else:
-            if self.fuse:
+            if self.fuse_attn_qkv:
                 q, k, v, cache = self._fuse_prepare_qkv(query, use_cache,
                                                         cache)
             else:
@@ -342,6 +345,7 @@ def __init__(self,
                  normalize_before=True,
                  weight_attr=None,
                  bias_attr=None,
+                 fuse_attn_qkv=False,
                  mesh=None,
                  mesh_idx=None):
         self._config = locals()
@@ -364,6 +368,7 @@ def __init__(self,
             dropout=attn_dropout,
             weight_attr=weight_attrs[0],
             bias_attr=bias_attrs[0],
+            fuse_attn_qkv=fuse_attn_qkv,
             mesh=mesh,
             mesh_idx=mesh_idx)
 
@@ -483,6 +488,7 @@ def __init__(self,
                  max_position_embeddings=512,
                  type_vocab_size=16,
                  initializer_range=0.02,
+                 fuse_attn_qkv=False,
                  mesh=None):
 
         super(GPTModelAuto, self).__init__()
@@ -518,6 +524,7 @@ def __init__(self,
                         initializer=nn.initializer.Normal(
                             mean=0.0, std=self.initializer_range)),
                     bias_attr=None,
+                    fuse_attn_qkv=fuse_attn_qkv,
                     mesh=self.mesh,
                     mesh_idx=stages[i]))
 

projects/gpt/auto_gpt_345M_single_card.sh

@@ -1,4 +1,3 @@
-
 #! /bin/bash
 
 # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.

projects/gpt/docs/auto_parallel.md

@@ -88,6 +88,7 @@ Engine训练设置完成模型训练/验证/推理等过程中的参数设置，
     type_vocab_size: 16
     initializer_range: 0.02
     use_recompute: True
+    fuse_attn_qkv: True
 ```
 
 其中参数对应的释义如下：
@@ -105,6 +106,7 @@ Engine训练设置完成模型训练/验证/推理等过程中的参数设置，
 | type_vocab_size              | 词表类型                   |
 | initializer_range            | 参数初始化的范围               |
 | use_recompute                | 是否使用recompute训练，重计算全部transformer  |
+| fuse_attn_qkv                | 是否对attention层中qkv计算使用fuse代替传统Linear加速训练 |
 
 
 ### 数据集