update comment

paddlenlp/transformers/qwen2moe/modeling.py

@@ -30,8 +30,8 @@
 
 try:
     from paddle.incubate.nn.functional import fused_rotary_position_embedding
 except ImportError:
     fused_rotary_position_embedding = None
 from paddle.distributed.fleet.utils.sequence_parallel_utils import (
     ColumnSequenceParallelLinear,
     GatherOp,
@@ -56,8 +56,8 @@
 
 try:
     from paddle.nn.functional.flash_attention import flash_attention
 except:
     flash_attention = None
 
 __all__ = [
     "QWen2MoEModel",
@@ -87,72 +87,72 @@
     Returns:
         The auxiliary loss.
     """
     if gate_logits is None or not isinstance(gate_logits, tuple):
         return 0
 
     if isinstance(gate_logits, tuple):
         concatenated_gate_logits = paddle.concat(
             gate_logits, axis=0
         )  # [num_hidden_layers X batch_size X sequence_length, num_experts]
 
     routing_weights = F.softmax(concatenated_gate_logits, axis=-1)
     _, selected_experts = paddle.topk(routing_weights, top_k, axis=-1)
     expert_mask = F.one_hot(
         selected_experts, num_classes=num_experts
     )  # [num_hidden_layers X batch_size X sequence_length, top_k, num_experts]
 
     if attention_mask is None or len(attention_mask.shape) == 4:
         # Only intokens strategy has 4-D attention_mask, we currently do not support excluding padding tokens.
         # Compute the percentage of tokens routed to each experts
         tokens_per_expert = paddle.mean(expert_mask.astype("float32"), axis=0)
 
         # Compute the average probability of routing to these experts
         router_prob_per_expert = paddle.mean(routing_weights, axis=0)
     else:
         # Exclude the load balancing loss of padding tokens.
         if len(attention_mask.shape) == 2:
             batch_size, sequence_length = attention_mask.shape
             num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
 
             # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
             expert_attention_mask = (
                 attention_mask[None, :, :, None, None]
                 .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
                 .reshape([-1, top_k, num_experts])
             )  # [num_hidden_layers * batch_size * sequence_length, top_k, num_experts]
 
             # Compute the percentage of tokens routed to each experts
             tokens_per_expert = paddle.sum(expert_mask.astype("float32") * expert_attention_mask, axis=0) / paddle.sum(
                 expert_attention_mask, axis=0
             )
 
             # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
             router_per_expert_attention_mask = (
                 attention_mask[None, :, :, None]
                 .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
                 .reshape([-1, num_experts])
             )
 
             # Compute the average probability of routing to these experts
             router_prob_per_expert = paddle.sum(
                 routing_weights * router_per_expert_attention_mask, axis=0
             ) / paddle.sum(router_per_expert_attention_mask, axis=0)
 
     overall_loss = paddle.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
     return overall_loss * num_experts
 
 
 def get_triangle_upper_mask(x, mask=None):
     if mask is not None:
         return mask
     # [bsz, n_head, q_len, kv_seq_len]
     shape = x.shape
     #  [bsz, 1, q_len, kv_seq_len]
     shape[1] = 1
     mask = paddle.full(shape, paddle.finfo(x.dtype).min, dtype=x.dtype)
     mask = paddle.triu(mask, diagonal=1)
     mask.stop_gradient = True
     return mask
 
 
 def assign_kv_heads(num_kv_heads: int, num_gpus: int):
@@ -168,20 +168,20 @@
         assign_kv_heads(num_kv_heads=2, num_gpus=4): [[0],[0],[1],[1]]
         assign_kv_heads(num_kv_heads=4, num_gpus=4): [[0],[1],[2],[3]]
     """
     assignment_list = [[] for _ in range(num_gpus)]
     # Case 1: more heads than cards
     if num_kv_heads > num_gpus:
         num_heads_per_card = num_kv_heads // num_gpus
         for i in range(num_gpus):
             for j in range(num_heads_per_card):
                 assignment_list[i].append(i * num_heads_per_card + j)
     # Case 2: more cards than heads. each card get only 1 head.
     else:
         num_card_per_heads = num_gpus // num_kv_heads
         for i in range(num_kv_heads):
             for j in range(num_card_per_heads):
                 assignment_list[i * num_card_per_heads + j].append(i)
     return assignment_list
 
 
 def parallel_matmul(x: Tensor, y: Tensor, tensor_parallel_output=True):
@@ -189,25 +189,25 @@
     tensor_parallel_degree = 1
     try:
         hcg = fleet.get_hybrid_communicate_group()
         model_parallel_group = hcg.get_model_parallel_group()
         tensor_parallel_degree = hcg.get_model_parallel_world_size()
     except:
         is_fleet_init = False
 
     if paddle.in_dynamic_mode():
         y_is_distributed = y.is_distributed
     else:
         y_is_distributed = tensor_parallel_degree > 1
 
     if is_fleet_init and tensor_parallel_degree > 1 and y_is_distributed:
         # if not running under distributed.launch, it will raise AttributeError: 'Fleet' object has no attribute '_hcg'
         input_parallel = paddle.distributed.collective._c_identity(x, group=model_parallel_group)
         logits = paddle.matmul(input_parallel, y, transpose_y=False)
 
         if tensor_parallel_output:
             return logits
 
         return paddle.distributed.collective._c_concat(logits, group=model_parallel_group)
 
     else:
         logits = paddle.matmul(x, y, transpose_y=False)
@@ -231,9 +231,9 @@
         # Paddle Flash Attention input [ bz, seqlen, nhead, head_dim]
         # Torch Flash Attention input [ bz, nhead, seqlen, head_dim]
 
         version = paddle.version.full_version
         if version != "0.0.0" and version <= "2.5.2":
             attn_output, attn_weights = flash_attention(
                 query_states,
                 key_states,
                 value_states,
@@ -241,7 +241,7 @@
                 return_softmax=output_attentions,
             )
         else:
             attn_output = F.scaled_dot_product_attention(
                 query_states,
                 key_states,
                 value_states,
@@ -250,13 +250,13 @@
                 dropout_p=config.attention_dropout if training else 0.0,
                 training=training,
             )
             attn_weights = None
 
         if sequence_parallel:
             attn_output = attn_output.reshape([bsz * q_len, head_dim * num_heads])
         else:
             attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
         return (attn_output, attn_weights) if output_attentions else attn_output
     else:
         #  [ bz, seqlen, nhead, head_dim] -> [bs, nhead, seq_len, head_dim]
         query_states = paddle.transpose(query_states, [0, 2, 1, 3])
@@ -268,22 +268,22 @@
         attn_weights = paddle.matmul(query_states / math.sqrt(head_dim), key_states.transpose([0, 1, 3, 2]))
 
         if attn_weights.shape != [bsz, num_heads, q_len, kv_seq_len]:
             raise ValueError(
                 f"Attention weights should be of shape {(bsz, num_heads, q_len, kv_seq_len)}, but is"
                 f" {attn_weights.shape}"
             )
 
         if attention_mask is None:
             attention_mask = get_triangle_upper_mask(attn_weights)
         attention_mask = attention_mask.reshape([bsz, 1, q_len, kv_seq_len])
         if attention_mask.shape != [bsz, 1, q_len, kv_seq_len]:
             raise ValueError(
                 f"Attention mask should be of shape {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.shape}"
             )
 
         attn_weights = attn_weights + attention_mask
         if not paddle.in_dynamic_mode():
             attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
         else:
             with paddle.amp.auto_cast(False):
                 attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
@@ -294,22 +294,22 @@
         attn_output = attn_output.transpose([0, 2, 1, 3])
 
         if sequence_parallel:
             attn_output = attn_output.reshape([bsz * q_len, head_dim * num_heads])
         else:
             attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
         return (attn_output, attn_weights) if output_attentions else attn_output
 
 
 def masked_fill(x, mask, value):
     y = paddle.full(x.shape, value, x.dtype)
     return paddle.where(mask, y, x)
 
 
 def is_casual_mask(attention_mask):
     """
     Upper triangular of attention_mask equals to attention_mask is casual
     """
     return (paddle.triu(attention_mask) == attention_mask).all().item()
 
 
 def _make_causal_mask(input_ids_shape, past_key_values_length):
@@ -322,7 +322,7 @@
 
     if past_key_values_length > 0:
         # [tgt_len, tgt_len + past_len]
         mask = paddle.concat([paddle.ones([target_length, past_key_values_length], dtype="bool"), mask], axis=-1)
 
     # [bs, 1, tgt_len, tgt_len + past_len]
     return mask[None, None, :, :].expand([batch_size, 1, target_length, target_length + past_key_values_length])
@@ -355,7 +355,7 @@
         self.config = config
 
         if config.sequence_parallel:
             mark_as_sequence_parallel_parameter(self.weight)
 
     def forward(self, hidden_states):
         if paddle.in_dynamic_mode():
@@ -364,12 +364,12 @@
                 variance = hidden_states.pow(2).mean(-1, keepdim=True)
                 hidden_states = paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
         else:
             hidden_states = hidden_states.astype("float32")
             variance = hidden_states.pow(2).mean(-1, keepdim=True)
             hidden_states = paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
 
         if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
             hidden_states = paddle.cast(hidden_states, self.weight.dtype)
         return hidden_states * self.weight
 
 
@@ -399,7 +399,7 @@
     def forward(self, x, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
         if seq_len > self.max_seq_len_cached:
             self._set_cos_sin_cache(seq_len)
         cos = self.cos_cached[:, :seq_len, :, :]
         sin = self.sin_cached[:, :seq_len, :, :]
         return (
@@ -416,30 +416,10 @@
 
 
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`):
-            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
-            used to pass offsetted position ids when working with a KV-cache.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
     if position_ids is None:
         # Note: Only for QWen2MoEForCausalLMPipe model pretraining
         cos = cos[:, : q.shape[1], :, :]  # [bs, seq_len, 1, dim]
         sin = sin[:, : q.shape[1], :, :]  # [bs, seq_len, 1, dim]
     else:
         cos = cos.squeeze(axis=[0, 2])  # [seq_len, dim]
         sin = sin.squeeze(axis=[0, 2])  # [seq_len, dim]
@@ -461,26 +441,26 @@
         self.tensor_parallel_degree = config.tensor_parallel_degree
 
         if config.sequence_parallel:
             ColumnParallelLinear = ColumnSequenceParallelLinear
             RowParallelLinear = RowSequenceParallelLinear
         else:
             ColumnParallelLinear = fleet.meta_parallel.ColumnParallelLinear
             RowParallelLinear = fleet.meta_parallel.RowParallelLinear
 
         if config.tensor_parallel_degree > 1:
             self.gate_proj = ColumnParallelLinear(
                 self.hidden_size,
                 self.intermediate_size,
                 gather_output=False,
                 has_bias=False,
             )
             self.up_proj = ColumnParallelLinear(
                 self.hidden_size,
                 self.intermediate_size,
                 gather_output=False,
                 has_bias=False,
             )
             self.down_proj = RowParallelLinear(
                 self.intermediate_size,
                 self.hidden_size,
                 input_is_parallel=True,
@@ -506,8 +486,8 @@
     if n_rep == 1:
         return hidden_states
 
     hidden_states = hidden_states.unsqueeze(-2).tile([1, 1, 1, n_rep, 1])
     return hidden_states.reshape([batch, slen, num_key_value_heads * n_rep, head_dim])
 
 
 class QWen2MoEAttention(nn.Layer):
@@ -543,41 +523,41 @@
         self.layerwise_recompute = layerwise_recompute
         self.recompute_granularity = config.recompute_granularity
         if config.tensor_parallel_degree > 1:
             assert (
                 self.num_heads % config.tensor_parallel_degree == 0
             ), f"num_heads: {self.num_heads}, tensor_parallel_degree: {config.tensor_parallel_degree}"
             self.num_heads = self.num_heads // config.tensor_parallel_degree
 
             assert (
                 self.num_key_value_heads % config.tensor_parallel_degree == 0
             ), f"num_key_value_heads: {self.num_key_value_heads}, tensor_parallel_degree: {config.tensor_parallel_degree}"
             self.num_key_value_heads = self.num_key_value_heads // config.tensor_parallel_degree
 
         self.use_fused_rope = config.use_fused_rope
         if self.use_fused_rope:
             if "gpu" not in paddle.device.get_device() or fused_rotary_position_embedding is None:
                 warnings.warn(
                     "Enable fuse rope in the config, but fuse rope is not available. "
                     "Will disable fuse rope. Try using latest gpu version of Paddle."
                 )
                 self.use_fused_rope = False
 
         if config.sequence_parallel:
             ColumnParallelLinear = ColumnSequenceParallelLinear
             RowParallelLinear = RowSequenceParallelLinear
         else:
             ColumnParallelLinear = fleet.meta_parallel.ColumnParallelLinear
             RowParallelLinear = fleet.meta_parallel.RowParallelLinear
 
         if config.tensor_parallel_degree > 1:
             self.q_proj = ColumnParallelLinear(self.hidden_size, self.hidden_size, has_bias=True, gather_output=False)
             self.k_proj = ColumnParallelLinear(
                 self.hidden_size, self.config.num_key_value_heads * self.head_dim, has_bias=True, gather_output=False
             )
             self.v_proj = ColumnParallelLinear(
                 self.hidden_size, self.config.num_key_value_heads * self.head_dim, has_bias=True, gather_output=False
             )
             self.o_proj = RowParallelLinear(self.hidden_size, self.hidden_size, has_bias=False, input_is_parallel=True)
         else:
             self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias_attr=True)
             self.k_proj = nn.Linear(self.hidden_size, self.config.num_key_value_heads * self.head_dim, bias_attr=True)
@@ -610,8 +590,8 @@
         value_states = self.v_proj(hidden_states)
 
         if self.sequence_parallel:
             target_query_shape = [-1, self.seq_length, self.num_heads, self.head_dim]
             target_key_value_shape = [-1, self.seq_length, self.num_key_value_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]
@@ -625,9 +605,9 @@
             kv_seq_len += past_key_value[0].shape[-3]
 
         if self.use_fused_rope:
             assert past_key_value is None, "fuse rotary not support cache kv for now"
             cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
             query_states, key_states, _ = fused_rotary_position_embedding(
                 query_states,
                 key_states,
                 v=None,
@@ -658,7 +638,7 @@
             and has_gradient
             and self.recompute_granularity == "core_attn"
         ):
             outputs = recompute(
                 scaled_dot_product_attention,
                 query_states,
                 self.config,
@@ -730,7 +710,7 @@
             routing_weights = F.softmax(router_logits.astype("float32"), axis=1)
         routing_weights, selected_experts = paddle.topk(routing_weights, self.top_k, axis=-1)
         if self.norm_topk_prob:  # Note: Mixtral is set norm as default, QWen2MoE is set to no norm
             routing_weights /= routing_weights.sum(axis=-1, keepdim=True)
         # we cast back to input dtype
         routing_weights = routing_weights.astype(hidden_states.dtype)
 
@@ -782,7 +762,7 @@
             self.mlp = QWen2MoESparseMoEBlock(config)
         else:
             # num_experts == 0 or this layer is not sparse layer
             self.mlp = QWen2MoEMLP(config)
 
         self.input_layernorm = QWen2MoERMSNorm(config)
         self.post_attention_layernorm = QWen2MoERMSNorm(config)
@@ -835,7 +815,7 @@
             and has_gradient
             and self.recompute_granularity == "full_attn"
         ):
             outputs = recompute(
                 self.self_attn,
                 hidden_states,
                 position_ids,
@@ -875,7 +855,7 @@
         if isinstance(hidden_states, tuple):
             hidden_states, router_logits = hidden_states
         else:
             router_logits = None
 
         hidden_states = residual + hidden_states
 
@@ -888,7 +868,7 @@
             outputs += (present_key_value,)
 
         if output_router_logits:
             outputs += (router_logits,)
 
         if type(outputs) is tuple and len(outputs) == 1:
             outputs = outputs[0]
@@ -950,80 +930,80 @@
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: QWen2MoEConfig, is_split=True):
         from paddlenlp.transformers.conversion_utils import split_or_merge_func
 
         fn = split_or_merge_func(
             is_split=is_split,
             tensor_parallel_degree=config.tensor_parallel_degree,
             tensor_parallel_rank=config.tensor_parallel_rank,
             num_attention_heads=config.num_attention_heads,
         )
 
         def get_tensor_parallel_split_mappings(num_layers, num_experts):
             final_actions = {}
 
             base_actions = {
                 "lm_head.weight": partial(fn, is_column=True),
                 # Row Linear
                 "embed_tokens.weight": partial(fn, is_column=False),
                 "layers.0.self_attn.o_proj.weight": partial(fn, is_column=False),
             }
 
             if not config.vocab_size % config.tensor_parallel_degree == 0:
                 base_actions.pop("lm_head.weight")
                 base_actions.pop("embed_tokens.weight")
 
             # Column Linear
             base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
             base_actions["layers.0.self_attn.q_proj.bias"] = partial(fn, is_column=True)
             # if we have enough num_key_value_heads to split, then split it.
             if config.num_key_value_heads % config.tensor_parallel_degree == 0:
                 base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.k_proj.bias"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.v_proj.bias"] = partial(fn, is_column=True)
 
             for key, action in base_actions.items():
                 if "layers.0." in key:
                     for i in range(num_layers):
                         final_actions[key.replace("layers.0.", f"layers.{i}.")] = action
                 final_actions[key] = action
 
             # Add tp split for expert params.
             base_actions = {
                 "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
                 "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
                 "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
             }
             for key, action in base_actions.items():
                 for i in range(num_layers):
                     newkey = key.replace("layers.0.", f"layers.{i}.")
                     for j in range(num_experts):
                         newkey2 = newkey.replace("experts.0.", f"experts.{j}.")
                         final_actions[newkey2] = action
 
             # Add tp split for shared expert params.
             base_actions = {
                 "layers.0.mlp.shared_expert.gate_proj.weight": partial(fn, is_column=True),
                 "layers.0.mlp.shared_expert.up_proj.weight": partial(fn, is_column=True),
                 "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
             }
             for key, action in base_actions.items():
                 if "layers.0." in key:
                     for i in range(num_layers):
                         final_actions[key.replace("layers.0.", f"layers.{i}.")] = action
                 final_actions[key] = action
 
             return final_actions
 
         mappings = get_tensor_parallel_split_mappings(config.num_hidden_layers, config.num_experts)
 
         return mappings
 
     def _init_weights(self, layer):
         """Initialization hook"""
         if self.config.tensor_parallel_degree > 1:
             rng_tracker = get_rng_state_tracker().rng_state
         if isinstance(
             layer,
             (
@@ -1041,8 +1021,8 @@
             # and reset the `state_dict` to update parameter in static mode.
             if isinstance(layer.weight, paddle.Tensor):
                 if layer.weight.is_distributed:
                     with rng_tracker():
                         layer.weight.set_value(
                             paddle.tensor.normal(
                                 mean=0.0,
                                 std=self.config.initializer_range
@@ -1095,7 +1075,7 @@
         # Recompute defaults to False and is controlled by Trainer
         self.enable_recompute = False
         if config.tensor_parallel_degree > 1 and config.vocab_size % config.tensor_parallel_degree == 0:
             self.embed_tokens = mpu.VocabParallelEmbedding(
                 self.vocab_size,
                 self.hidden_size,
                 weight_attr=paddle.ParamAttr(initializer=nn.initializer.XavierNormal()),
@@ -1140,7 +1120,7 @@
             else:
                 expanded_attn_mask = attention_mask
         else:
             expanded_attn_mask = _make_causal_mask(
                 input_shape,
                 past_key_values_length=past_key_values_length,
             )
@@ -1160,13 +1140,13 @@
         past_key_value: Tensor,
         use_cache: bool,
     ):
         def create_custom_forward(module):
             def custom_forward(*inputs):
                 return module(*inputs)
 
             return custom_forward
 
         hidden_states = recompute(
             create_custom_forward(layer_module),
             hidden_states,
             position_ids,
@@ -1178,7 +1158,7 @@
             use_reentrant=self.config.recompute_use_reentrant,
         )
 
         return hidden_states
 
     def forward(
         self,
@@ -1195,7 +1175,7 @@
         **kwargs,
     ):
         if self.sequence_parallel and use_cache:
             raise ValueError("We currently only support sequence parallel without cache.")
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
 
@@ -1211,13 +1191,13 @@
 
         # retrieve input_ids and inputs_embeds
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
         elif input_ids is not None:
             batch_size, seq_length = input_ids.shape
         elif inputs_embeds is not None:
             batch_size, seq_length, _ = inputs_embeds.shape
         else:
             raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
         if past_key_values is None:
             past_key_values = tuple([None] * len(self.layers))
@@ -1234,10 +1214,10 @@
 
         if self.sequence_parallel:
             # [bs, seq_len, num_head * head_dim] -> [bs * seq_len, num_head * head_dim]
             bs, seq_len, hidden_size = inputs_embeds.shape
             inputs_embeds = paddle.reshape_(inputs_embeds, [bs * seq_len, hidden_size])
             # [seq_len * bs / n, num_head * head_dim] (n is mp parallelism)
             inputs_embeds = ScatterOp.apply(inputs_embeds)
 
         # embed positions
         if attention_mask is None:
@@ -1251,9 +1231,9 @@
             attention_mask, (batch_size, seq_length), cache_length, inputs_embeds.dtype
         )  # [bs, 1, seq_len, seq_len]
         if self.config.use_flash_attention:
             is_casual = is_casual_mask(attention_mask)
             if is_casual:
                 attention_mask = None
         hidden_states = inputs_embeds
 
         # decoder layers
@@ -1274,7 +1254,7 @@
                 and has_gradient
                 and self.recompute_granularity == "full"
             ):
                 layer_outputs = self.recompute_training_full(
                     decoder_layer,
                     hidden_states,
                     position_ids,
@@ -1309,7 +1289,7 @@
                 next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
 
             if output_router_logits:
                 all_router_logits += (layer_outputs[-1],)
 
         hidden_states = self.norm(hidden_states)
 
@@ -1347,26 +1327,26 @@
         self.enable_parallel_cross_entropy = config.tensor_parallel_degree > 1 and config.tensor_parallel_output
 
         if self.enable_parallel_cross_entropy:  # and False: # and lm_head is distributed
             self.loss_func = mpu.ParallelCrossEntropy(ignore_index=self.ignore_index)
         else:
             self.loss_func = paddle.nn.CrossEntropyLoss(reduction="none", ignore_index=self.ignore_index)
 
     def forward(self, prediction_scores, masked_lm_labels):
         if self.enable_parallel_cross_entropy:
             if prediction_scores.shape[-1] == self.config.vocab_size:
                 warnings.warn(
                     f"enable_parallel_cross_entropy, the vocab_size should be splited: {prediction_scores.shape[-1]}, {self.config.vocab_size}"
                 )
                 self.loss_func = paddle.nn.CrossEntropyLoss(reduction="none", ignore_index=self.ignore_index)
 
         with paddle.amp.auto_cast(False):
             masked_lm_loss = self.loss_func(prediction_scores.astype("float32"), masked_lm_labels.unsqueeze(2))
 
             # skip ignore_index which loss == 0
             masked_lm_loss = masked_lm_loss[masked_lm_loss > 0]
             loss = paddle.mean(masked_lm_loss)
 
         return loss
 
 
 class QWen2MoELMHead(nn.Layer):
@@ -1374,7 +1354,7 @@
         super(QWen2MoELMHead, self).__init__()
         self.config = config
         if config.tensor_parallel_degree > 1 and config.vocab_size % config.tensor_parallel_degree == 0:
             vocab_size = config.vocab_size // config.tensor_parallel_degree
         else:
             vocab_size = config.vocab_size
 
@@ -1385,16 +1365,16 @@
         # Must set distributed attr for Tensor Parallel !
         self.weight.is_distributed = True if (vocab_size != config.vocab_size) else False
         if self.weight.is_distributed:
             self.weight.split_axis = 1
 
     def forward(self, hidden_states, tensor_parallel_output=None):
         if self.config.sequence_parallel:
             hidden_states = GatherOp.apply(hidden_states)
             seq_length = self.config.seq_length
             hidden_states = paddle.reshape_(hidden_states, [-1, seq_length, self.config.hidden_size])
 
         if tensor_parallel_output is None:
             tensor_parallel_output = self.config.tensor_parallel_output
 
         logits = parallel_matmul(hidden_states, self.weight, tensor_parallel_output=tensor_parallel_output)
         return logits
@@ -1427,16 +1407,16 @@
         self.qwen2moe.embed_tokens = value
 
     def get_output_embeddings(self):
         return self.lm_head
 
     def set_output_embeddings(self, new_embeddings):
         self.lm_head = new_embeddings
 
     def set_decoder(self, decoder):
         self.qwen2moe = decoder
 
     def get_decoder(self):
         return self.qwen2moe
 
     def prepare_inputs_for_generation(
         self,
@@ -1457,7 +1437,7 @@
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
             model_inputs = {"input_ids": input_ids}
 
@@ -1473,7 +1453,7 @@
         return model_inputs
 
     def _get_model_inputs_spec(self, dtype: str):
         return {
             "input_ids": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
             "attention_mask": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
             "position_ids": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
@@ -1486,12 +1466,12 @@
             model_kwargs["past_key_values"] = outputs[1]
 
         if isinstance(outputs, MoECausalLMOutputWithPast) and "past_key_values" in outputs:
             model_kwargs["past_key_values"] = outputs.past_key_values
 
         # update position_ids
         if "position_ids" in model_kwargs and model_kwargs["position_ids"] is not None:
             position_ids = model_kwargs["position_ids"]
             model_kwargs["position_ids"] = paddle.concat([position_ids, position_ids[..., -1:] + 1], axis=-1)
 
         if not is_encoder_decoder and "attention_mask" in model_kwargs:
             attention_mask = model_kwargs["attention_mask"]
@@ -1550,23 +1530,23 @@
 
         loss = None
         if labels is not None:
             loss = self.criterion(logits, labels)
 
         aux_loss = None
         if output_router_logits:
             aux_loss = load_balancing_loss_func(
                 outputs.router_logits if return_dict else outputs[-1],
                 self.num_experts,
                 self.num_experts_per_tok,
                 attention_mask,
             )
             if labels is not None:
                 loss += self.router_aux_loss_coef * aux_loss
 
         if not return_dict:
             output = (logits,) + outputs[1:]
             if output_router_logits:
                 output = (aux_loss,) + output
             return (loss,) + output if loss is not None else output
 
         return MoECausalLMOutputWithPast(