francislabounty commited on
Commit
bdf10ed
1 Parent(s): dff7fd9

Create modeling_sparsetral.py

Browse files
Files changed (1) hide show
  1. modeling_sparsetral.py +1622 -0
modeling_sparsetral.py ADDED
@@ -0,0 +1,1622 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch Mistral model."""
21
+ import inspect
22
+ import math
23
+ import warnings
24
+ from dataclasses import dataclass
25
+ from typing import List, Optional, Tuple, Union
26
+
27
+ import torch
28
+ import torch.nn.functional as F
29
+ import torch.utils.checkpoint
30
+ from torch import nn
31
+ from torch.nn import CrossEntropyLoss
32
+
33
+ from transformers.activations import ACT2FN
34
+ from transformers.cache_utils import Cache, DynamicCache
35
+ from transformers.modeling_attn_mask_utils import (
36
+ _prepare_4d_causal_attention_mask,
37
+ _prepare_4d_causal_attention_mask_for_sdpa,
38
+ )
39
+ from transformers.modeling_outputs import (
40
+ BaseModelOutputWithPast,
41
+ CausalLMOutputWithPast,
42
+ SequenceClassifierOutputWithPast,
43
+ )
44
+ from transformers.modeling_utils import PreTrainedModel
45
+ from transformers.utils import (
46
+ add_start_docstrings,
47
+ add_start_docstrings_to_model_forward,
48
+ is_flash_attn_2_available,
49
+ is_flash_attn_greater_or_equal_2_10,
50
+ logging,
51
+ replace_return_docstrings,
52
+ ModelOutput,
53
+ )
54
+ from .configuration_sparsetral import SparsetralConfig
55
+
56
+
57
+ if is_flash_attn_2_available():
58
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
59
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
60
+
61
+ _flash_supports_window_size = "window_size" in list(
62
+ inspect.signature(flash_attn_func).parameters
63
+ )
64
+
65
+
66
+ logger = logging.get_logger(__name__)
67
+
68
+ _CONFIG_FOR_DOC = "SparsetralConfig"
69
+
70
+
71
+ @dataclass
72
+ class MoEModelOutputWithPast(ModelOutput):
73
+ last_hidden_state: torch.FloatTensor = None
74
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
75
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
76
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
77
+ router_logits: Optional[Tuple[torch.FloatTensor]] = None
78
+
79
+
80
+ @dataclass
81
+ class MoECausalLMOutputWithPast(ModelOutput):
82
+ loss: Optional[torch.FloatTensor] = None
83
+ aux_loss: Optional[torch.FloatTensor] = None
84
+ logits: torch.FloatTensor = None
85
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
86
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
87
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
88
+ router_logits: Optional[Tuple[torch.FloatTensor]] = None
89
+
90
+
91
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
92
+ def _get_unpad_data(attention_mask):
93
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
94
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
95
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
96
+ cu_seqlens = F.pad(
97
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
98
+ )
99
+ return (
100
+ indices,
101
+ cu_seqlens,
102
+ max_seqlen_in_batch,
103
+ )
104
+
105
+
106
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral
107
+ class MistralRMSNorm(nn.Module):
108
+ def __init__(self, hidden_size, eps=1e-6):
109
+ """
110
+ MistralRMSNorm is equivalent to T5LayerNorm
111
+ """
112
+ super().__init__()
113
+ self.weight = nn.Parameter(torch.ones(hidden_size))
114
+ self.variance_epsilon = eps
115
+
116
+ def forward(self, hidden_states):
117
+ input_dtype = hidden_states.dtype
118
+ hidden_states = hidden_states.to(torch.float32)
119
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
120
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
121
+ return self.weight * hidden_states.to(input_dtype)
122
+
123
+
124
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Mistral
125
+ class MistralRotaryEmbedding(nn.Module):
126
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
127
+ super().__init__()
128
+
129
+ self.dim = dim
130
+ self.max_position_embeddings = max_position_embeddings
131
+ self.base = base
132
+ inv_freq = 1.0 / (
133
+ self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
134
+ )
135
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
136
+
137
+ # Build here to make `torch.jit.trace` work.
138
+ self._set_cos_sin_cache(
139
+ seq_len=max_position_embeddings,
140
+ device=self.inv_freq.device,
141
+ dtype=torch.get_default_dtype(),
142
+ )
143
+
144
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
145
+ self.max_seq_len_cached = seq_len
146
+ t = torch.arange(
147
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
148
+ )
149
+
150
+ freqs = torch.outer(t, self.inv_freq)
151
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
152
+ emb = torch.cat((freqs, freqs), dim=-1)
153
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
154
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
155
+
156
+ def forward(self, x, seq_len=None):
157
+ # x: [bs, num_attention_heads, seq_len, head_size]
158
+ if seq_len > self.max_seq_len_cached:
159
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
160
+
161
+ return (
162
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
163
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
164
+ )
165
+
166
+
167
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
168
+ def rotate_half(x):
169
+ """Rotates half the hidden dims of the input."""
170
+ x1 = x[..., : x.shape[-1] // 2]
171
+ x2 = x[..., x.shape[-1] // 2 :]
172
+ return torch.cat((-x2, x1), dim=-1)
173
+
174
+
175
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
176
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
177
+ """Applies Rotary Position Embedding to the query and key tensors.
178
+ Args:
179
+ q (`torch.Tensor`): The query tensor.
180
+ k (`torch.Tensor`): The key tensor.
181
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
182
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
183
+ position_ids (`torch.Tensor`):
184
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
185
+ used to pass offsetted position ids when working with a KV-cache.
186
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
187
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
188
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
189
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
190
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
191
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
192
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
193
+ Returns:
194
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
195
+ """
196
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
197
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
198
+ q_embed = (q * cos) + (rotate_half(q) * sin)
199
+ k_embed = (k * cos) + (rotate_half(k) * sin)
200
+ return q_embed, k_embed
201
+
202
+
203
+ # Mistral MoE
204
+ def load_balancing_loss_func(
205
+ gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2
206
+ ) -> float:
207
+ r"""
208
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
209
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
210
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
211
+ experts is too unbalanced.
212
+ Args:
213
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
214
+ Logits from the `gate`, should be a tuple of tensors. Shape: [batch_size, seqeunce_length, num_experts].
215
+ num_experts (`int`, *optional*):
216
+ Number of experts
217
+ Returns:
218
+ The auxiliary loss.
219
+ """
220
+ if gate_logits is None:
221
+ return 0
222
+
223
+ if isinstance(gate_logits, tuple):
224
+ # cat along the layers?
225
+ compute_device = gate_logits[0].device
226
+ gate_logits = torch.cat(
227
+ [gate.to(compute_device) for gate in gate_logits], dim=0
228
+ )
229
+
230
+ routing_weights, selected_experts = torch.topk(gate_logits, top_k, dim=-1)
231
+ routing_weights = routing_weights.softmax(dim=-1)
232
+
233
+ # cast the expert indices to int64, otherwise one-hot encoding will fail
234
+ if selected_experts.dtype != torch.int64:
235
+ selected_experts = selected_experts.to(torch.int64)
236
+
237
+ if len(selected_experts.shape) == 2:
238
+ selected_experts = selected_experts.unsqueeze(2)
239
+
240
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
241
+
242
+ # For a given token, determine if it was routed to a given expert.
243
+ expert_mask = torch.max(expert_mask, axis=-2).values
244
+
245
+ # cast to float32 otherwise mean will fail
246
+ expert_mask = expert_mask.to(torch.float32)
247
+ tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
248
+
249
+ router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-1)
250
+ return torch.mean(
251
+ tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)
252
+ ) * (num_experts**2)
253
+
254
+
255
+ class ParallelAdapterMLP(nn.Module):
256
+ def __init__(self, config, adapter_dim, adapter_scaling):
257
+ super().__init__()
258
+ self.config = config
259
+ self.intermediate_size = config.intermediate_size
260
+ self.hidden_size = config.hidden_size
261
+ self.adapter_down = nn.Linear(self.hidden_size, adapter_dim, bias=False)
262
+ self.adapter_up = nn.Linear(adapter_dim, self.hidden_size, bias=False)
263
+ self.adapter_act = nn.GELU()
264
+
265
+ self.adapter_dropout = nn.Dropout(p=config.adapter_dropout)
266
+ self.adapter_scaling = adapter_scaling
267
+
268
+ def forward(self, x):
269
+ x = self.adapter_dropout(x)
270
+ x = self.adapter_scaling * self.adapter_up(
271
+ self.adapter_act(self.adapter_down(x))
272
+ )
273
+ return x
274
+
275
+
276
+ class SparsetralGateAdapter(nn.Module):
277
+ def __init__(self, config: SparsetralConfig):
278
+ super().__init__()
279
+
280
+ self.intermediate_size = config.intermediate_size
281
+ self.hidden_size = config.hidden_size
282
+
283
+ # Step 1: Router
284
+ self.num_experts = config.num_experts
285
+ self.topk = config.topk
286
+ self.router = nn.Linear(config.hidden_size, self.num_experts, bias=False)
287
+ self.dtype = getattr(torch, config.moe_dtype)
288
+
289
+ # Step 2: Get the experts
290
+ self.expert_indicies = list(range(self.num_experts))
291
+ self.experts = nn.ModuleList(
292
+ [
293
+ ParallelAdapterMLP(config, config.adapter_dim, config.moe_scaling)
294
+ for _ in self.expert_indicies
295
+ ]
296
+ )
297
+
298
+ def forward(self, input_hidden_states, output_hidden_states, router_hidden_states):
299
+ orig_shape = output_hidden_states.shape
300
+ input_hidden_states = input_hidden_states.view(
301
+ -1, input_hidden_states.shape[-1]
302
+ )
303
+ output_hidden_states = output_hidden_states.view(
304
+ -1, output_hidden_states.shape[-1]
305
+ )
306
+ router_hidden_states = router_hidden_states.view(
307
+ -1, router_hidden_states.shape[-1]
308
+ )
309
+
310
+ router_logits = self.router(router_hidden_states)
311
+
312
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
313
+ routing_weights, selected_experts = torch.topk(
314
+ routing_weights, self.topk, dim=-1
315
+ )
316
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
317
+
318
+ final_hidden_states = None
319
+ for expert_idx in self.expert_indicies:
320
+ expert_layer = self.experts[expert_idx]
321
+ expert_mask = selected_experts == expert_idx
322
+ expert_weights = (routing_weights * expert_mask).sum(dim=-1, keepdim=True)
323
+
324
+ current_hidden_states = (
325
+ expert_layer(input_hidden_states)
326
+ .add_(output_hidden_states)
327
+ .mul_(expert_weights)
328
+ )
329
+ if final_hidden_states is None:
330
+ final_hidden_states = current_hidden_states
331
+ else:
332
+ final_hidden_states.add_(current_hidden_states)
333
+
334
+ return final_hidden_states.view(*orig_shape), router_logits
335
+
336
+
337
+ class MistralMLP(nn.Module):
338
+ def __init__(self, config):
339
+ super().__init__()
340
+ self.config = config
341
+ self.hidden_size = config.hidden_size
342
+ self.intermediate_size = config.intermediate_size
343
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
344
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
345
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
346
+ self.act_fn = ACT2FN[config.hidden_act]
347
+
348
+ self.moe_adapter = SparsetralGateAdapter(config)
349
+
350
+ def forward(self, x):
351
+ router_hidden_states = x
352
+ up_proj = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
353
+ down_proj = self.down_proj(up_proj)
354
+ down_proj, router_logits = self.moe_adapter(
355
+ down_proj, down_proj, router_hidden_states
356
+ )
357
+
358
+ return down_proj, router_logits
359
+
360
+
361
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
362
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
363
+ """
364
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
365
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
366
+ """
367
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
368
+ if n_rep == 1:
369
+ return hidden_states
370
+ hidden_states = hidden_states[:, :, None, :, :].expand(
371
+ batch, num_key_value_heads, n_rep, slen, head_dim
372
+ )
373
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
374
+
375
+
376
+ class MistralAttention(nn.Module):
377
+ """
378
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
379
+ and "Generating Long Sequences with Sparse Transformers".
380
+ """
381
+
382
+ def __init__(self, config: SparsetralConfig, layer_idx: Optional[int] = None):
383
+ super().__init__()
384
+ self.config = config
385
+ self.layer_idx = layer_idx
386
+ if layer_idx is None:
387
+ logger.warning_once(
388
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
389
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
390
+ "when creating this class."
391
+ )
392
+
393
+ self.hidden_size = config.hidden_size
394
+ self.num_heads = config.num_attention_heads
395
+ self.head_dim = self.hidden_size // self.num_heads
396
+ self.num_key_value_heads = config.num_key_value_heads
397
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
398
+ self.max_position_embeddings = config.max_position_embeddings
399
+ self.rope_theta = config.rope_theta
400
+ self.is_causal = True
401
+ self.attention_dropout = config.attention_dropout
402
+
403
+ if (self.head_dim * self.num_heads) != self.hidden_size:
404
+ raise ValueError(
405
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
406
+ f" and `num_heads`: {self.num_heads})."
407
+ )
408
+ self.q_proj = nn.Linear(
409
+ self.hidden_size, self.num_heads * self.head_dim, bias=False
410
+ )
411
+ self.k_proj = nn.Linear(
412
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
413
+ )
414
+ self.v_proj = nn.Linear(
415
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
416
+ )
417
+ self.o_proj = nn.Linear(
418
+ self.num_heads * self.head_dim, self.hidden_size, bias=False
419
+ )
420
+
421
+ self.rotary_emb = MistralRotaryEmbedding(
422
+ self.head_dim,
423
+ max_position_embeddings=self.max_position_embeddings,
424
+ base=self.rope_theta,
425
+ )
426
+
427
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
428
+ return (
429
+ tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
430
+ .transpose(1, 2)
431
+ .contiguous()
432
+ )
433
+
434
+ def forward(
435
+ self,
436
+ hidden_states: torch.Tensor,
437
+ attention_mask: Optional[torch.Tensor] = None,
438
+ position_ids: Optional[torch.LongTensor] = None,
439
+ past_key_value: Optional[Cache] = None,
440
+ output_attentions: bool = False,
441
+ use_cache: bool = False,
442
+ **kwargs,
443
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
444
+ if "padding_mask" in kwargs:
445
+ warnings.warn(
446
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
447
+ )
448
+ bsz, q_len, _ = hidden_states.size()
449
+
450
+ query_states = self.q_proj(hidden_states)
451
+ key_states = self.k_proj(hidden_states)
452
+ value_states = self.v_proj(hidden_states)
453
+
454
+ query_states = query_states.view(
455
+ bsz, q_len, self.num_heads, self.head_dim
456
+ ).transpose(1, 2)
457
+ key_states = key_states.view(
458
+ bsz, q_len, self.num_key_value_heads, self.head_dim
459
+ ).transpose(1, 2)
460
+ value_states = value_states.view(
461
+ bsz, q_len, self.num_key_value_heads, self.head_dim
462
+ ).transpose(1, 2)
463
+
464
+ kv_seq_len = key_states.shape[-2]
465
+ if past_key_value is not None:
466
+ if self.layer_idx is None:
467
+ raise ValueError(
468
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
469
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
470
+ "with a layer index."
471
+ )
472
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
473
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
474
+ query_states, key_states = apply_rotary_pos_emb(
475
+ query_states, key_states, cos, sin, position_ids
476
+ )
477
+
478
+ if past_key_value is not None:
479
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
480
+ key_states, value_states = past_key_value.update(
481
+ key_states, value_states, self.layer_idx, cache_kwargs
482
+ )
483
+
484
+ # repeat k/v heads if n_kv_heads < n_heads
485
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
486
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
487
+
488
+ attn_weights = torch.matmul(
489
+ query_states, key_states.transpose(2, 3)
490
+ ) / math.sqrt(self.head_dim)
491
+
492
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
493
+ raise ValueError(
494
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
495
+ f" {attn_weights.size()}"
496
+ )
497
+
498
+ if attention_mask is not None:
499
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
500
+ raise ValueError(
501
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
502
+ )
503
+
504
+ attn_weights = attn_weights + attention_mask
505
+
506
+ # upcast attention to fp32
507
+ attn_weights = nn.functional.softmax(
508
+ attn_weights, dim=-1, dtype=torch.float32
509
+ ).to(query_states.dtype)
510
+ attn_weights = nn.functional.dropout(
511
+ attn_weights, p=self.attention_dropout, training=self.training
512
+ )
513
+ attn_output = torch.matmul(attn_weights, value_states)
514
+
515
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
516
+ raise ValueError(
517
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
518
+ f" {attn_output.size()}"
519
+ )
520
+
521
+ attn_output = attn_output.transpose(1, 2).contiguous()
522
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
523
+
524
+ attn_output = self.o_proj(attn_output)
525
+
526
+ if not output_attentions:
527
+ attn_weights = None
528
+
529
+ return attn_output, attn_weights, past_key_value
530
+
531
+
532
+ class MistralFlashAttention2(MistralAttention):
533
+ """
534
+ Mistral flash attention module. This module inherits from `MistralAttention` as the weights of the module stays
535
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
536
+ flash attention and deal with padding tokens in case the input contains any of them.
537
+ """
538
+
539
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
540
+ def __init__(self, *args, **kwargs):
541
+ super().__init__(*args, **kwargs)
542
+
543
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
544
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, 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.
545
+ # 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).
546
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
547
+
548
+ def forward(
549
+ self,
550
+ hidden_states: torch.Tensor,
551
+ attention_mask: Optional[torch.Tensor] = None,
552
+ position_ids: Optional[torch.LongTensor] = None,
553
+ past_key_value: Optional[Cache] = None,
554
+ output_attentions: bool = False,
555
+ use_cache: bool = False,
556
+ **kwargs,
557
+ ):
558
+ if "padding_mask" in kwargs:
559
+ warnings.warn(
560
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
561
+ )
562
+
563
+ # overwrite attention_mask with padding_mask
564
+ attention_mask = kwargs.pop("padding_mask")
565
+ bsz, q_len, _ = hidden_states.size()
566
+
567
+ query_states = self.q_proj(hidden_states)
568
+ key_states = self.k_proj(hidden_states)
569
+ value_states = self.v_proj(hidden_states)
570
+
571
+ query_states = query_states.view(
572
+ bsz, q_len, self.num_heads, self.head_dim
573
+ ).transpose(1, 2)
574
+ key_states = key_states.view(
575
+ bsz, q_len, self.num_key_value_heads, self.head_dim
576
+ ).transpose(1, 2)
577
+ value_states = value_states.view(
578
+ bsz, q_len, self.num_key_value_heads, self.head_dim
579
+ ).transpose(1, 2)
580
+
581
+ kv_seq_len = key_states.shape[-2]
582
+ if past_key_value is not None:
583
+ if self.layer_idx is None:
584
+ raise ValueError(
585
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
586
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
587
+ "with a layer index."
588
+ )
589
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
590
+
591
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
592
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
593
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
594
+
595
+ query_states, key_states = apply_rotary_pos_emb(
596
+ query_states, key_states, cos, sin, position_ids
597
+ )
598
+
599
+ use_sliding_windows = (
600
+ _flash_supports_window_size
601
+ and getattr(self.config, "sliding_window", None) is not None
602
+ and kv_seq_len > self.config.sliding_window
603
+ )
604
+
605
+ if not _flash_supports_window_size:
606
+ logger.warning_once(
607
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
608
+ " make sure to upgrade flash-attn library."
609
+ )
610
+
611
+ if past_key_value is not None:
612
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
613
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
614
+ if (
615
+ getattr(self.config, "sliding_window", None) is not None
616
+ and kv_seq_len > self.config.sliding_window
617
+ and cache_has_contents
618
+ ):
619
+ slicing_tokens = 1 - self.config.sliding_window
620
+
621
+ past_key = past_key_value[self.layer_idx][0]
622
+ past_value = past_key_value[self.layer_idx][1]
623
+
624
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
625
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
626
+
627
+ if past_key.shape[-2] != self.config.sliding_window - 1:
628
+ raise ValueError(
629
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
630
+ f" {past_key.shape}"
631
+ )
632
+
633
+ if attention_mask is not None:
634
+ attention_mask = attention_mask[:, slicing_tokens:]
635
+ attention_mask = torch.cat(
636
+ [attention_mask, torch.ones_like(attention_mask[:, -1:])],
637
+ dim=-1,
638
+ )
639
+
640
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
641
+ key_states, value_states = past_key_value.update(
642
+ key_states, value_states, self.layer_idx, cache_kwargs
643
+ )
644
+
645
+ # repeat k/v heads if n_kv_heads < n_heads
646
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
647
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
648
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
649
+
650
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
651
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
652
+ # cast them back in float16 just to be sure everything works as expected.
653
+ input_dtype = query_states.dtype
654
+ if input_dtype == torch.float32:
655
+ if torch.is_autocast_enabled():
656
+ target_dtype = torch.get_autocast_gpu_dtype()
657
+ # Handle the case where the model is quantized
658
+ elif hasattr(self.config, "_pre_quantization_dtype"):
659
+ target_dtype = self.config._pre_quantization_dtype
660
+ else:
661
+ target_dtype = self.q_proj.weight.dtype
662
+
663
+ logger.warning_once(
664
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
665
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
666
+ f" {target_dtype}."
667
+ )
668
+
669
+ query_states = query_states.to(target_dtype)
670
+ key_states = key_states.to(target_dtype)
671
+ value_states = value_states.to(target_dtype)
672
+
673
+ # Reashape to the expected shape for Flash Attention
674
+ query_states = query_states.transpose(1, 2)
675
+ key_states = key_states.transpose(1, 2)
676
+ value_states = value_states.transpose(1, 2)
677
+
678
+ attn_output = self._flash_attention_forward(
679
+ query_states,
680
+ key_states,
681
+ value_states,
682
+ attention_mask,
683
+ q_len,
684
+ dropout=dropout_rate,
685
+ use_sliding_windows=use_sliding_windows,
686
+ )
687
+
688
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
689
+ attn_output = self.o_proj(attn_output)
690
+
691
+ if not output_attentions:
692
+ attn_weights = None
693
+
694
+ return attn_output, attn_weights, past_key_value
695
+
696
+ def _flash_attention_forward(
697
+ self,
698
+ query_states,
699
+ key_states,
700
+ value_states,
701
+ attention_mask,
702
+ query_length,
703
+ dropout=0.0,
704
+ softmax_scale=None,
705
+ use_sliding_windows=False,
706
+ ):
707
+ """
708
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
709
+ first unpad the input, then computes the attention scores and pad the final attention scores.
710
+ Args:
711
+ query_states (`torch.Tensor`):
712
+ Input query states to be passed to Flash Attention API
713
+ key_states (`torch.Tensor`):
714
+ Input key states to be passed to Flash Attention API
715
+ value_states (`torch.Tensor`):
716
+ Input value states to be passed to Flash Attention API
717
+ attention_mask (`torch.Tensor`):
718
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
719
+ position of padding tokens and 1 for the position of non-padding tokens.
720
+ dropout (`int`, *optional*):
721
+ Attention dropout
722
+ softmax_scale (`float`, *optional*):
723
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
724
+ use_sliding_windows (`bool`, *optional*):
725
+ Whether to activate sliding window attention.
726
+ """
727
+ if not self._flash_attn_uses_top_left_mask:
728
+ causal = self.is_causal
729
+ else:
730
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
731
+ causal = self.is_causal and query_length != 1
732
+
733
+ # Contains at least one padding token in the sequence
734
+ if attention_mask is not None:
735
+ batch_size = query_states.shape[0]
736
+ (
737
+ query_states,
738
+ key_states,
739
+ value_states,
740
+ indices_q,
741
+ cu_seq_lens,
742
+ max_seq_lens,
743
+ ) = self._upad_input(
744
+ query_states, key_states, value_states, attention_mask, query_length
745
+ )
746
+
747
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
748
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
749
+
750
+ if not use_sliding_windows:
751
+ attn_output_unpad = flash_attn_varlen_func(
752
+ query_states,
753
+ key_states,
754
+ value_states,
755
+ cu_seqlens_q=cu_seqlens_q,
756
+ cu_seqlens_k=cu_seqlens_k,
757
+ max_seqlen_q=max_seqlen_in_batch_q,
758
+ max_seqlen_k=max_seqlen_in_batch_k,
759
+ dropout_p=dropout,
760
+ softmax_scale=softmax_scale,
761
+ causal=causal,
762
+ )
763
+ else:
764
+ attn_output_unpad = flash_attn_varlen_func(
765
+ query_states,
766
+ key_states,
767
+ value_states,
768
+ cu_seqlens_q=cu_seqlens_q,
769
+ cu_seqlens_k=cu_seqlens_k,
770
+ max_seqlen_q=max_seqlen_in_batch_q,
771
+ max_seqlen_k=max_seqlen_in_batch_k,
772
+ dropout_p=dropout,
773
+ softmax_scale=softmax_scale,
774
+ causal=causal,
775
+ window_size=(
776
+ self.config.sliding_window,
777
+ self.config.sliding_window,
778
+ ),
779
+ )
780
+
781
+ attn_output = pad_input(
782
+ attn_output_unpad, indices_q, batch_size, query_length
783
+ )
784
+ else:
785
+ if not use_sliding_windows:
786
+ attn_output = flash_attn_func(
787
+ query_states,
788
+ key_states,
789
+ value_states,
790
+ dropout,
791
+ softmax_scale=softmax_scale,
792
+ causal=causal,
793
+ )
794
+ else:
795
+ attn_output = flash_attn_func(
796
+ query_states,
797
+ key_states,
798
+ value_states,
799
+ dropout,
800
+ softmax_scale=softmax_scale,
801
+ causal=causal,
802
+ window_size=(
803
+ self.config.sliding_window,
804
+ self.config.sliding_window,
805
+ ),
806
+ )
807
+
808
+ return attn_output
809
+
810
+ def _upad_input(
811
+ self, query_layer, key_layer, value_layer, attention_mask, query_length
812
+ ):
813
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
814
+
815
+ # On the first iteration we need to properly re-create the padding mask
816
+ # by slicing it on the proper place
817
+ if kv_seq_len != attention_mask.shape[-1]:
818
+ attention_mask_num_tokens = attention_mask.shape[-1]
819
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
820
+
821
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
822
+
823
+ key_layer = index_first_axis(
824
+ key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
825
+ )
826
+ value_layer = index_first_axis(
827
+ value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
828
+ )
829
+
830
+ if query_length == kv_seq_len:
831
+ query_layer = index_first_axis(
832
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim),
833
+ indices_k,
834
+ )
835
+ cu_seqlens_q = cu_seqlens_k
836
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
837
+ indices_q = indices_k
838
+ elif query_length == 1:
839
+ max_seqlen_in_batch_q = 1
840
+ cu_seqlens_q = torch.arange(
841
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
842
+ ) # There is a memcpy here, that is very bad.
843
+ indices_q = cu_seqlens_q[:-1]
844
+ query_layer = query_layer.squeeze(1)
845
+ else:
846
+ # The -q_len: slice assumes left padding.
847
+ attention_mask = attention_mask[:, -query_length:]
848
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
849
+ query_layer, attention_mask
850
+ )
851
+
852
+ return (
853
+ query_layer,
854
+ key_layer,
855
+ value_layer,
856
+ indices_q,
857
+ (cu_seqlens_q, cu_seqlens_k),
858
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
859
+ )
860
+
861
+
862
+ # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Mistral
863
+ class MistralSdpaAttention(MistralAttention):
864
+ """
865
+ Mistral attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
866
+ `MistralAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
867
+ SDPA API.
868
+ """
869
+
870
+ # Adapted from MistralAttention.forward
871
+ def forward(
872
+ self,
873
+ hidden_states: torch.Tensor,
874
+ attention_mask: Optional[torch.Tensor] = None,
875
+ position_ids: Optional[torch.LongTensor] = None,
876
+ past_key_value: Optional[Cache] = None,
877
+ output_attentions: bool = False,
878
+ use_cache: bool = False,
879
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
880
+ if output_attentions:
881
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
882
+ logger.warning_once(
883
+ "MistralModel is using MistralSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
884
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
885
+ )
886
+ return super().forward(
887
+ hidden_states=hidden_states,
888
+ attention_mask=attention_mask,
889
+ position_ids=position_ids,
890
+ past_key_value=past_key_value,
891
+ output_attentions=output_attentions,
892
+ use_cache=use_cache,
893
+ )
894
+
895
+ bsz, q_len, _ = hidden_states.size()
896
+
897
+ query_states = self.q_proj(hidden_states)
898
+ key_states = self.k_proj(hidden_states)
899
+ value_states = self.v_proj(hidden_states)
900
+
901
+ query_states = query_states.view(
902
+ bsz, q_len, self.num_heads, self.head_dim
903
+ ).transpose(1, 2)
904
+ key_states = key_states.view(
905
+ bsz, q_len, self.num_key_value_heads, self.head_dim
906
+ ).transpose(1, 2)
907
+ value_states = value_states.view(
908
+ bsz, q_len, self.num_key_value_heads, self.head_dim
909
+ ).transpose(1, 2)
910
+
911
+ kv_seq_len = key_states.shape[-2]
912
+ if past_key_value is not None:
913
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
914
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
915
+
916
+ query_states, key_states = apply_rotary_pos_emb(
917
+ query_states, key_states, cos, sin, position_ids
918
+ )
919
+
920
+ if past_key_value is not None:
921
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
922
+ key_states, value_states = past_key_value.update(
923
+ key_states, value_states, self.layer_idx, cache_kwargs
924
+ )
925
+
926
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
927
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
928
+
929
+ if attention_mask is not None:
930
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
931
+ raise ValueError(
932
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
933
+ )
934
+
935
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
936
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
937
+ if query_states.device.type == "cuda" and attention_mask is not None:
938
+ query_states = query_states.contiguous()
939
+ key_states = key_states.contiguous()
940
+ value_states = value_states.contiguous()
941
+
942
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
943
+ query_states,
944
+ key_states,
945
+ value_states,
946
+ attn_mask=attention_mask,
947
+ dropout_p=self.attention_dropout if self.training else 0.0,
948
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
949
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
950
+ )
951
+
952
+ attn_output = attn_output.transpose(1, 2).contiguous()
953
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
954
+
955
+ attn_output = self.o_proj(attn_output)
956
+
957
+ return attn_output, None, past_key_value
958
+
959
+
960
+ MISTRAL_ATTENTION_CLASSES = {
961
+ "eager": MistralAttention,
962
+ "flash_attention_2": MistralFlashAttention2,
963
+ "sdpa": MistralSdpaAttention,
964
+ }
965
+
966
+
967
+ class MistralDecoderLayer(nn.Module):
968
+ def __init__(self, config: SparsetralConfig, layer_idx: int):
969
+ super().__init__()
970
+ self.config = config
971
+ self.hidden_size = config.hidden_size
972
+
973
+ self.self_attn = MISTRAL_ATTENTION_CLASSES[config._attn_implementation](
974
+ config, layer_idx
975
+ )
976
+
977
+ self.mlp = MistralMLP(config)
978
+ self.input_layernorm = MistralRMSNorm(
979
+ config.hidden_size, eps=config.rms_norm_eps
980
+ )
981
+ self.post_attention_layernorm = MistralRMSNorm(
982
+ config.hidden_size, eps=config.rms_norm_eps
983
+ )
984
+
985
+ def forward(
986
+ self,
987
+ hidden_states: torch.Tensor,
988
+ attention_mask: Optional[torch.Tensor] = None,
989
+ position_ids: Optional[torch.LongTensor] = None,
990
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
991
+ output_attentions: Optional[bool] = False,
992
+ output_router_logits: Optional[bool] = False,
993
+ use_cache: Optional[bool] = False,
994
+ **kwargs,
995
+ ) -> Tuple[
996
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
997
+ ]:
998
+ if "padding_mask" in kwargs:
999
+ warnings.warn(
1000
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
1001
+ )
1002
+ """
1003
+ Args:
1004
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
1005
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
1006
+ `(batch, sequence_length)` where padding elements are indicated by 0.
1007
+ output_attentions (`bool`, *optional*):
1008
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
1009
+ returned tensors for more detail.
1010
+ use_cache (`bool`, *optional*):
1011
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1012
+ (see `past_key_values`).
1013
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
1014
+ """
1015
+
1016
+ residual = hidden_states
1017
+
1018
+ hidden_states = self.input_layernorm(hidden_states)
1019
+
1020
+ # Self Attention
1021
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
1022
+ hidden_states=hidden_states,
1023
+ attention_mask=attention_mask,
1024
+ position_ids=position_ids,
1025
+ past_key_value=past_key_value,
1026
+ output_attentions=output_attentions,
1027
+ use_cache=use_cache,
1028
+ )
1029
+ hidden_states = residual + hidden_states
1030
+
1031
+ # Fully Connected
1032
+ residual = hidden_states
1033
+ hidden_states = self.post_attention_layernorm(hidden_states)
1034
+ hidden_states, router_logits = self.mlp(hidden_states)
1035
+ hidden_states = residual + hidden_states
1036
+
1037
+ outputs = (hidden_states,)
1038
+
1039
+ if output_attentions:
1040
+ outputs += (self_attn_weights,)
1041
+
1042
+ if use_cache:
1043
+ outputs += (present_key_value,)
1044
+
1045
+ if output_router_logits:
1046
+ outputs += (router_logits,)
1047
+
1048
+ return outputs
1049
+
1050
+
1051
+ MISTRAL_START_DOCSTRING = r"""
1052
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
1053
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
1054
+ etc.)
1055
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
1056
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
1057
+ and behavior.
1058
+ Parameters:
1059
+ config ([`SparsetralConfig`]):
1060
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
1061
+ load the weights associated with the model, only the configuration. Check out the
1062
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
1063
+ """
1064
+
1065
+
1066
+ @add_start_docstrings(
1067
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
1068
+ MISTRAL_START_DOCSTRING,
1069
+ )
1070
+ class MistralPreTrainedModel(PreTrainedModel):
1071
+ config_class = SparsetralConfig
1072
+ base_model_prefix = "model"
1073
+ supports_gradient_checkpointing = True
1074
+ _no_split_modules = ["MistralDecoderLayer"]
1075
+ _skip_keys_device_placement = "past_key_values"
1076
+ _supports_flash_attn_2 = True
1077
+ _supports_sdpa = True
1078
+ _supports_cache_class = True
1079
+
1080
+ def _init_weights(self, module):
1081
+ std = self.config.initializer_range
1082
+ if isinstance(module, nn.Linear):
1083
+ module.weight.data.normal_(mean=0.0, std=std)
1084
+ if module.bias is not None:
1085
+ module.bias.data.zero_()
1086
+ elif isinstance(module, nn.Embedding):
1087
+ module.weight.data.normal_(mean=0.0, std=std)
1088
+ if module.padding_idx is not None:
1089
+ module.weight.data[module.padding_idx].zero_()
1090
+
1091
+ def _set_gradient_checkpointing(self, module, value=False):
1092
+ if isinstance(module, MistralModel):
1093
+ module.gradient_checkpointing = value
1094
+
1095
+
1096
+ MISTRAL_INPUTS_DOCSTRING = r"""
1097
+ Args:
1098
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1099
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1100
+ it.
1101
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1102
+ [`PreTrainedTokenizer.__call__`] for details.
1103
+ [What are input IDs?](../glossary#input-ids)
1104
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1105
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1106
+ - 1 for tokens that are **not masked**,
1107
+ - 0 for tokens that are **masked**.
1108
+ [What are attention masks?](../glossary#attention-mask)
1109
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1110
+ [`PreTrainedTokenizer.__call__`] for details.
1111
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
1112
+ `past_key_values`).
1113
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1114
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1115
+ information on the default strategy.
1116
+ - 1 indicates the head is **not masked**,
1117
+ - 0 indicates the head is **masked**.
1118
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1119
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1120
+ config.n_positions - 1]`.
1121
+ [What are position IDs?](../glossary#position-ids)
1122
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1123
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1124
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1125
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1126
+ Two formats are allowed:
1127
+ - a [`~cache_utils.Cache`] instance;
1128
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1129
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1130
+ cache format.
1131
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1132
+ legacy cache format will be returned.
1133
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1134
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1135
+ of shape `(batch_size, sequence_length)`.
1136
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1137
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1138
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1139
+ model's internal embedding lookup matrix.
1140
+ use_cache (`bool`, *optional*):
1141
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1142
+ `past_key_values`).
1143
+ output_attentions (`bool`, *optional*):
1144
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1145
+ tensors for more detail.
1146
+ output_hidden_states (`bool`, *optional*):
1147
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1148
+ more detail.
1149
+ return_dict (`bool`, *optional*):
1150
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1151
+ """
1152
+
1153
+
1154
+ @add_start_docstrings(
1155
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
1156
+ MISTRAL_START_DOCSTRING,
1157
+ )
1158
+ class MistralModel(MistralPreTrainedModel):
1159
+ """
1160
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MistralDecoderLayer`]
1161
+ Args:
1162
+ config: MistralConfig
1163
+ """
1164
+
1165
+ def __init__(self, config: SparsetralConfig):
1166
+ super().__init__(config)
1167
+ self.padding_idx = config.pad_token_id
1168
+ self.vocab_size = config.vocab_size
1169
+
1170
+ self.embed_tokens = nn.Embedding(
1171
+ config.vocab_size, config.hidden_size, self.padding_idx
1172
+ )
1173
+ self.layers = nn.ModuleList(
1174
+ [
1175
+ MistralDecoderLayer(config, layer_idx)
1176
+ for layer_idx in range(config.num_hidden_layers)
1177
+ ]
1178
+ )
1179
+ self._attn_implementation = config._attn_implementation
1180
+ self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1181
+
1182
+ self.gradient_checkpointing = False
1183
+ # Initialize weights and apply final processing
1184
+ self.post_init()
1185
+
1186
+ def get_input_embeddings(self):
1187
+ return self.embed_tokens
1188
+
1189
+ def set_input_embeddings(self, value):
1190
+ self.embed_tokens = value
1191
+
1192
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1193
+ def forward(
1194
+ self,
1195
+ input_ids: torch.LongTensor = None,
1196
+ attention_mask: Optional[torch.Tensor] = None,
1197
+ position_ids: Optional[torch.LongTensor] = None,
1198
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1199
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1200
+ use_cache: Optional[bool] = None,
1201
+ output_attentions: Optional[bool] = None,
1202
+ output_hidden_states: Optional[bool] = None,
1203
+ output_router_logits: Optional[bool] = None,
1204
+ return_dict: Optional[bool] = None,
1205
+ ) -> Union[Tuple, MoEModelOutputWithPast]:
1206
+ output_attentions = (
1207
+ output_attentions
1208
+ if output_attentions is not None
1209
+ else self.config.output_attentions
1210
+ )
1211
+ output_hidden_states = (
1212
+ output_hidden_states
1213
+ if output_hidden_states is not None
1214
+ else self.config.output_hidden_states
1215
+ )
1216
+ output_router_logits = (
1217
+ output_router_logits
1218
+ if output_router_logits is not None
1219
+ else self.config.output_router_logits
1220
+ )
1221
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1222
+
1223
+ return_dict = (
1224
+ return_dict if return_dict is not None else self.config.use_return_dict
1225
+ )
1226
+
1227
+ # retrieve input_ids and inputs_embeds
1228
+ if input_ids is not None and inputs_embeds is not None:
1229
+ raise ValueError(
1230
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
1231
+ )
1232
+ elif input_ids is not None:
1233
+ batch_size, seq_length = input_ids.shape
1234
+ elif inputs_embeds is not None:
1235
+ batch_size, seq_length, _ = inputs_embeds.shape
1236
+ else:
1237
+ raise ValueError(
1238
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
1239
+ )
1240
+
1241
+ if self.gradient_checkpointing and self.training:
1242
+ if use_cache:
1243
+ logger.warning_once(
1244
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1245
+ )
1246
+ use_cache = False
1247
+
1248
+ seq_length_with_past = seq_length
1249
+ past_key_values_length = 0
1250
+
1251
+ if past_key_values is not None:
1252
+ past_key_values_length = past_key_values[0][0].shape[2]
1253
+ seq_length_with_past = seq_length_with_past + past_key_values_length
1254
+
1255
+ if use_cache:
1256
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1257
+ if use_legacy_cache:
1258
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1259
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1260
+
1261
+ if position_ids is None:
1262
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1263
+ position_ids = torch.arange(
1264
+ past_key_values_length,
1265
+ seq_length + past_key_values_length,
1266
+ dtype=torch.long,
1267
+ device=device,
1268
+ )
1269
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1270
+ else:
1271
+ position_ids = position_ids.view(-1, seq_length).long()
1272
+
1273
+ if inputs_embeds is None:
1274
+ inputs_embeds = self.embed_tokens(input_ids)
1275
+
1276
+ if (
1277
+ attention_mask is not None
1278
+ and self._attn_implementation == "flash_attention_2"
1279
+ and use_cache
1280
+ ):
1281
+ is_padding_right = attention_mask[:, -1].sum().item() != batch_size
1282
+ if is_padding_right:
1283
+ raise ValueError(
1284
+ "You are attempting to perform batched generation with padding_side='right'"
1285
+ " this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
1286
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
1287
+ )
1288
+
1289
+ if self._attn_implementation == "flash_attention_2":
1290
+ # 2d mask is passed through the layers
1291
+ attention_mask = (
1292
+ attention_mask
1293
+ if (attention_mask is not None and 0 in attention_mask)
1294
+ else None
1295
+ )
1296
+ elif self._attn_implementation == "sdpa" and not output_attentions:
1297
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1298
+ # the manual implementation that requires a 4D causal mask in all cases.
1299
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1300
+ attention_mask,
1301
+ (batch_size, seq_length),
1302
+ inputs_embeds,
1303
+ past_key_values_length,
1304
+ )
1305
+ else:
1306
+ # 4d mask is passed through the layers
1307
+ attention_mask = _prepare_4d_causal_attention_mask(
1308
+ attention_mask,
1309
+ (batch_size, seq_length),
1310
+ inputs_embeds,
1311
+ past_key_values_length,
1312
+ sliding_window=self.config.sliding_window,
1313
+ )
1314
+
1315
+ hidden_states = inputs_embeds
1316
+
1317
+ # decoder layers
1318
+ all_hidden_states = () if output_hidden_states else None
1319
+ all_self_attns = () if output_attentions else None
1320
+ all_router_logits = () if output_router_logits else None
1321
+ next_decoder_cache = None
1322
+
1323
+ for decoder_layer in self.layers:
1324
+ if output_hidden_states:
1325
+ all_hidden_states += (hidden_states,)
1326
+
1327
+ if self.gradient_checkpointing and self.training:
1328
+
1329
+ def create_custom_forward(module):
1330
+ def custom_forward(*inputs):
1331
+ # None for past_key_value
1332
+ return module(
1333
+ *inputs, output_attentions, output_router_logits, None
1334
+ )
1335
+
1336
+ return custom_forward
1337
+
1338
+ layer_outputs = torch.utils.checkpoint.checkpoint(
1339
+ create_custom_forward(decoder_layer),
1340
+ hidden_states,
1341
+ attention_mask,
1342
+ position_ids,
1343
+ None,
1344
+ )
1345
+ else:
1346
+ layer_outputs = decoder_layer(
1347
+ hidden_states,
1348
+ attention_mask=attention_mask,
1349
+ position_ids=position_ids,
1350
+ past_key_value=past_key_values,
1351
+ output_attentions=output_attentions,
1352
+ output_router_logits=output_router_logits,
1353
+ use_cache=use_cache,
1354
+ )
1355
+
1356
+ hidden_states = layer_outputs[0]
1357
+
1358
+ if use_cache:
1359
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1360
+
1361
+ if output_attentions:
1362
+ all_self_attns += (layer_outputs[1],)
1363
+
1364
+ if output_router_logits:
1365
+ all_router_logits += (layer_outputs[-1],)
1366
+
1367
+ hidden_states = self.norm(hidden_states)
1368
+
1369
+ # add hidden states from the last decoder layer
1370
+ if output_hidden_states:
1371
+ all_hidden_states += (hidden_states,)
1372
+
1373
+ next_cache = None
1374
+ if use_cache:
1375
+ next_cache = (
1376
+ next_decoder_cache.to_legacy_cache()
1377
+ if use_legacy_cache
1378
+ else next_decoder_cache
1379
+ )
1380
+
1381
+ if not return_dict:
1382
+ return tuple(
1383
+ v
1384
+ for v in [
1385
+ hidden_states,
1386
+ next_cache,
1387
+ all_hidden_states,
1388
+ all_self_attns,
1389
+ all_router_logits,
1390
+ ]
1391
+ if v is not None
1392
+ )
1393
+ return MoEModelOutputWithPast(
1394
+ last_hidden_state=hidden_states,
1395
+ past_key_values=next_cache,
1396
+ hidden_states=all_hidden_states,
1397
+ attentions=all_self_attns,
1398
+ router_logits=all_router_logits,
1399
+ )
1400
+
1401
+
1402
+ class MistralForCausalLM(MistralPreTrainedModel):
1403
+ _tied_weights_keys = ["lm_head.weight"]
1404
+
1405
+ def __init__(self, config):
1406
+ super().__init__(config)
1407
+ self.config = config
1408
+ self.model = MistralModel(config)
1409
+ self.vocab_size = config.vocab_size
1410
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1411
+
1412
+ # Initialize weights and apply final processing
1413
+ self.post_init()
1414
+
1415
+ def get_input_embeddings(self):
1416
+ return self.model.embed_tokens
1417
+
1418
+ def set_input_embeddings(self, value):
1419
+ self.model.embed_tokens = value
1420
+
1421
+ def get_output_embeddings(self):
1422
+ return self.lm_head
1423
+
1424
+ def set_output_embeddings(self, new_embeddings):
1425
+ self.lm_head = new_embeddings
1426
+
1427
+ def set_decoder(self, decoder):
1428
+ self.model = decoder
1429
+
1430
+ def get_decoder(self):
1431
+ return self.model
1432
+
1433
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1434
+ @replace_return_docstrings(
1435
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
1436
+ )
1437
+ def forward(
1438
+ self,
1439
+ input_ids: torch.LongTensor = None,
1440
+ attention_mask: Optional[torch.Tensor] = None,
1441
+ position_ids: Optional[torch.LongTensor] = None,
1442
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1443
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1444
+ labels: Optional[torch.LongTensor] = None,
1445
+ use_cache: Optional[bool] = None,
1446
+ output_attentions: Optional[bool] = None,
1447
+ output_hidden_states: Optional[bool] = None,
1448
+ output_router_logits: Optional[bool] = None,
1449
+ return_dict: Optional[bool] = None,
1450
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1451
+ r"""
1452
+ Args:
1453
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1454
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1455
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1456
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1457
+ Returns:
1458
+ Example:
1459
+ ```python
1460
+ >>> from transformers import AutoTokenizer, MistralForCausalLM
1461
+ >>> model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
1462
+ >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
1463
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1464
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1465
+ >>> # Generate
1466
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1467
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1468
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1469
+ ```"""
1470
+
1471
+ output_attentions = (
1472
+ output_attentions
1473
+ if output_attentions is not None
1474
+ else self.config.output_attentions
1475
+ )
1476
+ output_hidden_states = (
1477
+ output_hidden_states
1478
+ if output_hidden_states is not None
1479
+ else self.config.output_hidden_states
1480
+ )
1481
+ output_router_logits = (
1482
+ output_router_logits
1483
+ if output_router_logits is not None
1484
+ else self.config.output_router_logits
1485
+ )
1486
+ return_dict = (
1487
+ return_dict if return_dict is not None else self.config.use_return_dict
1488
+ )
1489
+
1490
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1491
+ outputs = self.model(
1492
+ input_ids=input_ids,
1493
+ attention_mask=attention_mask,
1494
+ position_ids=position_ids,
1495
+ past_key_values=past_key_values,
1496
+ inputs_embeds=inputs_embeds,
1497
+ use_cache=use_cache,
1498
+ output_attentions=output_attentions,
1499
+ output_hidden_states=output_hidden_states,
1500
+ output_router_logits=output_router_logits,
1501
+ return_dict=return_dict,
1502
+ )
1503
+
1504
+ hidden_states = outputs[0]
1505
+ logits = self.lm_head(hidden_states)
1506
+ logits = logits.float()
1507
+
1508
+ loss = None
1509
+ if labels is not None:
1510
+ # Shift so that tokens < n predict n
1511
+ shift_logits = logits[..., :-1, :].contiguous()
1512
+ shift_labels = labels[..., 1:].contiguous()
1513
+ # Flatten the tokens
1514
+ loss_fct = CrossEntropyLoss()
1515
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1516
+ shift_labels = shift_labels.view(-1)
1517
+ # Enable model parallelism
1518
+ shift_labels = shift_labels.to(shift_logits.device)
1519
+ loss = loss_fct(shift_logits, shift_labels)
1520
+
1521
+ aux_loss = None
1522
+ if output_router_logits:
1523
+ aux_loss = load_balancing_loss_func(
1524
+ outputs.router_logits if return_dict else outputs[-1],
1525
+ self.config.num_experts,
1526
+ self.config.topk,
1527
+ )
1528
+ if labels is not None:
1529
+ loss += self.config.router_aux_loss_coef * aux_loss
1530
+
1531
+ if not return_dict:
1532
+ output = (logits,) + outputs[1:]
1533
+ if output_router_logits:
1534
+ output = (aux_loss,) + output
1535
+ return (loss,) + output if loss is not None else output
1536
+
1537
+ return MoECausalLMOutputWithPast(
1538
+ loss=loss,
1539
+ aux_loss=aux_loss,
1540
+ logits=logits,
1541
+ past_key_values=outputs.past_key_values,
1542
+ hidden_states=outputs.hidden_states,
1543
+ attentions=outputs.attentions,
1544
+ router_logits=outputs.router_logits,
1545
+ )
1546
+
1547
+ def prepare_inputs_for_generation(
1548
+ self,
1549
+ input_ids,
1550
+ past_key_values=None,
1551
+ attention_mask=None,
1552
+ inputs_embeds=None,
1553
+ **kwargs,
1554
+ ):
1555
+ # Omit tokens covered by past_key_values
1556
+ if past_key_values is not None:
1557
+ if isinstance(past_key_values, Cache):
1558
+ cache_length = past_key_values.get_seq_length()
1559
+ past_length = past_key_values.seen_tokens
1560
+ max_cache_length = past_key_values.get_max_length()
1561
+ else:
1562
+ cache_length = past_length = past_key_values[0][0].shape[2]
1563
+ max_cache_length = None
1564
+
1565
+ # Keep only the unprocessed tokens:
1566
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1567
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1568
+ # input)
1569
+ if (
1570
+ attention_mask is not None
1571
+ and attention_mask.shape[1] > input_ids.shape[1]
1572
+ ):
1573
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1574
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1575
+ # input_ids based on the past_length.
1576
+ elif past_length < input_ids.shape[1]:
1577
+ input_ids = input_ids[:, past_length:]
1578
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1579
+
1580
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1581
+ if (
1582
+ max_cache_length is not None
1583
+ and attention_mask is not None
1584
+ and cache_length + input_ids.shape[1] > max_cache_length
1585
+ ):
1586
+ attention_mask = attention_mask[:, -max_cache_length:]
1587
+
1588
+ position_ids = kwargs.get("position_ids", None)
1589
+ if attention_mask is not None and position_ids is None:
1590
+ # create position_ids on the fly for batch generation
1591
+ position_ids = attention_mask.long().cumsum(-1) - 1
1592
+ position_ids.masked_fill_(attention_mask == 0, 1)
1593
+ if past_key_values:
1594
+ position_ids = position_ids[:, -1].unsqueeze(-1)
1595
+
1596
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1597
+ if inputs_embeds is not None and past_key_values is None:
1598
+ model_inputs = {"inputs_embeds": inputs_embeds}
1599
+ else:
1600
+ model_inputs = {"input_ids": input_ids}
1601
+
1602
+ model_inputs.update(
1603
+ {
1604
+ "position_ids": position_ids,
1605
+ "past_key_values": past_key_values,
1606
+ "use_cache": kwargs.get("use_cache"),
1607
+ "attention_mask": attention_mask,
1608
+ }
1609
+ )
1610
+ return model_inputs
1611
+
1612
+ @staticmethod
1613
+ def _reorder_cache(past_key_values, beam_idx):
1614
+ reordered_past = ()
1615
+ for layer_past in past_key_values:
1616
+ reordered_past += (
1617
+ tuple(
1618
+ past_state.index_select(0, beam_idx.to(past_state.device))
1619
+ for past_state in layer_past
1620
+ ),
1621
+ )
1622
+ return reordered_past