# coding=utf-8 # Copyright 2024 The HuggingFace Inc. team and Gemma2MoE Contributors. # # 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. """ PyTorch Gemma2MoE model.""" import math from typing import List, Optional, Tuple, Union import torch import torch.nn.functional as F import torch.utils.checkpoint from torch import nn from torch.nn import CrossEntropyLoss from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_outputs import MoeModelOutputWithPast, MoeCausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from transformers.generation import GenerationMixin from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings, ) from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask from .configuration_gemma2moe import Gemma2MoeConfig logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "Gemma2MoeConfig" # --- Auxiliary Loss & Router Functions --- def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float: r""" Computes auxiliary load balancing loss as in Switch Transformer. """ if gate_logits is None or not isinstance(gate_logits, torch.Tensor): return 0.0 # gate_logits: [batch_size * seq_len, num_experts] assumed flattened or [batch, seq, experts] if gate_logits.dim() == 3: gate_logits = gate_logits.view(-1, gate_logits.shape[-1]) routing_weights = torch.softmax(gate_logits, dim=-1) # top_k indices _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # expert_mask: [num_tokens, num_experts] expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) if expert_mask.dim() == 3: expert_mask = expert_mask.sum(dim=1) # Sum over k selected experts # Normalize to get fraction of tokens per expert tokens_per_expert = torch.mean(expert_mask.float(), dim=0) # Mean probability per expert router_prob_per_expert = torch.mean(routing_weights, dim=0) # Loss = N * sum(f_i * P_i) overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert) * num_experts return overall_loss # --- Gemma 2 Components --- class Gemma2RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()) # Gemma 2 signature: output * (1 + weight) # Casting back to input dtype output = output * (1.0 + self.weight.float()) return output.type_as(x) class Gemma2RotaryEmbedding(nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base self.register_buffer("inv_freq", None, persistent=False) @torch.no_grad() def forward(self, x, position_ids, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if self.inv_freq is None: self.inv_freq = 1.0 / ( self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim) ) inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() # Use float32 for RoPE calculation to maintain precision with torch.autocast(device_type=x.device.type, enabled=False): freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() 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): """Applies Rotary Position Embedding to the query and key tensors.""" cos = cos.unsqueeze(1) # [bs, 1, seq_len, head_dim] sin = sin.unsqueeze(1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). Used for Grouped Query Attention (GQA). """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class Gemma2Attention(nn.Module): """ Multi-headed attention with Soft-capping, Sliding Window and GQA. """ def __init__(self, config: Gemma2MoeConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.head_dim self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = True # Gemma 2 scaling specific self.scaling = config.query_pre_attn_scalar ** -0.5 # Soft capping parameter self.attn_logit_soft_capping = config.attn_logit_soft_capping self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) self.rotary_emb = Gemma2RotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) self.sliding_window = config.sliding_window def forward( self, hidden_states: torch.Tensor, 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, cache_position: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids=position_ids, seq_len=None) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # cache_position for static cache, legacy for dynamic cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position, "sliding_window": self.sliding_window} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) # Scaled Dot Product Calculation attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling # Logit Soft Capping if self.attn_logit_soft_capping is not None: attn_weights = attn_weights / self.attn_logit_soft_capping attn_weights = torch.tanh(attn_weights) attn_weights = attn_weights * self.attn_logit_soft_capping if attention_mask is not None: # Mask should be broadcastable attn_weights = attn_weights + attention_mask # Softmax and Dropout attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): raise ValueError( f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" f" {attn_output.size()}" ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value # --- Expert & MoE Block --- class Gemma2MLP(nn.Module): """ Gemma 2 MLP: Gated GELU Tanh """ def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class Gemma2MoeBlock(nn.Module): """ Sparse MoE Block for Gemma 2. Uses Top-k gating and processes selected tokens through experts. """ def __init__(self, config: Gemma2MoeConfig): super().__init__() self.hidden_dim = config.hidden_size self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok self.jitter_noise = config.router_jitter_noise self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False) self.experts = nn.ModuleList([Gemma2MLP(config) for _ in range(self.num_experts)]) def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: batch_size, sequence_length, hidden_dim = hidden_states.shape hidden_states_flat = hidden_states.view(-1, hidden_dim) # Router Logits router_logits = self.gate(hidden_states_flat) if self.training and self.jitter_noise > 0: router_logits += torch.empty_like(router_logits).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) routing_weights = F.softmax(router_logits, dim=1) topk_weight, topk_idx = torch.topk(routing_weights, self.top_k, dim=-1, sorted=False) # Normalize weights topk_weight /= topk_weight.sum(dim=-1, keepdim=True) topk_weight = topk_weight.to(hidden_states.dtype) # Routing process # Using a loop here for clarity and simplicity in Python. # For extreme performance, Triton or CUDA kernels should be used. final_hidden_states = torch.zeros_like(hidden_states_flat) # Flatten indices to handle batching easier flat_topk_idx = topk_idx.view(-1) # We need to process each expert for i, expert in enumerate(self.experts): # Find tokens assigned to this expert (in any of the top-k slots) # This is a bit inefficient in pure PyTorch but ensures correctness without custom kernels # Create a mask for tokens where this expert is selected expert_mask = (topk_idx == i) if expert_mask.any(): # We need to collect inputs, process, and scatter back # This logic handles cases where an expert is selected multiple times (unlikely in top-k but possible conceptually) # But typically top-k implies distinct experts. # Get indices where this expert is used batch_indices, k_indices = torch.where(expert_mask) # Extract inputs inp = hidden_states_flat[batch_indices] # Forward pass out = expert(inp) # Weighting: We need the weight associated with this selection weights = topk_weight[batch_indices, k_indices] # Accumulate result # Ideally, scatter_add, but here we iterate. # Since batch_indices might repeat if we allowed k repetitions (we don't usually), # standard scatter_add_ is safer. weighted_out = out * weights.unsqueeze(-1) final_hidden_states.index_add_(0, batch_indices, weighted_out) final_hidden_states = final_hidden_states.view(batch_size, sequence_length, hidden_dim) return final_hidden_states, router_logits # --- Decoder Layer (Strict Gemma 2 Topology) --- class Gemma2MoeDecoderLayer(nn.Module): def __init__(self, config: Gemma2MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = Gemma2Attention(config, layer_idx) self.block_sparse_moe = Gemma2MoeBlock(config) # Gemma 2 uses 4 specific RMSNorms per layer self.input_layernorm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.pre_feedforward_layernorm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_feedforward_layernorm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: # --- Attention Path --- residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states # Residual Connection # --- MoE Path --- residual = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) # Using MoE instead of standard MLP hidden_states, router_logits = self.block_sparse_moe(hidden_states) hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states # Residual Connection outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) if output_router_logits: outputs += (router_logits,) return outputs # --- PreTrained Model Wrappers --- class Gemma2MoePreTrainedModel(PreTrainedModel): config_class = Gemma2MoeConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["Gemma2MoeDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = False # Keeping SDPA for broad compatibility logic implemented above _supports_sdpa = True _supports_cache_class = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class Gemma2MoeModel(Gemma2MoePreTrainedModel): def __init__(self, config: Gemma2MoeConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [Gemma2MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, MoeModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = output_router_logits if output_router_logits is not None else self.config.output_router_logits output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache() if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) # 4D Attention Mask Creation (handles Sliding Window if config requests it) causal_mask = _prepare_4d_causal_attention_mask( attention_mask, (inputs_embeds.shape[0], inputs_embeds.shape[1]), inputs_embeds, past_key_values.get_seq_length() if past_key_values is not None else 0, sliding_window=self.config.sliding_window, ) # Normalization (Gemma 2 embedding scaling) normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=inputs_embeds.dtype) hidden_states = inputs_embeds * normalizer all_hidden_states = () if output_hidden_states else None all_router_logits = () if output_router_logits else None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, causal_mask, position_ids, past_key_values, output_attentions, output_router_logits, use_cache, cache_position, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, output_router_logits=output_router_logits, use_cache=use_cache, cache_position=cache_position, ) hidden_states = layer_outputs[0] if output_router_logits: all_router_logits += (layer_outputs[-1],) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_router_logits] if v is not None) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden_states, router_logits=all_router_logits, ) class Gemma2MoeForCausalLM(Gemma2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.model = Gemma2MoeModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.router_aux_loss_coef = config.router_aux_loss_coef self.num_experts = config.num_local_experts self.num_experts_per_tok = config.num_experts_per_tok # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.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 forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, MoeCausalLMOutputWithPast]: output_router_logits = output_router_logits if output_router_logits is not None else self.config.output_router_logits return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_router_logits=output_router_logits, return_dict=return_dict, cache_position=cache_position, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) # Final Soft Capping (Gemma 2 Specific feature) # tanh(logits / cap) * cap if self.config.logit_soft_capping is not None: logits = logits / self.config.logit_soft_capping logits = torch.tanh(logits) logits = logits * self.config.logit_soft_capping logits = logits.float() loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1) loss = loss_fct(shift_logits, shift_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, ) 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( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, router_logits=outputs.router_logits, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs ): past_length = 0 if past_key_values is not None: if isinstance(past_key_values, Cache): past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length() # --- HATA DÜZELTMESİ BAŞLANGICI --- # get_max_length metodunun varlığını kontrol ediyoruz if hasattr(past_key_values, "get_max_length") and past_key_values.get_max_length() is not None: max_cache_length = torch.tensor(past_key_values.get_max_length(), device=input_ids.device) else: max_cache_length = None # --- HATA DÜZELTMESİ BİTİŞİ --- cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length) # Legacy Cache (Tuple formatı için) else: past_length = past_key_values[0][0].shape[2] max_cache_length = None # Keep only the unprocessed tokens: if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. if ( max_cache_length is not None and attention_mask is not None and cache_length + input_ids.shape[1] > max_cache_length ): attention_mask = attention_mask[:, -max_cache_length:] position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] # 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} if cache_position is None: # Inputs embeds veya input_ids hangisi varsa onun shape'ini al input_len = model_inputs.get("input_ids", inputs_embeds).shape[1] cache_position = torch.arange(past_length, past_length + input_len, device=input_ids.device) model_inputs.update( { "position_ids": position_ids, "cache_position": cache_position, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "attention_mask": attention_mask, } ) return model_inputs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += ( tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), ) return reordered_past