diff --git a/paddlenlp/transformers/gpt/modeling_auto.py b/paddlenlp/transformers/gpt/modeling_auto.py index e21067ba42..8b52ff0ed4 100644 --- a/paddlenlp/transformers/gpt/modeling_auto.py +++ b/paddlenlp/transformers/gpt/modeling_auto.py @@ -32,13 +32,6 @@ from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker 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 @@ class MultiHeadAttentionAuto(nn.Layer): ) 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 @@ class MultiHeadAttentionAuto(nn.Layer): 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 @@ class MultiHeadAttentionAuto(nn.Layer): # 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 @@ class TransformerDecoder(nn.Layer): 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 @@ class GPTDecoderLayerAuto(nn.Layer): 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 @@ class GPTDecoderLayerAuto(nn.Layer): # 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,7 +600,8 @@ class GPTDecoderLayerAuto(nn.Layer): 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): @@ -615,6 +609,8 @@ class GPTDecoderLayerAuto(nn.Layer): 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 @@ class GPTEmbeddingsAuto(nn.Layer): 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 @@ class GPTEmbeddingsAuto(nn.Layer): 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 @@ class GPTLMHeadAuto(nn.Layer): 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 diff --git a/scripts/distribute/ci_case_auto.sh b/scripts/distribute/ci_case_auto.sh index 25fd5af46d..cec0187905 100755 --- a/scripts/distribute/ci_case_auto.sh +++ b/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 ==========="