[Auto Parallel] Fix sp in gpt modeling_auto (#10835)

* Update modeling_auto.py

* Update pretrain-gpt3_13b_dynamic_auto.json

* update

* fix gpt sp model

* Update loss base

* update loss base

* disable benchmark SP for now

Author: Xing-lil

paddlenlp/transformers/gpt/modeling_auto.py

@@ -32,13 +32,6 @@
 from paddle.distributed.fleet.utils import recompute
 from paddle.utils import try_import
 
-try:
-    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-        mark_as_sequence_parallel_parameter,
-    )
-except:
-    pass
-
 from ...utils.converter import StateDictNameMapping
 from .. import PretrainedModel, register_base_model
 from ..model_outputs import BaseModelOutputWithPastAndCrossAttentions
@@ -209,19 +202,19 @@ def __init__(self, config, ipp=None):
         )
 
     def _fuse_prepare_qkv(self, query, use_cache=False, past_key_value=None):
-        if self.config.sequence_parallel:
-            # [bs, seq_len, num_head * head_dim] -> [bs / n, seq_len, num_head, head_dim] (n is model parallelism)
-            target_shape = [-1, self.config.seq_length, self.num_attention_heads, 3 * self.head_dim]
-        else:
-            target_shape = [0, 0, self.num_attention_heads, 3 * self.head_dim]
-
+        target_shape = [0, 0, self.num_attention_heads, 3 * self.head_dim]
         # bs, seq_len, num_head * 3*head_dim
         mix_layer = self.qkv_proj(query)
         # bs, seq_len, num_head, 3*head_dim
         mix_layer = paddle.reshape_(mix_layer, target_shape)
         # query_states, key_states, value_states => bs, seq_len, num_head, head_dim
         query_states, key_states, value_states = paddle.split(mix_layer, num_or_sections=3, axis=-1)
-
+        if self.config.sequence_parallel:
+            # [seq_len, bs, num_head * head_dim] -> [bs, seq_len, num_head * head_dim]  (if sequence_parallel)
+            # FA and rope not support sequence first
+            query_states = paddle.transpose(query_states, [1, 0, 2, 3])
+            key_states = paddle.transpose(key_states, [1, 0, 2, 3])
+            value_states = paddle.transpose(value_states, [1, 0, 2, 3])
         # [bs, seq_len, num_head, head_dim]
         if past_key_value is not None:
             # reuse k, v, self_attention
@@ -326,6 +319,8 @@ def forward(
         Applies multi-head attention to map queries and a set of key-value pairs
         to outputs.
         """
+        if self.config.sequence_parallel:
+            query = dist.reshard(query, get_mesh(self.ipp), [dist.Shard(1), dist.Replicate()])
         key = query if key is None else key
         value = query if value is None else value
         if self.config.fuse_attention_qkv:
@@ -363,11 +358,11 @@ def forward(
         # else their shape are [bs, q_len, num_head * head_dim / n], n is mp parallelism.
 
         if self.config.sequence_parallel:
-            bs, seq_len, dim = out.shape
-            out = out.reshape([bs * seq_len, dim])  # [bs, seq_len, dim / n] => [bs * seq_len, dim / n]
-
+            out = paddle.transpose(out, [1, 0, 2])
         # project to output
         out = self.out_proj(out)
+        if self.config.sequence_parallel:
+            out = dist.reshard(out, get_mesh(self.ipp), [dist.Shard(1), dist.Shard(0)])
         # if sequence_parallel is true, out shape are [bs * seq_len / n, dim]
         # else their shape are [bs, seq_len, dim], n is mp parallelism.
         outs = [out]
@@ -390,9 +385,6 @@ def __init__(self, config, decoder_layers, norm=None, hidden_size=None):
         self.layers = decoder_layers
 
         self.norm = GPTLayerNorm(config, config.hidden_size, epsilon=1e-5)
-        if config.sequence_parallel:
-            mark_as_sequence_parallel_parameter(self.norm.weight)
-            mark_as_sequence_parallel_parameter(self.norm.bias)
 
         # Note that we will actually perform a recompute only if both enable_recompute and layerwise_recompute are set to True
         # Enable_recompute defaults to False and is controlled by Trainer
@@ -536,11 +528,6 @@ def __init__(self, config: GPTConfig, ipp=None):
         self.norm1 = GPTLayerNorm(config, config.hidden_size, self.ipp, epsilon=1e-5, bias_attr=True)
         self.norm2 = GPTLayerNorm(config, config.hidden_size, self.ipp, epsilon=1e-5, bias_attr=True)
 
-        if config.sequence_parallel:
-            mark_as_sequence_parallel_parameter(self.norm1.weight)
-            mark_as_sequence_parallel_parameter(self.norm1.bias)
-            mark_as_sequence_parallel_parameter(self.norm2.weight)
-            mark_as_sequence_parallel_parameter(self.norm2.bias)
         if config.use_fused_dropout_add:
             self.fused_dropout_add1 = FusedDropoutAdd(config.attention_probs_dropout_prob, mode="upscale_in_train")
             self.fused_dropout_add2 = FusedDropoutAdd(config.hidden_dropout_prob, mode="upscale_in_train")
@@ -593,6 +580,12 @@ def forward(
 
         # Use a ternary operator for a more concise assignment of current_seed
         current_seed = "local_seed" if self.config.sequence_parallel else "global_seed"
+        if self.config.sequence_parallel:
+            hidden_states = dist.reshard(
+                hidden_states,
+                get_mesh(self.ipp),
+                [dist.Shard(1), dist.Shard(0)],
+            )
 
         # The 'with' block ensures the correct seed context is used
         with seed_guard_context(current_seed):
@@ -607,14 +600,17 @@ def forward(
         residual = hidden_states
         if self.config.normalize_before:
             hidden_states = self.norm2(hidden_states)
-
+        if self.config.sequence_parallel:
+            hidden_states = dist.reshard(hidden_states, get_mesh(self.ipp), [dist.Shard(1), dist.Replicate()])
         # when sequence_parallel=True:
         # hidden_states => [bs * seq_len / n, embed_dim]
         with seed_guard_context(current_seed):
             if not self.config.use_fused_dropout_add:
                 l_1 = self.linear1(hidden_states)
                 act = self.activation(l_1, approximate=True)
                 l_2 = self.linear2(act)
+                if self.config.sequence_parallel:
+                    l_2 = dist.reshard(l_2, get_mesh(self.ipp), [dist.Shard(1), dist.Shard(0)])
                 hidden_states = residual + self.dropout2(l_2)
             else:
                 hidden_states = self.fused_dropout_add2(
@@ -658,7 +654,7 @@ def __init__(
             config.hidden_size,
         )
         self.word_embeddings.weight = dist.shard_tensor(
-            self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Replicate()]
+            self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)]
         )
         self.position_embeddings.weight = dist.shard_tensor(
             self.position_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)]
@@ -685,18 +681,15 @@ def forward(self, input_ids, position_ids=None, inputs_embeddings=None):
             position_embeddings = self.position_embeddings(position_ids)
         embeddings = inputs_embeddings + position_embeddings
 
-        # exit()
-        if self.config.sequence_parallel:
-            # embeddings = dist.shard_tensor(embeddings,get_mesh(),[dist.Replicate(),dist.Replicate()])
-            bs, seq_len, hidden_size = embeddings.shape
-            # [bs, seq_len, dim] -> [bs * seq_len, dim]
-            embeddings = paddle.reshape_(embeddings, [bs * seq_len, hidden_size])
-            # [bs * seq_len / n, dim] (n is mp parallelism)
-            # embeddings = ScatterOp.apply(embeddings)
-            embeddings = dist.reshard(embeddings, get_mesh(), [dist.Replicate(), dist.Shard(0)])
         # Use a ternary operator for a more concise assignment of current_seed
         current_seed = "local_seed" if self.config.sequence_parallel else "global_seed"
         # The 'with' block ensures the correct seed context is used
+        if self.config.sequence_parallel:
+            # [B, S, H] -> [S, B, H]
+            embeddings = paddle.transpose(embeddings, [1, 0, 2])
+            embeddings = dist.reshard(embeddings, get_mesh(), [dist.Shard(1), dist.Shard(0)])
+        else:
+            embeddings = dist.reshard(embeddings, get_mesh(), [dist.Shard(0), dist.Replicate()])
         with seed_guard_context(current_seed):
             embeddings = self.dropout(embeddings)
         return embeddings
@@ -1176,13 +1169,16 @@ def __init__(self, config: GPTConfig, embedding_weights=None, ipp=None):
                 shape=[config.vocab_size, config.hidden_size],
                 dtype=paddle.get_default_dtype(),
             )
-            self.weight = dist.shard_tensor(self.weight, get_mesh(self.ipp), [dist.Replicate(), dist.Shard(0)])
+            self.weight = dist.shard_tensor(self.weight, get_mesh(self.ipp), [dist.Replicate(), dist.Shard(1)])
 
     def forward(self, hidden_states, tensor_parallel_output=None):
-
         if self.config.sequence_parallel:
-            hidden_states = dist.reshard(hidden_states, get_mesh(self.ipp), [dist.Replicate(), dist.Replicate()])
-            hidden_states = paddle.reshape(hidden_states, [-1, self.config.seq_length, self.config.hidden_size])
+            hidden_states = dist.reshard(
+                hidden_states,
+                get_mesh(self.ipp),
+                [dist.Shard(1), dist.Shard(0)],
+            )
+            hidden_states = paddle.transpose(hidden_states, [1, 0, 2])
 
         if tensor_parallel_output is None:
             tensor_parallel_output = self.config.tensor_parallel_output

scripts/distribute/ci_case_auto.sh

@@ -2492,11 +2492,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2() {
     ips=-1
     mem=-1
     echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
-    loss_base=10.55853653 # output of dropout is different after supporting spmd
+    loss_base=10.55727577 # output of dropout is different after supporting spmd
     ips_base=-1
     mem_base=-1
     if [ $IS_A100 -ne 0 ];then
-        loss_base=10.56019211 # after add dropout spmd
+        loss_base=10.56668472 # after add dropout spmd
     fi
     check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
     echo "=========== $FUNCNAME run  end ==========="
@@ -2564,11 +2564,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2-MP2() {
     ips=-1
     mem=-1
     echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
-    loss_base=10.5657959 # output of dropout is different after supporting spmd
+    loss_base=10.57985115 # output of dropout is different after supporting spmd
     ips_base=-1
     mem_base=-1
     if [ $IS_A100 -ne 0 ];then
-        loss_base=10.5760107 # after add dropout spmd
+        loss_base=10.57280159 # after add dropout spmd
     fi
     check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
     echo "=========== $FUNCNAME run  end ==========="
@@ -2637,11 +2637,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2-MP2-PP2() {
     mem=-1
     echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
     # loss_base=10.59993172     # note: need to debug
-    loss_base=10.57174778 # output of dropout is different after supporting spmd
+    loss_base=10.57274055 # output of dropout is different after supporting spmd
     ips_base=-1
     mem_base=-1
     if [ $IS_A100 -ne 0 ];then
-        loss_base=10.57701015 # after add dropout spmd
+        loss_base=10.57785797 # after add dropout spmd
     fi
     check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
     echo "=========== $FUNCNAME run  end ==========="
@@ -2710,11 +2710,11 @@ function llm_gpt_dygraph_auto_bs8_fp16_DP2-MP2-PP2() {
     mem=-1
     echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
     # loss_base=10.58456802     # note: need to debug
-    loss_base=10.57304478
+    loss_base=10.57409477
     ips_base=-1
     mem_base=-1
     if [ $IS_A100 -ne 0 ];then
-        loss_base=10.57861042 # after add dropout spmd
+        loss_base=10.57924652 # after add dropout spmd
     fi
     check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
     echo "=========== $FUNCNAME run  end ==========="