# coding=utf-8 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import sys import warnings from typing import Dict, List, Optional, Tuple import torch import torch.distributed as dist import torch.nn.functional as F import torch.utils.checkpoint import torch_npu from torch import nn from torch.distributed.distributed_c10d import _world from transformers.activations import ACT2FN from transformers.cache_utils import Cache from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask from transformers.modeling_utils import PreTrainedModel from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_13 from transformers.utils.import_utils import is_torch_fx_available sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from configuration_openpangu_moe import PanguUltraMoEConfig if is_torch_fx_available(): if not is_torch_greater_or_equal_than_1_13: import torch.fx _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) class PanguUltraMoERMSNorm(nn.Module): def __init__(self, hidden_dim, epsilon=1e-5): super().__init__() self.weight = nn.Parameter(torch.empty(hidden_dim)) self.epsilon = epsilon def forward(self, hidden_states, *args): if len(args) == 0: result = torch_npu.npu_rms_norm(hidden_states, self.weight, self.epsilon)[0] return result elif len(args) == 1 and args[0] is None: result = torch_npu.npu_rms_norm(hidden_states, self.weight, self.epsilon)[0] residual = hidden_states return (result, residual) elif len(args) == 1: residual = args[0] y, _, x = torch_npu.npu_add_rms_norm( residual, hidden_states, self.weight, self.epsilon ) return (y, x) else: raise NotImplementedError(f"PanguUltraMoERMSNorm inner error") class PanguUltraMoERotaryEmbedding(nn.Module): def __init__( self, dim, max_position_embeddings=131072, base=25600000.0, device=None ): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base self._set_cache( seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype(), ) def _set_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len dim = self.dim inv_freq = 1.0 / ( self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) ) self.register_buffer("inv_freq", inv_freq, persistent=False) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) def forward(self, x, kv_len, max_seq_len=None): if max_seq_len is None: self._set_cache(seq_len=kv_len, device=x.device, dtype=x.dtype) elif max_seq_len > self.max_seq_len_cached: self._set_cache(seq_len=max_seq_len, device=x.device, dtype=x.dtype) batch_size = x.shape[0] seq_len = x.shape[1] if seq_len == 1: cos = ( torch.index_select(self.cos_cached, dim=0, index=kv_len) .unsqueeze(1) .unsqueeze(1) ) sin = ( torch.index_select(self.sin_cached, dim=0, index=kv_len) .unsqueeze(1) .unsqueeze(1) ) else: cos = ( self.cos_cached[:seq_len] .unsqueeze(0) .unsqueeze(2) .repeat(batch_size, 1, 1, 1) ) sin = ( self.sin_cached[:seq_len] .unsqueeze(0) .unsqueeze(2) .repeat(batch_size, 1, 1, 1) ) cos = cos[0, :, 0, :] sin = sin[0, :, 0, :] return ( cos.to(dtype=x.dtype), sin.to(dtype=x.dtype), ) def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): """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. """ cos = cos[position_ids].unsqueeze(unsqueeze_dim) sin = sin[position_ids].unsqueeze(unsqueeze_dim) b, h, s, d = q.shape q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) b, h, s, d = k.shape k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class MLP(nn.Module): def __init__(self, config, runner_config, hidden_size=None, intermediate_size=None): super().__init__() self.runner_config = runner_config self.moe_tp_size = self.runner_config.get("parallel_config").get( "moe_tp_size", 1 ) self.hidden_size = config.hidden_size if hidden_size is None else hidden_size self.intermediate_size = ( config.intermediate_size if intermediate_size is None else intermediate_size ) self.intermediate_size_per_rank = self.intermediate_size // self.moe_tp_size self.merge_up_gate_proj = nn.Linear( self.hidden_size, self.intermediate_size_per_rank * 2, bias=False ) self.down_proj = nn.Linear( self.intermediate_size_per_rank, self.hidden_size, bias=False ) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): merged_x = self.merge_up_gate_proj(x) gate_state, up_state = merged_x.chunk(2, dim=-1) intermediate_hidden_states = self.act_fn(gate_state) * up_state down_proj = self.down_proj(intermediate_hidden_states) if self.moe_tp_size > 1: dist.all_reduce(down_proj) return down_proj class MoE(nn.Module): def __init__(self, config, runner_config, hidden_size=None, intermediate_size=None): super().__init__() self.runner_config = runner_config self.moe_tp_size = self.runner_config.get("parallel_config").get( "moe_tp_size", 1 ) self.num_experts = config.num_routed_experts self.hidden_size = config.hidden_size if hidden_size is None else hidden_size self.intermediate_size = ( config.intermediate_size if intermediate_size is None else intermediate_size ) self.intermediate_size_per_rank = self.intermediate_size // self.moe_tp_size self.act_fn = ACT2FN[config.hidden_act] self.group_w1_w3 = nn.Parameter( torch.ones( self.num_experts, self.intermediate_size_per_rank * 2, self.hidden_size ), requires_grad=False, ) self.group_w2 = nn.Parameter( torch.ones( self.num_experts, self.hidden_size, self.intermediate_size_per_rank ), requires_grad=False, ) def forward(self, hidden_states, expert_tokens, seq_len=None): mm1_mm3 = torch_npu.npu_grouped_matmul( [hidden_states], [torch.transpose(self.group_w1_w3, 1, 2)], group_list=expert_tokens, group_type=0, split_item=3, )[0] mm1, mm3 = mm1_mm3.chunk(2, dim=-1) intermediate_hidden_states = self.act_fn(mm1) * mm3 hidden_states = torch_npu.npu_grouped_matmul( [intermediate_hidden_states], [torch.transpose(self.group_w2, 1, 2)], group_list=expert_tokens, group_type=0, split_item=3, )[0] return hidden_states class MoEGate(nn.Module): def __init__(self, config): super().__init__() self.top_k = config.num_experts_per_tok self.routed_scaling_factor = config.routed_scaling_factor self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter( torch.empty((config.num_routed_experts, config.hidden_size)) ) def forward(self, hidden_states): bsz, seq_len, h = hidden_states.shape hidden_states = hidden_states.view(-1, h) logits = F.linear( hidden_states.to(torch.float32), self.weight.to(torch.float32), None ) scores = logits.sigmoid() scores_for_choice = scores.view(bsz * seq_len, -1) _, topk_idx = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) topk_weight = scores.gather(1, topk_idx) if self.top_k > 1 and self.norm_topk_prob: denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 topk_weight = topk_weight / denominator topk_weight = topk_weight * self.routed_scaling_factor return topk_idx, topk_weight class PanguUltraMoE(nn.Module): def __init__(self, config, runner_config): super().__init__() self.runner_config = runner_config self.hidden_dim = config.hidden_size self.moe_tp_size = self.runner_config.get("parallel_config").get( "moe_tp_size", 1 ) self.batch_size_decode = self.runner_config.get("data_config").get( "batch_size", 1 ) self.batch_size_prefill = self.batch_size_decode self.num_experts_per_tok = config.num_experts_per_tok self.num_experts = config.num_routed_experts self.num_shared_experts = config.num_shared_experts self.top_k = config.num_experts_per_tok self.experts_per_rank = config.num_routed_experts self.experts = MoE( config, self.runner_config, intermediate_size=config.moe_intermediate_size ) self.gate = MoEGate(config) if self.num_shared_experts is not None: intermediate_size = config.moe_intermediate_size * self.num_shared_experts self.shared_experts = MLP( config, self.runner_config, intermediate_size=intermediate_size ) self.row_idx_decode_len = self.batch_size_decode * self.top_k self.row_idx_decode = ( torch.arange(0, self.row_idx_decode_len, dtype=torch.int32) .view(self.top_k, -1) .permute(1, 0) .int() .contiguous() .npu() ) def forward(self, hidden_states): identity = hidden_states topk_idx, topk_weight = self.gate(hidden_states) y = self.moe_npu(hidden_states, topk_idx, topk_weight) if self.num_shared_experts is not None: y = y + self.shared_experts(identity) return y def moe_npu(self, x, topk_ids, topk_weight): batch_size, sequence_length, h = x.shape hidden_states = x.view(-1, x.shape[-1]) routing_weights = topk_weight.to(x.dtype) expert_idx = topk_ids.int() if sequence_length == 1: row_idx = self.row_idx_decode else: row_idx_prefill_len = self.batch_size_prefill * sequence_length * self.top_k row_idx = ( torch.arange( 0, row_idx_prefill_len, dtype=torch.int32, device=topk_weight.device ) .view(self.top_k, -1) .permute(1, 0) .int() .contiguous() ) active_num = batch_size * sequence_length expanded_x, expanded_row_idx, expanded_expert_idx = ( torch_npu.npu_moe_init_routing( hidden_states, row_idx=row_idx, expert_idx=expert_idx, active_num=active_num, ) ) expert_tokens = torch_npu.npu_moe_compute_expert_tokens( expanded_expert_idx, self.num_experts ) expert_tokens = expert_tokens.to(torch.int64) hidden_states_ordered_by_experts = self.experts( expanded_x, expert_tokens, seq_len=sequence_length ) hidden_states = torch_npu.npu_moe_finalize_routing( hidden_states_ordered_by_experts, skip1=None, skip2=None, bias=None, scales=routing_weights, expanded_src_to_dst_row=expanded_row_idx, export_for_source_row=expert_idx, ) if self.moe_tp_size > 1: dist.all_reduce(hidden_states) hidden_states = hidden_states.view(batch_size, -1, self.hidden_dim) return hidden_states class PanguUltraMoEAttention(nn.Module): def __init__( self, config: PanguUltraMoEConfig, layer_idx: Optional[int] = None, runner_config: Optional[Dict] = None, ): super().__init__() if runner_config is not None: self.attn_tp_size = runner_config.get("parallel_config").get( "attn_tp_size", 1 ) else: self.attn_tp_size = 1 self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.num_heads_per_rank = self.num_heads // self.attn_tp_size self.num_key_value_heads_per_rank = self.num_heads_per_rank self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.attention_q_lora_dim = config.attention_q_lora_dim self.attention_qk_rope_dim = config.attention_qk_rope_dim self.attention_kv_lora_dim = config.attention_kv_lora_dim self.attention_v_dim = config.attention_v_dim self.attention_qk_dim = config.attention_qk_dim self.q_head_dim = config.attention_qk_dim + config.attention_qk_rope_dim if self.attention_q_lora_dim is None: self.q_proj = nn.Linear( self.hidden_size, self.num_heads_per_rank * self.q_head_dim, bias=False ) else: self.q_a_proj = nn.Linear( self.hidden_size, config.attention_q_lora_dim, bias=False ) self.q_a_layernorm = PanguUltraMoERMSNorm(config.attention_q_lora_dim) self.q_b_proj = nn.Linear( config.attention_q_lora_dim, self.num_heads_per_rank * self.q_head_dim, bias=False, ) self.kv_a_proj_with_mqa = nn.Linear( self.hidden_size, config.attention_kv_lora_dim + config.attention_qk_rope_dim, bias=False, ) self.kv_a_layernorm = PanguUltraMoERMSNorm(config.attention_kv_lora_dim) self.kv_b_proj_w_k = nn.Parameter( torch.zeros( self.num_heads_per_rank, self.attention_qk_dim, self.attention_kv_lora_dim, ) ) self.kv_b_proj_w_v = nn.Parameter( torch.zeros( self.num_heads_per_rank, self.attention_kv_lora_dim, self.attention_v_dim, ) ) self.o_proj = nn.Linear( self.num_heads_per_rank * self.attention_v_dim, self.hidden_size, bias=False, ) self.softmax_scale = self.q_head_dim ** (-0.5) def bmm_5d(self, x, y): b, s, n, _, d = x.shape x = x.view(b * s, n, d).transpose(0, 1) output = torch.matmul(x, y) output = output.transpose(1, 0).view(b, s, n, -1) return output def prepare_qkv( self, hidden_states: torch.Tensor, cos_sin: torch.Tensor = None, kv_len: torch.IntTensor = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, **kwargs, ): bsz, q_len, _ = hidden_states.size() if self.attention_q_lora_dim is None: q = self.q_proj(hidden_states) else: q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) compressed_kv, k_pe = torch.split( compressed_kv, [self.attention_kv_lora_dim, self.attention_qk_rope_dim], dim=-1, ) q = q.view(bsz, q_len, self.num_heads_per_rank, self.q_head_dim) q_nope, q_pe = torch.split( q, [self.attention_qk_dim, self.attention_qk_rope_dim], dim=-1 ) q_pe = q_pe.transpose(1, 2) q_nope = self.bmm_5d( q_nope.view(bsz, q_len, self.num_heads_per_rank, 1, self.attention_qk_dim), self.kv_b_proj_w_k, ) q_nope = q_nope.view( bsz, q_len, self.num_heads_per_rank, self.attention_kv_lora_dim ) q_nope = q_nope.transpose(1, 2) k_pe = k_pe.view(bsz, q_len, 1, self.attention_qk_rope_dim).transpose(1, 2) k_nope = ( self.kv_a_layernorm(compressed_kv) .view(bsz, -1, 1, self.attention_kv_lora_dim) .transpose(1, 2) ) cos, sin = cos_sin q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) query_states = torch.cat([q_nope, q_pe], dim=-1) key_states = torch.cat([k_nope, k_pe], dim=-1) kv_seq_len = k_nope.shape[-2] if past_key_value is not None: past_key_states = past_key_value[self.layer_idx][0] torch_npu.scatter_update_(past_key_states, kv_len, key_states, -2) if q_len == 1: key_states = past_key_states kv_seq_len = past_key_value[0][0].size()[-2] value_states = key_states return query_states, key_states, value_states, kv_seq_len def apply_attention_npu( self, query_states, key_states, value_states, kv_seq_len, attention_mask: Optional[torch.Tensor] = None, actual_seq_lengths_kv: list = None, output_attentions: bool = False, past_key_value: Optional[Cache] = None, ): # repeat k/v heads if n_kv_heads < n_heads bsz, _, q_len, _ = query_states.size() attn_weights = ( torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale ) if attention_mask is not None: attn_weights = attn_weights + attention_mask else: raise ValueError("attention mask must not be None") attn_weights = nn.functional.softmax( attn_weights, dim=-1, dtype=torch.float32 ).to(query_states.dtype) value_states = value_states[..., : self.attention_kv_lora_dim] attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = self.bmm_5d(attn_output.unsqueeze(3), self.kv_b_proj_w_v) attn_output = self.o_proj(attn_output.reshape(bsz, q_len, -1)) if self.attn_tp_size > 1: dist.all_reduce(attn_output) return attn_output def forward( self, hidden_states: torch.Tensor, kv_len: torch.IntTensor = None, actual_seq_lengths_kv: list = None, cos_sin: torch.Tensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if "padding_mask" in kwargs: warnings.warn( "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" ) query_states, key_states, value_states, kv_seq_len = self.prepare_qkv( hidden_states=hidden_states, cos_sin=cos_sin, kv_len=kv_len, position_ids=position_ids, past_key_value=past_key_value, ) output = self.apply_attention_npu( query_states=query_states, key_states=key_states, value_states=value_states, kv_seq_len=kv_seq_len, actual_seq_lengths_kv=actual_seq_lengths_kv, attention_mask=attention_mask, output_attentions=output_attentions, past_key_value=past_key_value, ) return output class PanguUltraMoEDecoderLayer(nn.Module): def __init__( self, config: PanguUltraMoEConfig, runner_config: Dict, layer_idx: int ): super().__init__() self.runner_config = runner_config self.hidden_size = config.hidden_size self.self_attn = PanguUltraMoEAttention( config=config, runner_config=self.runner_config, layer_idx=layer_idx ) self.mlp = ( PanguUltraMoE(config, self.runner_config) if ( config.num_routed_experts is not None and layer_idx >= config.num_dense_layers ) else MLP(config, self.runner_config) ) self.input_layernorm = PanguUltraMoERMSNorm( config.hidden_size, epsilon=config.rms_norm_eps ) self.post_attention_layernorm = PanguUltraMoERMSNorm( config.hidden_size, epsilon=config.rms_norm_eps ) if getattr(config, "sandwich_norm", False): self.sandwich_norm = True self.pre_mlp_layernorm = PanguUltraMoERMSNorm( config.hidden_size, epsilon=config.rms_norm_eps ) self.post_mlp_layernorm = PanguUltraMoERMSNorm( config.hidden_size, epsilon=config.rms_norm_eps ) else: self.sandwich_norm = False def forward( self, hidden_states: torch.Tensor, kv_len: torch.IntTensor, actual_seq_lengths_kv: list, cos_sin: torch.Tensor, past_residual: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.FloatTensor]: hidden_states, residual = self.input_layernorm(hidden_states, past_residual) # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, kv_len=kv_len, actual_seq_lengths_kv=actual_seq_lengths_kv, cos_sin=cos_sin, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, ) if self.sandwich_norm: hidden_states = self.post_attention_layernorm(hidden_states) hidden_states, residual = self.pre_mlp_layernorm(hidden_states, residual) else: hidden_states, residual = self.post_attention_layernorm( hidden_states, residual ) hidden_states = self.mlp(hidden_states) if self.sandwich_norm: hidden_states = self.post_mlp_layernorm(hidden_states) outputs = (residual, hidden_states) return outputs class PanguUltraMoEPreTrainedModel(PreTrainedModel): config_class = PanguUltraMoEConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["PanguUltraMoEDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_cache_class = True def _init_weights(self, module): pass class PanguUltraMoEModel(PanguUltraMoEPreTrainedModel): def __init__(self, config: PanguUltraMoEConfig, runner_config: Dict): super().__init__(config) self.config = config self.runner_config = runner_config self.local_rank = int(os.getenv("LOCAL_RANK", "0")) self.rank_offset = int(os.getenv("RANK_OFFSET", "0")) self.global_rank = self.local_rank + self.rank_offset self.embed_tp_size = self.runner_config.get("parallel_config").get( "embed_tp_size", 1 ) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.vocab_size_per_rank = self.vocab_size // self.embed_tp_size self.embed_tokens = nn.Embedding( self.vocab_size_per_rank, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ PanguUltraMoEDecoderLayer(config, self.runner_config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = PanguUltraMoERMSNorm(config.hidden_size, epsilon=config.rms_norm_eps) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() self.rotary_emb = PanguUltraMoERotaryEmbedding( self.config.attention_qk_rope_dim, max_position_embeddings=self.config.max_position_embeddings, base=self.config.rope_theta, ) def forward( self, input_ids: torch.LongTensor, kv_len: torch.IntTensor = None, actual_seq_lengths_kv: list = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, ): batch_size, seq_length = input_ids.shape past_key_values_length = past_key_values[0][0].size()[-2] if position_ids is None: device = input_ids.device position_ids = torch.arange( past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device, ) position_ids = position_ids.unsqueeze(0).view(-1, seq_length) else: position_ids = position_ids.view(-1, seq_length).long() if self.embed_tp_size > 1: new_input_ids = input_ids - self.global_rank * self.vocab_size_per_rank mask = (new_input_ids >= 0) & ( new_input_ids < self.vocab_size_per_rank ) # (bs, qlen) new_input_ids_per_rank = new_input_ids * mask inputs_embeds = self.embed_tokens(new_input_ids_per_rank) * mask.unsqueeze( -1 ) dist.all_reduce(inputs_embeds) else: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds cos_sin = self.rotary_emb( hidden_states, kv_len, self.config.max_position_embeddings ) residual = None for decoder_layer in self.layers: residual, hidden_states = decoder_layer( hidden_states, kv_len, actual_seq_lengths_kv, cos_sin=cos_sin, past_residual=residual, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values, ) hidden_states, _ = self.norm(hidden_states, residual) return hidden_states class PanguUltraMoEForCausalLM(PanguUltraMoEPreTrainedModel): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config, runner_config): super().__init__(config) self.config = config self.runner_config = runner_config self.embed_tp_size = self.runner_config.get("parallel_config").get( "embed_tp_size", 1 ) self.model = PanguUltraMoEModel(config, self.runner_config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear( config.hidden_size, config.vocab_size // self.embed_tp_size, bias=False ) def forward( self, input_ids: torch.LongTensor = None, kv_len: torch.IntTensor = None, actual_seq_lengths_kv: list = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, ): outputs = self.model( input_ids=input_ids, kv_len=kv_len, actual_seq_lengths_kv=actual_seq_lengths_kv, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, ) hidden_states = outputs if hidden_states.size()[1] > 1: gather_index, _ = torch.max(position_ids, dim=-1) gather_index = ( gather_index.unsqueeze(1) .unsqueeze(2) .repeat(1, 1, hidden_states.shape[-1]) ) hidden_states = torch.gather(hidden_states, 1, gather_index) logits = self.lm_head(hidden_states) if self.embed_tp_size > 1: new_logits = torch.zeros_like(logits).repeat(self.embed_tp_size, 1, 1) dist.all_gather_into_tensor(new_logits, logits, group=_world._default_pg) new_logits = new_logits.reshape( self.embed_tp_size, logits.shape[0], logits.shape[1], -1 ).permute(1, 2, 0, 3) logits = new_logits.reshape(logits.shape[0], logits.shape[1], -1) logits = logits.float() return logits def init_cache(self, input_ids): batch_size, seq_len = input_ids.size() cache_seq_len = self.config.max_position_embeddings past_key_values = () cache_key_shape = ( batch_size, 1, cache_seq_len, self.config.attention_kv_lora_dim + self.config.attention_qk_rope_dim, ) dtype = self.config.torch_dtype for _ in range(self.config.num_hidden_layers): key_cache = torch.zeros( cache_key_shape, dtype=dtype, device=input_ids.device ) past_key_values += ((key_cache,),) return past_key_values def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, is_prefill=None, kv_len=None, share_mask_tril=None, **kwargs, ): batch_size, seq_len = input_ids.size() if past_key_values is None: past_key_values = self.init_cache(input_ids) if is_prefill: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) attention_mask = share_mask_tril kv_len = torch.zeros( (position_ids.size()[0]), dtype=torch.int32, device=input_ids.device ) actual_seq_lengths_kv = None past_key_values_length = 0 input_mask = None else: attention_mask = None position_ids = kv_len.unsqueeze(1) actual_seq_lengths_kv = (kv_len + 1).cpu().detach().numpy().tolist() past_key_values_length = self.config.max_position_embeddings - seq_len input_mask = share_mask_tril attention_mask = _prepare_4d_causal_attention_mask( input_mask, (batch_size, seq_len), input_ids.float(), past_key_values_length ) model_inputs = {} model_inputs.update( { "input_ids": input_ids, "position_ids": position_ids, "past_key_values": past_key_values, "attention_mask": attention_mask, "kv_len": kv_len, "actual_seq_lengths_kv": actual_seq_lengths_kv, } ) return model_inputs