Aluode/PerceptionLabPortable

1#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2#           This file was automatically generated from src/transformers/models/diffllama/modular_diffllama.py.3#               Do NOT edit this file manually as any edits will be overwritten by the generation of4#             the file from the modular. If any change should be done, please apply the change to the5#                          modular_diffllama.py file directly. One of our CI enforces this.6#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2024 weak-kajuma and the HuggingFace Inc. team. All rights reserved.9#10# This code is based on Llama implementations in this library and Microsoft's11# Differential Transformer implementations.12 13# Licensed under the Apache License, Version 2.0 (the "License");14# you may not use this file except in compliance with the License.15# You may obtain a copy of the License at16#17#     http://www.apache.org/licenses/LICENSE-2.018#19# Unless required by applicable law or agreed to in writing, software20# distributed under the License is distributed on an "AS IS" BASIS,21# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.22# See the License for the specific language governing permissions and23# limitations under the License.24import math25from typing import Optional, Union26 27import torch28from torch import nn29 30from ...activations import ACT2FN31from ...cache_utils import Cache, DynamicCache, StaticCache32from ...generation import GenerationMixin33from ...integrations import use_kernel_forward_from_hub34from ...masking_utils import create_causal_mask35from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask36from ...modeling_layers import (37    GenericForQuestionAnswering,38    GenericForSequenceClassification,39    GenericForTokenClassification,40    GradientCheckpointingLayer,41)42from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast43from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update44from ...modeling_utils import PreTrainedModel45from ...processing_utils import Unpack46from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging47from ...utils.deprecation import deprecate_kwarg48from ...utils.generic import check_model_inputs49from .configuration_diffllama import DiffLlamaConfig50 51 52logger = logging.get_logger(__name__)53 54 55class DiffLlamaMLP(nn.Module):56    def __init__(self, config):57        super().__init__()58        self.config = config59        self.hidden_size = config.hidden_size60        self.intermediate_size = config.intermediate_size61        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)62        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)63        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)64        self.act_fn = ACT2FN[config.hidden_act]65 66    def forward(self, x):67        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))68        return down_proj69 70 71def rotate_half(x):72    """Rotates half the hidden dims of the input."""73    x1 = x[..., : x.shape[-1] // 2]74    x2 = x[..., x.shape[-1] // 2 :]75    return torch.cat((-x2, x1), dim=-1)76 77 78def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):79    """Applies Rotary Position Embedding to the query and key tensors.80 81    Args:82        q (`torch.Tensor`): The query tensor.83        k (`torch.Tensor`): The key tensor.84        cos (`torch.Tensor`): The cosine part of the rotary embedding.85        sin (`torch.Tensor`): The sine part of the rotary embedding.86        position_ids (`torch.Tensor`, *optional*):87            Deprecated and unused.88        unsqueeze_dim (`int`, *optional*, defaults to 1):89            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and90            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note91            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and92            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes93            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have94            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.95    Returns:96        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.97    """98    cos = cos.unsqueeze(unsqueeze_dim)99    sin = sin.unsqueeze(unsqueeze_dim)100    q_embed = (q * cos) + (rotate_half(q) * sin)101    k_embed = (k * cos) + (rotate_half(k) * sin)102    return q_embed, k_embed103 104 105def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:106    """107    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,108    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)109    """110    batch, num_key_value_heads, slen, head_dim = hidden_states.shape111    if n_rep == 1:112        return hidden_states113    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)114    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)115 116 117def lambda_init_fn(layer_idx):118    return 0.8 - 0.6 * math.exp(-0.3 * layer_idx)119 120 121class DiffLlamaAttention(nn.Module):122    """Multi-headed attention from 'Attention Is All You Need' paper"""123 124    def __init__(self, config: DiffLlamaConfig, layer_idx: Optional[int] = None):125        super().__init__()126        self.config = config127        self.layer_idx = layer_idx128        if layer_idx is None:129            logger.warning_once(130                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "131                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "132                "when creating this class."133            )134 135        self.attention_dropout = config.attention_dropout136        self.hidden_size = config.hidden_size137        self.num_heads = config.num_attention_heads138        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)139        self.num_key_value_heads = config.num_key_value_heads140        self.num_key_value_groups = self.num_heads // self.num_key_value_heads141        # under this are not used142        self.max_position_embeddings = config.max_position_embeddings143        self.rope_theta = config.rope_theta144        self.is_causal = True145 146        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)147        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)148        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)149        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)150 151        self.lambda_init = lambda_init_fn(layer_idx)152        self.lambda_q1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,)))153        self.lambda_k1 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,)))154        self.lambda_q2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,)))155        self.lambda_k2 = nn.Parameter(torch.normal(0, config.lambda_std_dev, size=(self.head_dim,)))156        self.groupnorm = nn.RMSNorm(2 * self.head_dim, eps=config.rms_norm_eps, elementwise_affine=False)157 158    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")159    def forward(160        self,161        hidden_states: torch.Tensor,162        position_embeddings: tuple[torch.Tensor, torch.Tensor],163        attention_mask: Optional[torch.Tensor] = None,164        position_ids: Optional[torch.LongTensor] = None,165        past_key_values: Optional[Cache] = None,166        use_cache: bool = False,167        cache_position: Optional[torch.LongTensor] = None,168        **kwargs,169    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:170        bsz, target_len, _ = hidden_states.size()171        q_len = target_len172 173        query_states = self.q_proj(hidden_states)174        key_states = self.k_proj(hidden_states)175        value_states = self.v_proj(hidden_states)176 177        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)178        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)179        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)180 181        cos, sin = position_embeddings182        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)183 184        if past_key_values is not None:185            # sin and cos are specific to RoPE models; cache_position needed for the static cache186            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}187            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)188 189        key_states = repeat_kv(key_states, self.num_key_value_groups)190        value_states = repeat_kv(value_states, self.num_key_value_groups)191        value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1)192        value_states = value_states.repeat(1, 2, 1, 1)193 194        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)195 196        if attention_mask is not None:  # no matter the length, we just slice it197            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]198            attn_weights = attn_weights + causal_mask199 200        # upcast attention to fp32201        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)202        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)203        lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to(204            query_states.dtype205        )206        lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to(207            query_states.dtype208        )209        lambda_full = lambda_1 - lambda_2 + self.lambda_init210 211        attn_output = torch.matmul(attn_weights, value_states)212        attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1)213 214        attn_output = attn_output1 - lambda_full * attn_output2215        attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output)216        attn_output = attn_output.transpose(1, 2).contiguous()217        attn_output = attn_output.reshape(bsz, q_len, -1)218        attn_output = self.o_proj(attn_output)219        return attn_output, attn_weights220 221 222class DiffLlamaFlashAttention2(DiffLlamaAttention):223    """224    DiffLlama flash attention module. This module inherits from `DiffLlamaAttention` as the weights of the module stays225    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of226    flash attention and deal with padding tokens in case the input contains any of them.227    """228 229    def __init__(self, *args, **kwargs):230        super().__init__(*args, **kwargs)231 232        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.233        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.234        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).235        self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask()236 237    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")238    def forward(239        self,240        hidden_states: torch.Tensor,241        position_embeddings: tuple[torch.Tensor, torch.Tensor],242        attention_mask: Optional[torch.LongTensor] = None,243        position_ids: Optional[torch.LongTensor] = None,244        past_key_values: Optional[Cache] = None,245        use_cache: bool = False,246        cache_position: Optional[torch.LongTensor] = None,247    ) -> tuple[torch.Tensor, None]:248        if isinstance(past_key_values, StaticCache):249            raise ValueError(250                "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "251                "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"252            )253 254        bsz, q_len, _ = hidden_states.size()255 256        query_states = self.q_proj(hidden_states)257        key_states = self.k_proj(hidden_states)258        value_states = self.v_proj(hidden_states)259 260        # Flash attention requires the input to have the shape261        # batch_size x seq_length x head_dim x hidden_dim262        # therefore we just need to keep the original shape263        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)264        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)265        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)266 267        if position_embeddings is None:268            logger.warning_once(269                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "270                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "271                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "272                "removed and `position_embeddings` will be mandatory."273            )274            cos, sin = self.rotary_emb(value_states, position_ids)275        else:276            cos, sin = position_embeddings277        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)278 279        if past_key_values is not None:280            # sin and cos are specific to RoPE models; cache_position needed for the static cache281            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}282            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)283 284        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache285        # to be able to avoid many of these transpose/reshape/view.286        query_states = query_states.transpose(1, 2)287        key_states = key_states.transpose(1, 2)288        value_states = value_states.transpose(1, 2)289 290        dropout_rate = self.attention_dropout if self.training else 0.0291 292        # In PEFT, usually we cast the layer norms in float32 for training stability reasons293        # therefore the input hidden states gets silently casted in float32. Hence, we need294        # cast them back in the correct dtype just to be sure everything works as expected.295        # This might slowdown training & inference so it is recommended to not cast the LayerNorms296        # in fp32. (DiffLlamaRMSNorm handles it correctly)297 298        input_dtype = query_states.dtype299        device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"300        if input_dtype == torch.float32:301            if torch.is_autocast_enabled():302                target_dtype = (303                    torch.get_autocast_dtype(device_type)304                    if hasattr(torch, "get_autocast_dtype")305                    else torch.get_autocast_gpu_dtype()306                )307            # Handle the case where the model is quantized308            elif hasattr(self.config, "_pre_quantization_dtype"):309                target_dtype = self.config._pre_quantization_dtype310            else:311                target_dtype = self.q_proj.weight.dtype312 313            logger.warning_once(314                f"The input hidden states seems to be silently casted in float32, this might be related to"315                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"316                f" {target_dtype}."317            )318 319            query_states = query_states.to(target_dtype)320            key_states = key_states.to(target_dtype)321            value_states = value_states.to(target_dtype)322 323        value_states1, value_states2 = torch.chunk(value_states, 2, dim=2)324        value_states1 = value_states1.repeat(1, 1, 2, 1)325        value_states2 = value_states2.repeat(1, 1, 2, 1)326 327        attn_output1 = _flash_attention_forward(328            query_states,329            key_states,330            value_states1,331            attention_mask,332            q_len,333            position_ids=position_ids,334            dropout=dropout_rate,335            sliding_window=getattr(self, "sliding_window", None),336            use_top_left_mask=self._flash_attn_uses_top_left_mask,337            is_causal=self.is_causal,338        )339 340        attn_output2 = _flash_attention_forward(341            query_states,342            key_states,343            value_states2,344            attention_mask,345            q_len,346            position_ids=position_ids,347            dropout=dropout_rate,348            sliding_window=getattr(self, "sliding_window", None),349            use_top_left_mask=self._flash_attn_uses_top_left_mask,350            is_causal=self.is_causal,351        )352 353        attn_output = torch.cat([attn_output1, attn_output2], dim=-1)354        attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=2)355 356        lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to(357            query_states.dtype358        )359        lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to(360            query_states.dtype361        )362        lambda_full = lambda_1 - lambda_2 + self.lambda_init363 364        attn_output = attn_output1 - lambda_full * attn_output2365        attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output)366        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()367        attn_output = self.o_proj(attn_output)368        return attn_output, None369 370 371class DiffLlamaSdpaAttention(DiffLlamaAttention):372    """373    DiffLlama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from374    `DiffLlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to375    SDPA API.376    """377 378    # Adapted from DiffLlamaAttention.forward379    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")380    def forward(381        self,382        hidden_states: torch.Tensor,383        position_embeddings: tuple[torch.Tensor, torch.Tensor],384        attention_mask: Optional[torch.Tensor] = None,385        position_ids: Optional[torch.LongTensor] = None,386        past_key_values: Optional[Cache] = None,387        use_cache: bool = False,388        cache_position: Optional[torch.LongTensor] = None,389        **kwargs,390    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:391        bsz, q_len, _ = hidden_states.size()392 393        query_states = self.q_proj(hidden_states)394        key_states = self.k_proj(hidden_states)395        value_states = self.v_proj(hidden_states)396 397        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)398        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)399        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)400 401        cos, sin = position_embeddings402        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)403 404        if past_key_values is not None:405            # sin and cos are specific to RoPE models; cache_position needed for the static cache406            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}407            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)408 409        key_states = repeat_kv(key_states, self.num_key_value_groups)410        value_states = repeat_kv(value_states, self.num_key_value_groups)411        value_states = torch.cat(torch.chunk(value_states, 2, dim=1), dim=-1)412        value_states = value_states.repeat(1, 2, 1, 1)413 414        causal_mask = attention_mask415        if attention_mask is not None:416            causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]417 418        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,419        # Reference: https://github.com/pytorch/pytorch/issues/112577.420        if query_states.device.type == "cuda" and causal_mask is not None:421            query_states = query_states.contiguous()422            key_states = key_states.contiguous()423            value_states = value_states.contiguous()424 425        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment426        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.427        is_causal = causal_mask is None and q_len > 1428 429        attn_output = torch.nn.functional.scaled_dot_product_attention(430            query_states,431            key_states,432            value_states,433            attn_mask=causal_mask,434            dropout_p=self.attention_dropout if self.training else 0.0,435            is_causal=is_causal,436        )437 438        attn_output1, attn_output2 = torch.chunk(attn_output, 2, dim=1)439 440        lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1, dtype=torch.float32)).to(441            query_states.dtype442        )443        lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1, dtype=torch.float32)).to(444            query_states.dtype445        )446        lambda_full = lambda_1 - lambda_2 + self.lambda_init447 448        attn_output = attn_output1 - lambda_full * attn_output2449        attn_output = (1 - self.lambda_init) * self.groupnorm(attn_output)450        attn_output = attn_output.transpose(1, 2).contiguous()451        attn_output = attn_output.view(bsz, q_len, -1)452        attn_output = self.o_proj(attn_output)453        return attn_output, None454 455 456@use_kernel_forward_from_hub("RMSNorm")457class DiffLlamaRMSNorm(nn.Module):458    def __init__(self, hidden_size, eps=1e-6):459        """460        DiffLlamaRMSNorm is equivalent to T5LayerNorm461        """462        super().__init__()463        self.weight = nn.Parameter(torch.ones(hidden_size))464        self.variance_epsilon = eps465 466    def forward(self, hidden_states):467        input_dtype = hidden_states.dtype468        hidden_states = hidden_states.to(torch.float32)469        variance = hidden_states.pow(2).mean(-1, keepdim=True)470        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)471        return self.weight * hidden_states.to(input_dtype)472 473    def extra_repr(self):474        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"475 476 477DIFFLLAMA_ATTENTION_CLASSES = {478    "eager": DiffLlamaAttention,479    "flash_attention_2": DiffLlamaFlashAttention2,480    "sdpa": DiffLlamaSdpaAttention,481}482 483 484class DiffLlamaDecoderLayer(GradientCheckpointingLayer):485    def __init__(self, config: DiffLlamaConfig, layer_idx: int):486        super().__init__()487        self.hidden_size = config.hidden_size488 489        self.self_attn = DIFFLLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)490 491        self.mlp = DiffLlamaMLP(config)492        self.input_layernorm = DiffLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)493        self.post_attention_layernorm = DiffLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)494 495    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")496    def forward(497        self,498        hidden_states: torch.Tensor,499        attention_mask: Optional[torch.Tensor] = None,500        position_ids: Optional[torch.LongTensor] = None,501        past_key_values: Optional[Cache] = None,502        use_cache: Optional[bool] = False,503        cache_position: Optional[torch.LongTensor] = None,504        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC505        **kwargs: Unpack[TransformersKwargs],506    ) -> torch.Tensor:507        residual = hidden_states508        hidden_states = self.input_layernorm(hidden_states)509        # Self Attention510        hidden_states, _ = self.self_attn(511            hidden_states=hidden_states,512            attention_mask=attention_mask,513            position_ids=position_ids,514            past_key_values=past_key_values,515            use_cache=use_cache,516            cache_position=cache_position,517            position_embeddings=position_embeddings,518            **kwargs,519        )520        hidden_states = residual + hidden_states521 522        # Fully Connected523        residual = hidden_states524        hidden_states = self.post_attention_layernorm(hidden_states)525        hidden_states = self.mlp(hidden_states)526        hidden_states = residual + hidden_states527        return hidden_states528 529 530@auto_docstring531class DiffLlamaPreTrainedModel(PreTrainedModel):532    config: DiffLlamaConfig533    base_model_prefix = "model"534    supports_gradient_checkpointing = True535    _no_split_modules = ["DiffLlamaDecoderLayer"]536    _skip_keys_device_placement = ["past_key_values"]537    _supports_flash_attn = True538    _supports_sdpa = True539    _supports_flex_attn = False540 541    _can_compile_fullgraph = True542    _supports_attention_backend = False543    _can_record_outputs = {544        "hidden_states": DiffLlamaDecoderLayer,545        "attentions": DiffLlamaAttention,546    }547 548    def _init_weights(self, module):549        super()._init_weights(module)550        if isinstance(module, DiffLlamaAttention):551            module.lambda_q1.data.normal_(0, self.config.lambda_std_dev)552            module.lambda_k1.data.normal_(0, self.config.lambda_std_dev)553            module.lambda_q2.data.normal_(0, self.config.lambda_std_dev)554            module.lambda_k2.data.normal_(0, self.config.lambda_std_dev)555 556 557class DiffLlamaRotaryEmbedding(nn.Module):558    inv_freq: torch.Tensor  # fix linting for `register_buffer`559 560    def __init__(self, config: DiffLlamaConfig, device=None):561        super().__init__()562        # BC: "rope_type" was originally "type"563        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):564            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))565        else:566            self.rope_type = "default"567        self.max_seq_len_cached = config.max_position_embeddings568        self.original_max_seq_len = config.max_position_embeddings569 570        self.config = config571        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]572 573        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)574        self.register_buffer("inv_freq", inv_freq, persistent=False)575        self.original_inv_freq = self.inv_freq576 577    @torch.no_grad()578    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)579    def forward(self, x, position_ids):580        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)581        position_ids_expanded = position_ids[:, None, :].float()582 583        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"584        with torch.autocast(device_type=device_type, enabled=False):  # Force float32585            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)586            emb = torch.cat((freqs, freqs), dim=-1)587            cos = emb.cos() * self.attention_scaling588            sin = emb.sin() * self.attention_scaling589 590        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)591 592 593@auto_docstring594class DiffLlamaModel(DiffLlamaPreTrainedModel):595    def __init__(self, config: DiffLlamaConfig):596        super().__init__(config)597        self.padding_idx = config.pad_token_id598        self.vocab_size = config.vocab_size599 600        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)601        self.layers = nn.ModuleList(602            [DiffLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]603        )604        self.norm = DiffLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)605        self.rotary_emb = DiffLlamaRotaryEmbedding(config=config)606        self.gradient_checkpointing = False607 608        # Initialize weights and apply final processing609        self.post_init()610 611    @check_model_inputs()612    @auto_docstring613    def forward(614        self,615        input_ids: Optional[torch.LongTensor] = None,616        attention_mask: Optional[torch.Tensor] = None,617        position_ids: Optional[torch.LongTensor] = None,618        past_key_values: Optional[Cache] = None,619        inputs_embeds: Optional[torch.FloatTensor] = None,620        cache_position: Optional[torch.LongTensor] = None,621        use_cache: Optional[bool] = None,622        **kwargs: Unpack[TransformersKwargs],623    ) -> BaseModelOutputWithPast:624        if (input_ids is None) ^ (inputs_embeds is not None):625            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")626 627        if inputs_embeds is None:628            inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)629 630        if use_cache and past_key_values is None:631            past_key_values = DynamicCache(config=self.config)632 633        if cache_position is None:634            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0635            cache_position: torch.Tensor = torch.arange(636                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device637            )638 639        if position_ids is None:640            position_ids = cache_position.unsqueeze(0)641 642        causal_mask = create_causal_mask(643            config=self.config,644            input_embeds=inputs_embeds,645            attention_mask=attention_mask,646            cache_position=cache_position,647            past_key_values=past_key_values,648            position_ids=position_ids,649        )650 651        hidden_states = inputs_embeds652        position_embeddings = self.rotary_emb(hidden_states, position_ids)653 654        for decoder_layer in self.layers[: self.config.num_hidden_layers]:655            hidden_states = decoder_layer(656                hidden_states,657                attention_mask=causal_mask,658                position_ids=position_ids,659                past_key_values=past_key_values,660                cache_position=cache_position,661                position_embeddings=position_embeddings,662                **kwargs,663            )664 665        hidden_states = self.norm(hidden_states)666        return BaseModelOutputWithPast(667            last_hidden_state=hidden_states,668            past_key_values=past_key_values,669        )670 671 672@auto_docstring673class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin):674    _tied_weights_keys = ["lm_head.weight"]675    _tp_plan = {"lm_head": "colwise_rep"}676    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}677 678    def __init__(self, config):679        super().__init__(config)680        self.model = DiffLlamaModel(config)681        self.vocab_size = config.vocab_size682        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)683 684        # Initialize weights and apply final processing685        self.post_init()686 687    @can_return_tuple688    @auto_docstring689    def forward(690        self,691        input_ids: Optional[torch.LongTensor] = None,692        attention_mask: Optional[torch.Tensor] = None,693        position_ids: Optional[torch.LongTensor] = None,694        past_key_values: Optional[Cache] = None,695        inputs_embeds: Optional[torch.FloatTensor] = None,696        labels: Optional[torch.LongTensor] = None,697        use_cache: Optional[bool] = None,698        cache_position: Optional[torch.LongTensor] = None,699        logits_to_keep: Union[int, torch.Tensor] = 0,700        **kwargs: Unpack[TransformersKwargs],701    ) -> CausalLMOutputWithPast:702        r"""703        Example:704 705        ```python706        >>> from transformers import AutoTokenizer, DiffLlamaForCausalLM707 708        >>> model = DiffLlamaForCausalLM.from_pretrained("google/diffllama-7b")709        >>> tokenizer = AutoTokenizer.from_pretrained("google/diffllama-7b")710 711        >>> prompt = "What is your favorite condiment?"712        >>> inputs = tokenizer(prompt, return_tensors="pt")713 714        >>> # Generate715        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)716        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]717        "What is your favorite condiment?"718        ```"""719        outputs: BaseModelOutputWithPast = self.model(720            input_ids=input_ids,721            attention_mask=attention_mask,722            position_ids=position_ids,723            past_key_values=past_key_values,724            inputs_embeds=inputs_embeds,725            use_cache=use_cache,726            cache_position=cache_position,727            **kwargs,728        )729 730        hidden_states = outputs.last_hidden_state731        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss732        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep733        logits = self.lm_head(hidden_states[:, slice_indices, :])734 735        loss = None736        if labels is not None:737            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)738 739        return CausalLMOutputWithPast(740            loss=loss,741            logits=logits,742            past_key_values=outputs.past_key_values,743            hidden_states=outputs.hidden_states,744            attentions=outputs.attentions,745        )746 747 748class DiffLlamaForSequenceClassification(GenericForSequenceClassification, DiffLlamaPreTrainedModel):749    pass750 751 752class DiffLlamaForQuestionAnswering(GenericForQuestionAnswering, DiffLlamaPreTrainedModel):753    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`754 755 756class DiffLlamaForTokenClassification(GenericForTokenClassification, DiffLlamaPreTrainedModel):757    pass758 759 760__all__ = [761    "DiffLlamaPreTrainedModel",762    "DiffLlamaModel",763    "DiffLlamaForCausalLM",764    "DiffLlamaForSequenceClassification",765    "DiffLlamaForQuestionAnswering",766    "DiffLlamaForTokenClassification",767]768