rpDungeon/Loopstral-4B-Experimental

1#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2#           This file was automatically generated from src/transformers/models/mistral/modular_mistral.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_mistral.py file directly. One of our CI enforces this.6#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7import copy8from typing import Callable, Optional, Union9 10import torch11from torch import nn12from torch.nn import CrossEntropyLoss13 14#from transformers.modeling_utils import check_model_inputs15 16from transformers.activations import ACT2FN17from transformers.cache_utils import Cache, DynamicCache, DynamicLayer18from transformers.generation import GenerationMixin19from transformers.integrations import use_kernel_forward_from_hub20from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask21from transformers.modeling_flash_attention_utils import FlashAttentionKwargs22from transformers.modeling_layers import (23    GenericForQuestionAnswering,24    GenericForSequenceClassification,25    GenericForTokenClassification,26    GradientCheckpointingLayer,27)28from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast29from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update30from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel31from transformers.processing_utils import Unpack32from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple33#from transformers.utils.doc import auto_docstring34from transformers.utils.deprecation import deprecate_kwarg35from .configuration_loopstral import LoopstralConfig36 37 38class MistralMLP(nn.Module):39    def __init__(self, config):40        super().__init__()41        self.config = config42        self.hidden_size = config.hidden_size43        self.intermediate_size = config.intermediate_size44        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)45        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)46        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)47        self.act_fn = ACT2FN[config.hidden_act]48 49    def forward(self, x):50        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))51        return down_proj52 53 54def rotate_half(x):55    """Rotates half the hidden dims of the input."""56    x1 = x[..., : x.shape[-1] // 2]57    x2 = x[..., x.shape[-1] // 2 :]58    return torch.cat((-x2, x1), dim=-1)59 60 61def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):62    """Applies Rotary Position Embedding to the query and key tensors.63 64    Args:65        q (`torch.Tensor`): The query tensor.66        k (`torch.Tensor`): The key tensor.67        cos (`torch.Tensor`): The cosine part of the rotary embedding.68        sin (`torch.Tensor`): The sine part of the rotary embedding.69        position_ids (`torch.Tensor`, *optional*):70            Deprecated and unused.71        unsqueeze_dim (`int`, *optional*, defaults to 1):72            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and73            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note74            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and75            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes76            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have77            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.78    Returns:79        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.80    """81    cos = cos.unsqueeze(unsqueeze_dim)82    sin = sin.unsqueeze(unsqueeze_dim)83    q_embed = (q * cos) + (rotate_half(q) * sin)84    k_embed = (k * cos) + (rotate_half(k) * sin)85    return q_embed, k_embed86 87 88def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:89    """90    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,91    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)92    """93    batch, num_key_value_heads, slen, head_dim = hidden_states.shape94    if n_rep == 1:95        return hidden_states96    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)97    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)98 99 100def eager_attention_forward(101    module: nn.Module,102    query: torch.Tensor,103    key: torch.Tensor,104    value: torch.Tensor,105    attention_mask: Optional[torch.Tensor],106    scaling: float,107    dropout: float = 0.0,108    **kwargs: Unpack[TransformersKwargs],109):110    key_states = repeat_kv(key, module.num_key_value_groups)111    value_states = repeat_kv(value, module.num_key_value_groups)112 113    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling114    if attention_mask is not None:115        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]116        attn_weights = attn_weights + causal_mask117 118    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)119    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)120    attn_output = torch.matmul(attn_weights, value_states)121    attn_output = attn_output.transpose(1, 2).contiguous()122 123    return attn_output, attn_weights124 125 126class MistralAttention(nn.Module):127    """Multi-headed attention from 'Attention Is All You Need' paper"""128 129    def __init__(self, config: LoopstralConfig, layer_idx: int):130        super().__init__()131        self.config = config132        self.layer_idx = layer_idx133        self.head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads134        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads135        self.scaling = self.head_dim**-0.5136        self.attention_dropout = config.attention_dropout137        self.is_causal = True138        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)139        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)140        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)141        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)142 143    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")144    def forward(145        self,146        hidden_states: torch.Tensor,147        position_embeddings: tuple[torch.Tensor, torch.Tensor],148        attention_mask: Optional[torch.Tensor],149        past_key_values: Optional[Cache] = None,150        cache_position: Optional[torch.LongTensor] = None,151        cache_slot_idx: Optional[int] = None,152        **kwargs: Unpack[FlashAttentionKwargs],153    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:154        input_shape = hidden_states.shape[:-1]155        hidden_shape = (*input_shape, -1, self.head_dim)156 157        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)158        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)159        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)160 161        cos, sin = position_embeddings162        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)163 164        if past_key_values is not None:165            # sin and cos are specific to RoPE models; cache_position needed for the static cache166            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}167            # Use cache_slot_idx (position in layer sequence) instead of layer_idx168            # This allows each visit to a repeated layer to have its own cache slot169            slot_idx = cache_slot_idx if cache_slot_idx is not None else self.layer_idx170            key_states, value_states = past_key_values.update(key_states, value_states, slot_idx, cache_kwargs)171 172        attention_interface: Callable = eager_attention_forward173        if self.config._attn_implementation != "eager":174            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]175 176        attn_output, attn_weights = attention_interface(177            self,178            query_states,179            key_states,180            value_states,181            attention_mask,182            dropout=0.0 if not self.training else self.attention_dropout,183            scaling=self.scaling,184            sliding_window=getattr(self.config, "sliding_window", None),  # main diff with Llama185            **kwargs,186        )187 188        attn_output = attn_output.reshape(*input_shape, -1).contiguous()189        attn_output = self.o_proj(attn_output)190        return attn_output, attn_weights191 192 193@use_kernel_forward_from_hub("RMSNorm")194class MistralRMSNorm(nn.Module):195    def __init__(self, hidden_size, eps=1e-6):196        """197        MistralRMSNorm is equivalent to T5LayerNorm198        """199        super().__init__()200        self.weight = nn.Parameter(torch.ones(hidden_size))201        self.variance_epsilon = eps202 203    def forward(self, hidden_states):204        input_dtype = hidden_states.dtype205        hidden_states = hidden_states.to(torch.float32)206        variance = hidden_states.pow(2).mean(-1, keepdim=True)207        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)208        return self.weight * hidden_states.to(input_dtype)209 210    def extra_repr(self):211        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"212 213 214class MistralDecoderLayer(GradientCheckpointingLayer):215    def __init__(self, config: LoopstralConfig, layer_idx: int):216        super().__init__()217        self.hidden_size = config.hidden_size218        self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)219        self.mlp = MistralMLP(config)220        self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)221        self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)222 223    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")224    def forward(225        self,226        hidden_states: torch.Tensor,227        attention_mask: Optional[torch.Tensor] = None,228        position_ids: Optional[torch.LongTensor] = None,229        past_key_values: Optional[Cache] = None,230        use_cache: Optional[bool] = False,231        cache_position: Optional[torch.LongTensor] = None,232        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC233        cache_slot_idx: Optional[int] = None,234        **kwargs: Unpack[TransformersKwargs],235    ) -> torch.Tensor:236        residual = hidden_states237        hidden_states = self.input_layernorm(hidden_states)238        # Self Attention239        hidden_states, _ = self.self_attn(240            hidden_states=hidden_states,241            attention_mask=attention_mask,242            position_ids=position_ids,243            past_key_values=past_key_values,244            use_cache=use_cache,245            cache_position=cache_position,246            position_embeddings=position_embeddings,247            cache_slot_idx=cache_slot_idx,248            **kwargs,249        )250        hidden_states = residual + hidden_states251 252        # Fully Connected253        residual = hidden_states254        hidden_states = self.post_attention_layernorm(hidden_states)255        hidden_states = self.mlp(hidden_states)256        hidden_states = residual + hidden_states257        return hidden_states258 259 260@auto_docstring261class MistralPreTrainedModel(PreTrainedModel):262    config: LoopstralConfig263    base_model_prefix = "model"264    supports_gradient_checkpointing = True265    _no_split_modules = ["MistralDecoderLayer"]266    _skip_keys_device_placement = ["past_key_values"]267    _supports_flash_attn = True268    _supports_sdpa = True269    _supports_flex_attn = True270 271    _can_compile_fullgraph = True272    _supports_attention_backend = True273    _can_record_outputs = {274        "hidden_states": MistralDecoderLayer,275        "attentions": MistralAttention,276    }277 278 279class MistralRotaryEmbedding(nn.Module):280    inv_freq: torch.Tensor  # fix linting for `register_buffer`281 282    def __init__(self, config: LoopstralConfig, device=None):283        super().__init__()284        # BC: "rope_type" was originally "type"285        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):286            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))287        else:288            self.rope_type = "default"289        self.max_seq_len_cached = config.max_position_embeddings290        self.original_max_seq_len = config.max_position_embeddings291 292        self.config = config293        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]294 295        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)296        self.register_buffer("inv_freq", inv_freq, persistent=False)297        self.original_inv_freq = self.inv_freq298 299    @torch.no_grad()300    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)301    def forward(self, x, position_ids):302        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)303        position_ids_expanded = position_ids[:, None, :].float()304 305        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"306        with torch.autocast(device_type=device_type, enabled=False):  # Force float32307            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)308            emb = torch.cat((freqs, freqs), dim=-1)309            cos = emb.cos() * self.attention_scaling310            sin = emb.sin() * self.attention_scaling311 312        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)313 314 315def _expand_layer_sequence(layer_sequence, num_hidden_layers):316    """Expand layer_sequence config into a flat list of layer indices."""317    l_seq = []318    for item in layer_sequence:319        if isinstance(item, int):320            # Single layer index: 5 -> [5]321            l_seq.append(item)322        elif isinstance(item, list):323            if len(item) == 2:324                # Range without repeat: [4, 20] -> range(4, 20)325                start, end = item326                l_seq += list(range(start, min(end, num_hidden_layers)))327            elif len(item) == 3:328                # Range with repeat: [4, 20, 2] -> range(4, 20) repeated 2 times329                start, end, repeats = item330                l_seq += list(range(start, min(end, num_hidden_layers))) * repeats331            else:332                raise ValueError(f"Invalid layer_sequence item: {item}. Expected int, [start, end], or [start, end, repeats]")333        else:334            raise ValueError(f"Invalid layer_sequence item type: {type(item)}. Expected int or list.")335    return l_seq336 337 338@auto_docstring339class LoopstralModel(MistralPreTrainedModel):340    def __init__(self, config: LoopstralConfig):341        super().__init__(config)342        self.padding_idx = config.pad_token_id343        self.vocab_size = config.vocab_size344 345        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)346        self.layers = nn.ModuleList(347            [MistralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]348        )349        self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)350        self.rotary_emb = MistralRotaryEmbedding(config=config)351        self.gradient_checkpointing = False352 353        # Pre-compute the expanded layer sequence for the looping mechanism354        self._layer_sequence = _expand_layer_sequence(config.layer_sequence, config.num_hidden_layers)355        # Number of cache slots needed (one per position in layer sequence)356        self._num_cache_slots = len(self._layer_sequence)357 358        # Initialize weights and apply final processing359        self.post_init()360 361    #@check_model_inputs362    @auto_docstring363    def forward(364        self,365        input_ids: Optional[torch.LongTensor] = None,366        attention_mask: Optional[torch.Tensor] = None,367        position_ids: Optional[torch.LongTensor] = None,368        past_key_values: Optional[Cache] = None,369        inputs_embeds: Optional[torch.FloatTensor] = None,370        use_cache: Optional[bool] = None,371        cache_position: Optional[torch.LongTensor] = None,372        **kwargs: Unpack[TransformersKwargs],373    ) -> BaseModelOutputWithPast:374        if (input_ids is None) ^ (inputs_embeds is not None):375            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")376 377        if inputs_embeds is None:378            inputs_embeds = self.embed_tokens(input_ids)379 380        if use_cache:381            if past_key_values is None:382                # Create cache with enough slots for the full layer sequence383                # (more than num_hidden_layers if layers are repeated)384                cache_config = copy.copy(self.config)385                cache_config.num_hidden_layers = self._num_cache_slots386                past_key_values = DynamicCache(config=cache_config)387            elif isinstance(past_key_values, DynamicCache) and len(past_key_values.layers) < self._num_cache_slots:388                # Cache was created externally (e.g., by generate()) with fewer slots389                # Extend it to have enough slots for our layer sequence390                while len(past_key_values.layers) < self._num_cache_slots:391                    past_key_values.layers.append(DynamicLayer())392 393        if cache_position is None:394            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0395            cache_position = torch.arange(396                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device397            )398 399        if position_ids is None:400            position_ids = cache_position.unsqueeze(0)401 402        mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask403        causal_mask = mask_function(404            config=self.config,405            input_embeds=inputs_embeds,406            attention_mask=attention_mask,407            cache_position=cache_position,408            past_key_values=past_key_values,409            position_ids=position_ids,410        )411 412        hidden_states = inputs_embeds413        position_embeddings = self.rotary_emb(hidden_states, position_ids)414 415        # Execute layers in the configured sequence416        # Each position in the sequence gets its own cache slot, allowing417        # repeated layers to maintain separate KV caches for each visit418        for cache_slot_idx, layer_idx in enumerate(self._layer_sequence):419            decoder_layer = self.layers[layer_idx]420            hidden_states = decoder_layer(421                hidden_states,422                attention_mask=causal_mask,423                position_ids=position_ids,424                past_key_values=past_key_values,425                use_cache=use_cache,426                cache_position=cache_position,427                position_embeddings=position_embeddings,428                cache_slot_idx=cache_slot_idx,429                **kwargs,430            )431        hidden_states = self.norm(hidden_states)432        return BaseModelOutputWithPast(433            last_hidden_state=hidden_states,434            past_key_values=past_key_values if use_cache else None,435        )436 437 438@auto_docstring439class LoopstralForCausalLM(MistralPreTrainedModel, GenerationMixin):440    _tied_weights_keys = ["lm_head.weight"]441    _tp_plan = {"lm_head": "colwise_rep"}442    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}443 444    def __init__(self, config):445        super().__init__(config)446        self.model = LoopstralModel(config)447        self.vocab_size = config.vocab_size448        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)449 450        # Initialize weights and apply final processing451        self.post_init()452 453    @can_return_tuple454    @auto_docstring455    def forward(456        self,457        input_ids: Optional[torch.LongTensor] = None,458        attention_mask: Optional[torch.Tensor] = None,459        position_ids: Optional[torch.LongTensor] = None,460        past_key_values: Optional[Cache] = None,461        inputs_embeds: Optional[torch.FloatTensor] = None,462        labels: Optional[torch.LongTensor] = None,463        use_cache: Optional[bool] = None,464        cache_position: Optional[torch.LongTensor] = None,465        logits_to_keep: Union[int, torch.Tensor] = 0,466        **kwargs: Unpack[TransformersKwargs],467    ) -> CausalLMOutputWithPast:468        r"""469        Example:470 471        ```python472        >>> from transformers import AutoTokenizer, MistralForCausalLM473 474        >>> model = MistralForCausalLM.from_pretrained("meta-mistral/Mistral-2-7b-hf")475        >>> tokenizer = AutoTokenizer.from_pretrained("meta-mistral/Mistral-2-7b-hf")476 477        >>> prompt = "Hey, are you conscious? Can you talk to me?"478        >>> inputs = tokenizer(prompt, return_tensors="pt")479 480        >>> # Generate481        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)482        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]483        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."484        ```"""485        outputs: BaseModelOutputWithPast = self.model(486            input_ids=input_ids,487            attention_mask=attention_mask,488            position_ids=position_ids,489            past_key_values=past_key_values,490            inputs_embeds=inputs_embeds,491            use_cache=use_cache,492            cache_position=cache_position,493            **kwargs,494        )495 496        hidden_states = outputs.last_hidden_state497        logits = self.lm_head(hidden_states)498 499        loss = None500        if labels is not None:501            # THE FIX IS HERE: Standard loss calculation502            # Shift so that tokens < n predict n503            shift_logits = logits[..., :-1, :].contiguous()504            shift_labels = labels[..., 1:].contiguous()505            # Flatten the tokens506            loss_fct = CrossEntropyLoss()507            shift_logits = shift_logits.view(-1, self.config.vocab_size)508            shift_labels = shift_labels.view(-1)509            # Enable model parallelism510            shift_labels = shift_labels.to(shift_logits.device)511            loss = loss_fct(shift_logits, shift_labels)512 513        return CausalLMOutputWithPast(514            loss=loss,515            logits=logits,516            past_key_values=outputs.past_key_values,517            hidden_states=outputs.hidden_states,518            attentions=outputs.attentions,519        )520 521 522class MistralForTokenClassification(GenericForTokenClassification, MistralPreTrainedModel):523    pass524 525 526class MistralForSequenceClassification(GenericForSequenceClassification, MistralPreTrainedModel):527    pass528 529 530class MistralForQuestionAnswering(GenericForQuestionAnswering, MistralPreTrainedModel): ...531 532 533__all__ = [534    "LoopstralForCausalLM",535    "MistralForQuestionAnswering",536    "LoopstralModel",537    "MistralPreTrainedModel",538    "MistralForSequenceClassification",539    "MistralForTokenClassification",540]541