emozilla commited on
Commit
6196be7
1 Parent(s): ec4e634

Upload modeling_mistral_yarn.py

Browse files
Files changed (1) hide show
  1. modeling_mistral_yarn.py +1489 -0
modeling_mistral_yarn.py ADDED
@@ -0,0 +1,1489 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from typing import List, Optional, Tuple, Union
24
+
25
+ import torch
26
+ import torch.nn.functional as F
27
+ import torch.utils.checkpoint
28
+ from torch import nn
29
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
30
+
31
+ from transformers.activations import ACT2FN
32
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
33
+ from transformers.modeling_utils import PreTrainedModel
34
+ from transformers.utils import (
35
+ add_start_docstrings,
36
+ add_start_docstrings_to_model_forward,
37
+ is_flash_attn_2_available,
38
+ logging,
39
+ replace_return_docstrings,
40
+ )
41
+ from .configuration_mistral import MistralConfig
42
+
43
+
44
+ if is_flash_attn_2_available():
45
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
46
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
47
+
48
+ _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
49
+
50
+
51
+ logger = logging.get_logger(__name__)
52
+
53
+ _CONFIG_FOR_DOC = "MistralConfig"
54
+
55
+
56
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
57
+ def _get_unpad_data(padding_mask):
58
+ seqlens_in_batch = padding_mask.sum(dim=-1, dtype=torch.int32)
59
+ indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten()
60
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
61
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
62
+ return (
63
+ indices,
64
+ cu_seqlens,
65
+ max_seqlen_in_batch,
66
+ )
67
+
68
+
69
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
70
+ def _make_causal_mask(
71
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
72
+ ):
73
+ """
74
+ Make causal mask used for bi-directional self-attention.
75
+ """
76
+ bsz, tgt_len = input_ids_shape
77
+ mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
78
+ mask_cond = torch.arange(mask.size(-1), device=device)
79
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
80
+ mask = mask.to(dtype)
81
+
82
+ if past_key_values_length > 0:
83
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
84
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
85
+
86
+ def _make_sliding_window_causal_mask(
87
+ input_ids_shape: torch.Size,
88
+ dtype: torch.dtype,
89
+ device: torch.device,
90
+ past_key_values_length: int = 0,
91
+ sliding_window: int = 4096,
92
+ ):
93
+ """
94
+ Make causal mask used for sliding window attention
95
+ """
96
+ bsz, tgt_len = input_ids_shape
97
+
98
+ tensor = torch.full(
99
+ (tgt_len, tgt_len),
100
+ fill_value=1,
101
+ device=device,
102
+ )
103
+ mask = torch.tril(tensor, diagonal=0)
104
+ # make the mask banded to account for sliding window
105
+ mask = torch.triu(mask, diagonal=-sliding_window)
106
+ mask = torch.log(mask).to(dtype)
107
+
108
+ if past_key_values_length > 0:
109
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
110
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
111
+
112
+
113
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
114
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
115
+ """
116
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
117
+ """
118
+ bsz, src_len = mask.size()
119
+ tgt_len = tgt_len if tgt_len is not None else src_len
120
+
121
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
122
+
123
+ inverted_mask = 1.0 - expanded_mask
124
+
125
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
126
+
127
+ # Inverse dim formula to find dim based on number of rotations
128
+ def _yarn_find_correction_dim(num_rotations, dim, base=10000, max_position_embeddings=2048):
129
+ return (dim * math.log(max_position_embeddings/(num_rotations * 2 * math.pi)))/(2 * math.log(base))
130
+
131
+ # Find dim range bounds based on rotations
132
+ def _yarn_find_correction_range(low_rot, high_rot, dim, base=10000, max_position_embeddings=2048):
133
+ low = math.floor(_yarn_find_correction_dim(
134
+ low_rot, dim, base, max_position_embeddings))
135
+ high = math.ceil(_yarn_find_correction_dim(
136
+ high_rot, dim, base, max_position_embeddings))
137
+ return max(low, 0), min(high, dim-1) # Clamp values just in case
138
+
139
+ def _yarn_linear_ramp_mask(min, max, dim):
140
+ if min == max:
141
+ max += 0.001 # Prevent singularity
142
+
143
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
144
+ ramp_func = torch.clamp(linear_func, 0, 1)
145
+ return ramp_func
146
+
147
+ def _yarn_get_mscale(scale=1):
148
+ if scale <= 1:
149
+ return 1.0
150
+ return 0.07 * math.log(scale) + 1.0
151
+
152
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral
153
+ class MistralRMSNorm(nn.Module):
154
+ def __init__(self, hidden_size, eps=1e-6):
155
+ """
156
+ MistralRMSNorm is equivalent to T5LayerNorm
157
+ """
158
+ super().__init__()
159
+ self.weight = nn.Parameter(torch.ones(hidden_size))
160
+ self.variance_epsilon = eps
161
+
162
+ def forward(self, hidden_states):
163
+ input_dtype = hidden_states.dtype
164
+ hidden_states = hidden_states.to(torch.float32)
165
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
166
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
167
+ return self.weight * hidden_states.to(input_dtype)
168
+
169
+
170
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Mistral
171
+ class MistralRotaryEmbedding(nn.Module):
172
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
173
+ super().__init__()
174
+
175
+ self.dim = dim
176
+ self.max_position_embeddings = max_position_embeddings
177
+ self.base = base
178
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
179
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
180
+
181
+ # Build here to make `torch.jit.trace` work.
182
+ self._set_cos_sin_cache(
183
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
184
+ )
185
+
186
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
187
+ self.max_seq_len_cached = seq_len
188
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
189
+
190
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
191
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
192
+ emb = torch.cat((freqs, freqs), dim=-1)
193
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
194
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
195
+
196
+ def forward(self, x, seq_len=None):
197
+ # x: [bs, num_attention_heads, seq_len, head_size]
198
+ if seq_len > self.max_seq_len_cached:
199
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
200
+
201
+ return (
202
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
203
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
204
+ )
205
+
206
+ class MistralLinearScalingRotaryEmbedding(MistralRotaryEmbedding):
207
+ """MistralRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
208
+
209
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
210
+ self.scaling_factor = scaling_factor
211
+ super().__init__(dim, max_position_embeddings, base, device)
212
+
213
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
214
+ self.max_seq_len_cached = seq_len
215
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
216
+ t = t / self.scaling_factor
217
+
218
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
219
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
220
+ emb = torch.cat((freqs, freqs), dim=-1)
221
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
222
+ self.register_buffer("sin_cached", emb.cos().to(dtype), persistent=False)
223
+
224
+
225
+ class MistralDynamicNTKScalingRotaryEmbedding(MistralRotaryEmbedding):
226
+ """MistralRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
227
+
228
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
229
+ self.scaling_factor = scaling_factor
230
+ super().__init__(dim, max_position_embeddings, base, device)
231
+
232
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
233
+ self.max_seq_len_cached = seq_len
234
+
235
+ if seq_len > self.max_position_embeddings:
236
+ base = self.base * (
237
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
238
+ ) ** (self.dim / (self.dim - 2))
239
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
240
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
241
+
242
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
243
+
244
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
245
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
246
+ emb = torch.cat((freqs, freqs), dim=-1)
247
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
248
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
249
+
250
+
251
+ class MistralYaRNScaledRotaryEmbedding(torch.nn.Module):
252
+ """MistralRotaryEmbedding extended with YaRN. See: https://arxiv.org/abs/2309.00071"""
253
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, scale=1, original_max_position_embeddings=2048,
254
+ extrapolation_factor=1, attn_factor=1, beta_fast=128, beta_slow=2, finetuned=False, device=None):
255
+ super().__init__()
256
+
257
+ self.dim = dim
258
+ self.max_position_embeddings = max_position_embeddings
259
+ self.base = base
260
+ self.scale = scale
261
+ self.original_max_position_embeddings = original_max_position_embeddings
262
+ self.extrapolation_factor = extrapolation_factor
263
+ self.attn_factor = attn_factor
264
+ self.beta_fast = beta_fast
265
+ self.beta_slow = beta_slow
266
+
267
+ self.yarn(device)
268
+
269
+ # Build here to make `torch.jit.trace` work.
270
+ self.max_seq_len_cached = max_position_embeddings
271
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
272
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
273
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
274
+ emb = torch.cat((freqs, freqs), dim=-1)
275
+ dtype = torch.get_default_dtype()
276
+
277
+ self.register_buffer("cos_cached", (emb.cos() * self.mscale).to(dtype), persistent=False)
278
+ self.register_buffer("sin_cached", (emb.sin() * self.mscale).to(dtype), persistent=False)
279
+
280
+ def forward(self, x, seq_len=None):
281
+ # x: [bs, num_attention_heads, seq_len, head_size]
282
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
283
+ if seq_len > self.max_seq_len_cached:
284
+ self.max_seq_len_cached = seq_len
285
+
286
+ t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
287
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
288
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
289
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
290
+
291
+ self.register_buffer("cos_cached", (emb.cos() * self.mscale).to(x.dtype), persistent=False)
292
+ self.register_buffer("sin_cached", (emb.sin() * self.mscale).to(x.dtype), persistent=False)
293
+ return (
294
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
295
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
296
+ )
297
+
298
+ def yarn(self, device):
299
+ pos_freqs = self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
300
+ inv_freq_extrapolation = 1.0 / pos_freqs
301
+ inv_freq_interpolation = 1.0 / (self.scale * pos_freqs)
302
+
303
+ low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow, self.dim, self.base, self.original_max_position_embeddings)
304
+ inv_freq_mask = (1 - _yarn_linear_ramp_mask(low, high, self.dim // 2).float().to(device)) * self.extrapolation_factor # Get n-d rotational scaling corrected for extrapolation
305
+ inv_freq = inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
306
+
307
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
308
+ self.mscale = float(_yarn_get_mscale(self.scale) * self.attn_factor) # Get n-d magnitude scaling corrected for interpolation
309
+
310
+
311
+ class MistralDynamicYaRNScaledRotaryEmbedding(torch.nn.Module):
312
+ """MistralRotaryEmbedding extended with Dynamic YaRN. See: https://arxiv.org/abs/2309.00071"""
313
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, original_max_position_embeddings=2048,
314
+ extrapolation_factor=1, attn_factor=1, beta_fast=128, beta_slow=2, finetuned=False, device=None):
315
+ super().__init__()
316
+
317
+ self.dim = dim
318
+ self.max_position_embeddings = max_position_embeddings
319
+ self.base = base
320
+ self.original_max_position_embeddings = original_max_position_embeddings
321
+ self.extrapolation_factor = extrapolation_factor
322
+ self.attn_factor = attn_factor
323
+ self.beta_fast = beta_fast
324
+ self.beta_slow = beta_slow
325
+
326
+ if finetuned:
327
+ self.yarn(self.max_position_embeddings / self.original_max_position_embeddings, device)
328
+ else:
329
+ inv_freq = 1.0 / \
330
+ (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
331
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
332
+ self.mscale = 1
333
+
334
+ # Build here to make `torch.jit.trace` work.
335
+ self.max_seq_len_cached = max_position_embeddings
336
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
337
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
338
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
339
+ emb = torch.cat((freqs, freqs), dim=-1)
340
+ dtype = torch.get_default_dtype()
341
+
342
+ self.register_buffer("cos_cached", (emb.cos() * self.mscale).to(dtype), persistent=False)
343
+ self.register_buffer("sin_cached", (emb.sin() * self.mscale).to(dtype), persistent=False)
344
+
345
+ def forward(self, x, seq_len=None):
346
+ # x: [bs, num_attention_heads, seq_len, head_size]
347
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
348
+ if seq_len > self.max_seq_len_cached:
349
+ self.max_seq_len_cached = seq_len
350
+
351
+ self.yarn(seq_len / self.max_position_embeddings, x.device)
352
+
353
+ t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
354
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
355
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
356
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
357
+
358
+ self.register_buffer("cos_cached", (emb.cos() * self.mscale).to(x.dtype), persistent=False)
359
+ self.register_buffer("sin_cached", (emb.sin() * self.mscale).to(x.dtype), persistent=False)
360
+ return (
361
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
362
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
363
+ )
364
+
365
+ def yarn(self, scale, device):
366
+ pos_freqs = self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
367
+ inv_freq_extrapolation = 1.0 / pos_freqs
368
+ inv_freq_interpolation = 1.0 / (scale * pos_freqs)
369
+
370
+ low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow, self.dim, self.base, self.original_max_position_embeddings)
371
+ inv_freq_mask = (1 - _yarn_linear_ramp_mask(low, high, self.dim // 2).float().to(device)) * self.extrapolation_factor # Get n-d rotational scaling corrected for extrapolation
372
+ inv_freq = inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
373
+
374
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
375
+ self.mscale = float(_yarn_get_mscale(scale) * self.attn_factor) # Get n-d magnitude scaling corrected for interpolation
376
+
377
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
378
+ def rotate_half(x):
379
+ """Rotates half the hidden dims of the input."""
380
+ x1 = x[..., : x.shape[-1] // 2]
381
+ x2 = x[..., x.shape[-1] // 2 :]
382
+ return torch.cat((-x2, x1), dim=-1)
383
+
384
+
385
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
386
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
387
+ cos = cos[position_ids].unsqueeze(1) # [seq_len, dim] -> [batch_size, 1, seq_len, head_dim]
388
+ sin = sin[position_ids].unsqueeze(1)
389
+ q_embed = (q * cos) + (rotate_half(q) * sin)
390
+ k_embed = (k * cos) + (rotate_half(k) * sin)
391
+ return q_embed, k_embed
392
+
393
+
394
+ class MistralMLP(nn.Module):
395
+ def __init__(self, config):
396
+ super().__init__()
397
+ self.config = config
398
+ self.hidden_size = config.hidden_size
399
+ self.intermediate_size = config.intermediate_size
400
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
401
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
402
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
403
+ self.act_fn = ACT2FN[config.hidden_act]
404
+
405
+ def forward(self, x):
406
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
407
+
408
+
409
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
410
+ """
411
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
412
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
413
+ """
414
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
415
+ if n_rep == 1:
416
+ return hidden_states
417
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
418
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
419
+
420
+
421
+ class MistralAttention(nn.Module):
422
+ """
423
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
424
+ and "Generating Long Sequences with Sparse Transformers".
425
+ """
426
+
427
+ def __init__(self, config: MistralConfig):
428
+ super().__init__()
429
+ self.config = config
430
+ self.hidden_size = config.hidden_size
431
+ self.num_heads = config.num_attention_heads
432
+ self.head_dim = self.hidden_size // self.num_heads
433
+ self.num_key_value_heads = config.num_key_value_heads
434
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
435
+ self.max_position_embeddings = config.max_position_embeddings
436
+ self.rope_theta = config.rope_theta
437
+
438
+ if (self.head_dim * self.num_heads) != self.hidden_size:
439
+ raise ValueError(
440
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
441
+ f" and `num_heads`: {self.num_heads})."
442
+ )
443
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
444
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
445
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
446
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
447
+
448
+ self._init_rope()
449
+
450
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
451
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
452
+
453
+ def _init_rope(self):
454
+ if self.config.rope_scaling is None:
455
+ self.rotary_emb = MistralRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta)
456
+ else:
457
+ scaling_type = self.config.rope_scaling["type"]
458
+ scaling_factor = self.config.rope_scaling["factor"]
459
+ if scaling_type == "linear":
460
+ self.rotary_emb = MistralLinearScalingRotaryEmbedding(
461
+ self.head_dim, max_position_embeddings=self.max_position_embeddings,
462
+ scaling_factor=scaling_factor, base=self.rope_theta,
463
+ )
464
+ elif scaling_type == "dynamic":
465
+ self.rotary_emb = MistralDynamicNTKScalingRotaryEmbedding(
466
+ self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor,
467
+ base=self.rope_theta,
468
+ )
469
+ elif scaling_type == "yarn":
470
+ original_max_position_embeddings = self.config.rope_scaling["original_max_position_embeddings"]
471
+ self.rotary_emb = MistralYaRNScaledRotaryEmbedding(
472
+ self.head_dim, max_position_embeddings=self.max_position_embeddings, scale=scaling_factor,
473
+ original_max_position_embeddings=original_max_position_embeddings, base=self.rope_theta,
474
+ )
475
+ elif scaling_type == "dynamic-yarn":
476
+ original_max_position_embeddings = self.config.rope_scaling["original_max_position_embeddings"]
477
+ self.rotary_emb = MistralDynamicYaRNScaledRotaryEmbedding(
478
+ self.head_dim, max_position_embeddings=self.max_position_embeddings,
479
+ original_max_position_embeddings=original_max_position_embeddings, base=self.rope_theta,
480
+ )
481
+ else:
482
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
483
+
484
+ def forward(
485
+ self,
486
+ hidden_states: torch.Tensor,
487
+ attention_mask: Optional[torch.Tensor] = None,
488
+ position_ids: Optional[torch.LongTensor] = None,
489
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
490
+ output_attentions: bool = False,
491
+ use_cache: bool = False,
492
+ padding_mask: Optional[torch.Tensor] = None,
493
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
494
+ bsz, q_len, _ = hidden_states.size()
495
+
496
+ query_states = self.q_proj(hidden_states)
497
+ key_states = self.k_proj(hidden_states)
498
+ value_states = self.v_proj(hidden_states)
499
+
500
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
501
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
502
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
503
+
504
+ kv_seq_len = key_states.shape[-2]
505
+ if past_key_value is not None:
506
+ kv_seq_len += past_key_value[0].shape[-2]
507
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
508
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
509
+
510
+ if past_key_value is not None:
511
+ # reuse k, v, self_attention
512
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
513
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
514
+
515
+ past_key_value = (key_states, value_states) if use_cache else None
516
+
517
+ # repeat k/v heads if n_kv_heads < n_heads
518
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
519
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
520
+
521
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
522
+
523
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
524
+ raise ValueError(
525
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
526
+ f" {attn_weights.size()}"
527
+ )
528
+
529
+ if attention_mask is not None:
530
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
531
+ raise ValueError(
532
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
533
+ )
534
+
535
+ attn_weights = attn_weights + attention_mask
536
+
537
+ # upcast attention to fp32
538
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
539
+ attn_output = torch.matmul(attn_weights, value_states)
540
+
541
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
542
+ raise ValueError(
543
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
544
+ f" {attn_output.size()}"
545
+ )
546
+
547
+ attn_output = attn_output.transpose(1, 2).contiguous()
548
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
549
+
550
+ attn_output = self.o_proj(attn_output)
551
+
552
+ if not output_attentions:
553
+ attn_weights = None
554
+
555
+ return attn_output, attn_weights, past_key_value
556
+
557
+
558
+ class MistralFlashAttention2(MistralAttention):
559
+ """
560
+ Mistral flash attention module. This module inherits from `MistralAttention` as the weights of the module stays
561
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
562
+ flash attention and deal with padding tokens in case the input contains any of them.
563
+ """
564
+
565
+ def forward(
566
+ self,
567
+ hidden_states: torch.Tensor,
568
+ attention_mask: Optional[torch.Tensor] = None,
569
+ position_ids: Optional[torch.LongTensor] = None,
570
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
571
+ output_attentions: bool = False,
572
+ use_cache: bool = False,
573
+ padding_mask: Optional[torch.LongTensor] = None,
574
+ ):
575
+ bsz, q_len, _ = hidden_states.size()
576
+
577
+ query_states = self.q_proj(hidden_states)
578
+ key_states = self.k_proj(hidden_states)
579
+ value_states = self.v_proj(hidden_states)
580
+
581
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
582
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
583
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
584
+
585
+ kv_seq_len = key_states.shape[-2]
586
+ if past_key_value is not None:
587
+ kv_seq_len += past_key_value[0].shape[-2]
588
+
589
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
590
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
591
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
592
+
593
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
594
+
595
+ use_sliding_windows = (
596
+ _flash_supports_window_size
597
+ and hasattr(self.config, "sliding_window")
598
+ and kv_seq_len > self.config.sliding_window
599
+ )
600
+
601
+ if not _flash_supports_window_size:
602
+ logger.warning_once(
603
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
604
+ " make sure to upgrade flash-attn library."
605
+ )
606
+
607
+ if past_key_value is not None:
608
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
609
+ if hasattr(self.config, "sliding_window") and kv_seq_len > self.config.sliding_window:
610
+ slicing_tokens = kv_seq_len - self.config.sliding_window
611
+
612
+ past_key = past_key_value[0]
613
+ past_value = past_key_value[1]
614
+
615
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
616
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
617
+
618
+ if past_key.shape[-2] != self.config.sliding_window - 1:
619
+ raise ValueError(
620
+ f"past key much have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
621
+ f" {past_key.shape}"
622
+ )
623
+
624
+ past_key_value = (past_key, past_value)
625
+
626
+ if padding_mask is not None:
627
+ padding_mask = padding_mask[:, slicing_tokens:]
628
+ padding_mask = torch.cat([padding_mask, torch.ones_like(padding_mask[:, -1:])], dim=-1)
629
+
630
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
631
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
632
+
633
+ past_key_value = (key_states, value_states) if use_cache else None
634
+
635
+ # repeat k/v heads if n_kv_heads < n_heads
636
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
637
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
638
+
639
+ # TODO: Mistral does not have dropout in the config??
640
+ # It is recommended to use dropout with FA according to the docs
641
+ # when training.
642
+ dropout_rate = 0.0 # if not self.training else self.attn_dropout
643
+
644
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
645
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
646
+ # cast them back in float16 just to be sure everything works as expected.
647
+ input_dtype = query_states.dtype
648
+ if input_dtype == torch.float32:
649
+ logger.warning_once(
650
+ "The input hidden states seems to be silently casted in float32, this might be related to"
651
+ " the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
652
+ " float16."
653
+ )
654
+
655
+ query_states = query_states.to(torch.float16)
656
+ key_states = key_states.to(torch.float16)
657
+ value_states = value_states.to(torch.float16)
658
+
659
+ # Reashape to the expected shape for Flash Attention
660
+ query_states = query_states.transpose(1, 2)
661
+ key_states = key_states.transpose(1, 2)
662
+ value_states = value_states.transpose(1, 2)
663
+
664
+ attn_output = self._flash_attention_forward(
665
+ query_states,
666
+ key_states,
667
+ value_states,
668
+ padding_mask,
669
+ q_len,
670
+ dropout=dropout_rate,
671
+ use_sliding_windows=use_sliding_windows,
672
+ )
673
+
674
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
675
+ attn_output = self.o_proj(attn_output)
676
+
677
+ if not output_attentions:
678
+ attn_weights = None
679
+
680
+ return attn_output, attn_weights, past_key_value
681
+
682
+ def _flash_attention_forward(
683
+ self,
684
+ query_states,
685
+ key_states,
686
+ value_states,
687
+ padding_mask,
688
+ query_length,
689
+ dropout=0.0,
690
+ softmax_scale=None,
691
+ use_sliding_windows=False,
692
+ ):
693
+ """
694
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
695
+ first unpad the input, then computes the attention scores and pad the final attention scores.
696
+
697
+ Args:
698
+ query_states (`torch.Tensor`):
699
+ Input query states to be passed to Flash Attention API
700
+ key_states (`torch.Tensor`):
701
+ Input key states to be passed to Flash Attention API
702
+ value_states (`torch.Tensor`):
703
+ Input value states to be passed to Flash Attention API
704
+ padding_mask (`torch.Tensor`):
705
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
706
+ position of padding tokens and 1 for the position of non-padding tokens.
707
+ dropout (`int`, *optional*):
708
+ Attention dropout
709
+ softmax_scale (`float`, *optional*):
710
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
711
+ use_sliding_windows (`bool`, *optional*):
712
+ Whether to activate sliding window attention.
713
+ """
714
+ # Contains at least one padding token in the sequence
715
+ if padding_mask is not None:
716
+ batch_size = query_states.shape[0]
717
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
718
+ query_states, key_states, value_states, padding_mask, query_length
719
+ )
720
+
721
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
722
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
723
+
724
+ if not use_sliding_windows:
725
+ attn_output_unpad = flash_attn_varlen_func(
726
+ query_states,
727
+ key_states,
728
+ value_states,
729
+ cu_seqlens_q=cu_seqlens_q,
730
+ cu_seqlens_k=cu_seqlens_k,
731
+ max_seqlen_q=max_seqlen_in_batch_q,
732
+ max_seqlen_k=max_seqlen_in_batch_k,
733
+ dropout_p=dropout,
734
+ softmax_scale=softmax_scale,
735
+ causal=True,
736
+ )
737
+ else:
738
+ attn_output_unpad = flash_attn_varlen_func(
739
+ query_states,
740
+ key_states,
741
+ value_states,
742
+ cu_seqlens_q=cu_seqlens_q,
743
+ cu_seqlens_k=cu_seqlens_k,
744
+ max_seqlen_q=max_seqlen_in_batch_q,
745
+ max_seqlen_k=max_seqlen_in_batch_k,
746
+ dropout_p=dropout,
747
+ softmax_scale=softmax_scale,
748
+ causal=True,
749
+ window_size=(self.config.sliding_window, self.config.sliding_window),
750
+ )
751
+
752
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
753
+ else:
754
+ if not use_sliding_windows:
755
+ attn_output = flash_attn_func(
756
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=True
757
+ )
758
+ else:
759
+ attn_output = flash_attn_func(
760
+ query_states,
761
+ key_states,
762
+ value_states,
763
+ dropout,
764
+ softmax_scale=softmax_scale,
765
+ causal=True,
766
+ window_size=(self.config.sliding_window, self.config.sliding_window),
767
+ )
768
+
769
+ return attn_output
770
+
771
+ def _upad_input(self, query_layer, key_layer, value_layer, padding_mask, query_length):
772
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
773
+
774
+ # On the first iteration we need to properly re-create the padding mask
775
+ # by slicing it on the proper place
776
+ if kv_seq_len != padding_mask.shape[-1]:
777
+ padding_mask_num_tokens = padding_mask.shape[-1]
778
+ padding_mask = padding_mask[:, padding_mask_num_tokens - kv_seq_len :]
779
+
780
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(padding_mask)
781
+
782
+ key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
783
+ value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
784
+
785
+ if query_length == kv_seq_len:
786
+ query_layer = index_first_axis(
787
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
788
+ )
789
+ cu_seqlens_q = cu_seqlens_k
790
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
791
+ indices_q = indices_k
792
+ elif query_length == 1:
793
+ max_seqlen_in_batch_q = 1
794
+ cu_seqlens_q = torch.arange(
795
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
796
+ ) # There is a memcpy here, that is very bad.
797
+ indices_q = cu_seqlens_q[:-1]
798
+ query_layer = query_layer.squeeze(1)
799
+ else:
800
+ # The -q_len: slice assumes left padding.
801
+ padding_mask = padding_mask[:, -query_length:]
802
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, padding_mask)
803
+
804
+ return (
805
+ query_layer,
806
+ key_layer,
807
+ value_layer,
808
+ indices_q,
809
+ (cu_seqlens_q, cu_seqlens_k),
810
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
811
+ )
812
+
813
+
814
+ class MistralDecoderLayer(nn.Module):
815
+ def __init__(self, config: MistralConfig):
816
+ super().__init__()
817
+ self.hidden_size = config.hidden_size
818
+ self.self_attn = (
819
+ MistralAttention(config=config)
820
+ if not getattr(config, "_flash_attn_2_enabled", False)
821
+ else MistralFlashAttention2(config)
822
+ )
823
+ self.mlp = MistralMLP(config)
824
+ self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
825
+ self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
826
+
827
+ def forward(
828
+ self,
829
+ hidden_states: torch.Tensor,
830
+ attention_mask: Optional[torch.Tensor] = None,
831
+ position_ids: Optional[torch.LongTensor] = None,
832
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
833
+ output_attentions: Optional[bool] = False,
834
+ use_cache: Optional[bool] = False,
835
+ padding_mask: Optional[torch.Tensor] = None,
836
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
837
+ """
838
+ Args:
839
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
840
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
841
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
842
+ output_attentions (`bool`, *optional*):
843
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
844
+ returned tensors for more detail.
845
+ use_cache (`bool`, *optional*):
846
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
847
+ (see `past_key_values`).
848
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
849
+ """
850
+
851
+ residual = hidden_states
852
+
853
+ hidden_states = self.input_layernorm(hidden_states)
854
+
855
+ # Self Attention
856
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
857
+ hidden_states=hidden_states,
858
+ attention_mask=attention_mask,
859
+ position_ids=position_ids,
860
+ past_key_value=past_key_value,
861
+ output_attentions=output_attentions,
862
+ use_cache=use_cache,
863
+ padding_mask=padding_mask,
864
+ )
865
+ hidden_states = residual + hidden_states
866
+
867
+ # Fully Connected
868
+ residual = hidden_states
869
+ hidden_states = self.post_attention_layernorm(hidden_states)
870
+ hidden_states = self.mlp(hidden_states)
871
+ hidden_states = residual + hidden_states
872
+
873
+ outputs = (hidden_states,)
874
+
875
+ if output_attentions:
876
+ outputs += (self_attn_weights,)
877
+
878
+ if use_cache:
879
+ outputs += (present_key_value,)
880
+
881
+ return outputs
882
+
883
+
884
+ MISTRAL_START_DOCSTRING = r"""
885
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
886
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
887
+ etc.)
888
+
889
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
890
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
891
+ and behavior.
892
+
893
+ Parameters:
894
+ config ([`MistralConfig`]):
895
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
896
+ load the weights associated with the model, only the configuration. Check out the
897
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
898
+ """
899
+
900
+
901
+ @add_start_docstrings(
902
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
903
+ MISTRAL_START_DOCSTRING,
904
+ )
905
+ class MistralPreTrainedModel(PreTrainedModel):
906
+ config_class = MistralConfig
907
+ base_model_prefix = "model"
908
+ supports_gradient_checkpointing = True
909
+ _no_split_modules = ["MistralDecoderLayer"]
910
+ _skip_keys_device_placement = "past_key_values"
911
+ _supports_flash_attn_2 = True
912
+
913
+ def _init_weights(self, module):
914
+ std = self.config.initializer_range
915
+ if isinstance(module, nn.Linear):
916
+ module.weight.data.normal_(mean=0.0, std=std)
917
+ if module.bias is not None:
918
+ module.bias.data.zero_()
919
+ elif isinstance(module, nn.Embedding):
920
+ module.weight.data.normal_(mean=0.0, std=std)
921
+ if module.padding_idx is not None:
922
+ module.weight.data[module.padding_idx].zero_()
923
+
924
+ def _set_gradient_checkpointing(self, module, value=False):
925
+ if isinstance(module, MistralModel):
926
+ module.gradient_checkpointing = value
927
+
928
+
929
+ MISTRAL_INPUTS_DOCSTRING = r"""
930
+ Args:
931
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
932
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
933
+ it.
934
+
935
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
936
+ [`PreTrainedTokenizer.__call__`] for details.
937
+
938
+ [What are input IDs?](../glossary#input-ids)
939
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
940
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
941
+
942
+ - 1 for tokens that are **not masked**,
943
+ - 0 for tokens that are **masked**.
944
+
945
+ [What are attention masks?](../glossary#attention-mask)
946
+
947
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
948
+ [`PreTrainedTokenizer.__call__`] for details.
949
+
950
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
951
+ `past_key_values`).
952
+
953
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
954
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
955
+ information on the default strategy.
956
+
957
+ - 1 indicates the head is **not masked**,
958
+ - 0 indicates the head is **masked**.
959
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
960
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
961
+ config.n_positions - 1]`.
962
+
963
+ [What are position IDs?](../glossary#position-ids)
964
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
965
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
966
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
967
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
968
+
969
+ Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
970
+ blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
971
+
972
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
973
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
974
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
975
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
976
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
977
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
978
+ model's internal embedding lookup matrix.
979
+ use_cache (`bool`, *optional*):
980
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
981
+ `past_key_values`).
982
+ output_attentions (`bool`, *optional*):
983
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
984
+ tensors for more detail.
985
+ output_hidden_states (`bool`, *optional*):
986
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
987
+ more detail.
988
+ return_dict (`bool`, *optional*):
989
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
990
+ """
991
+
992
+
993
+ @add_start_docstrings(
994
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
995
+ MISTRAL_START_DOCSTRING,
996
+ )
997
+ class MistralModel(MistralPreTrainedModel):
998
+ """
999
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MistralDecoderLayer`]
1000
+
1001
+ Args:
1002
+ config: MistralConfig
1003
+ """
1004
+
1005
+ def __init__(self, config: MistralConfig):
1006
+ super().__init__(config)
1007
+ self.padding_idx = config.pad_token_id
1008
+ self.vocab_size = config.vocab_size
1009
+
1010
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
1011
+ self.layers = nn.ModuleList([MistralDecoderLayer(config) for _ in range(config.num_hidden_layers)])
1012
+ self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1013
+
1014
+ self.gradient_checkpointing = False
1015
+ # Initialize weights and apply final processing
1016
+ self.post_init()
1017
+
1018
+ def get_input_embeddings(self):
1019
+ return self.embed_tokens
1020
+
1021
+ def set_input_embeddings(self, value):
1022
+ self.embed_tokens = value
1023
+
1024
+ def _prepare_decoder_attention_mask(
1025
+ self, attention_mask, input_shape, inputs_embeds, past_key_values_length, sliding_window
1026
+ ):
1027
+ # create causal mask
1028
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
1029
+ combined_attention_mask = None
1030
+ if input_shape[-1] > 1:
1031
+ if sliding_window is not None:
1032
+ combined_attention_mask = _make_sliding_window_causal_mask(
1033
+ input_shape,
1034
+ inputs_embeds.dtype,
1035
+ device=inputs_embeds.device,
1036
+ past_key_values_length=past_key_values_length,
1037
+ sliding_window=sliding_window,
1038
+ )
1039
+ else:
1040
+ combined_attention_mask = _make_causal_mask(
1041
+ input_shape,
1042
+ inputs_embeds.dtype,
1043
+ device=inputs_embeds.device,
1044
+ past_key_values_length=past_key_values_length,
1045
+ )
1046
+
1047
+ if attention_mask is not None:
1048
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
1049
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
1050
+ inputs_embeds.device
1051
+ )
1052
+ combined_attention_mask = (
1053
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
1054
+ )
1055
+
1056
+ return combined_attention_mask
1057
+
1058
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1059
+ def forward(
1060
+ self,
1061
+ input_ids: torch.LongTensor = None,
1062
+ attention_mask: Optional[torch.Tensor] = None,
1063
+ position_ids: Optional[torch.LongTensor] = None,
1064
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1065
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1066
+ use_cache: Optional[bool] = None,
1067
+ output_attentions: Optional[bool] = None,
1068
+ output_hidden_states: Optional[bool] = None,
1069
+ return_dict: Optional[bool] = None,
1070
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
1071
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1072
+ output_hidden_states = (
1073
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1074
+ )
1075
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1076
+
1077
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1078
+
1079
+ # retrieve input_ids and inputs_embeds
1080
+ if input_ids is not None and inputs_embeds is not None:
1081
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
1082
+ elif input_ids is not None:
1083
+ batch_size, seq_length = input_ids.shape
1084
+ elif inputs_embeds is not None:
1085
+ batch_size, seq_length, _ = inputs_embeds.shape
1086
+ else:
1087
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
1088
+
1089
+ seq_length_with_past = seq_length
1090
+ past_key_values_length = 0
1091
+
1092
+ if past_key_values is not None:
1093
+ past_key_values_length = past_key_values[0][0].shape[2]
1094
+ seq_length_with_past = seq_length_with_past + past_key_values_length
1095
+
1096
+ if position_ids is None:
1097
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1098
+ position_ids = torch.arange(
1099
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1100
+ )
1101
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1102
+ else:
1103
+ position_ids = position_ids.view(-1, seq_length).long()
1104
+
1105
+ if inputs_embeds is None:
1106
+ inputs_embeds = self.embed_tokens(input_ids)
1107
+
1108
+ padding_mask = None
1109
+
1110
+ # embed positions
1111
+ if attention_mask is None:
1112
+ attention_mask = torch.ones(
1113
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
1114
+ )
1115
+ elif 0 in attention_mask:
1116
+ padding_mask = attention_mask
1117
+
1118
+ if (
1119
+ padding_mask is not None
1120
+ and hasattr(self.config, "_flash_attn_2_enabled")
1121
+ and self.config._flash_attn_2_enabled
1122
+ ):
1123
+ is_padding_right = padding_mask[:, -1].sum().item() != batch_size
1124
+ if is_padding_right:
1125
+ raise ValueError(
1126
+ "You are attempting to perform batched generation with padding_side='right'"
1127
+ " this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
1128
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
1129
+ )
1130
+
1131
+ attention_mask = self._prepare_decoder_attention_mask(
1132
+ attention_mask,
1133
+ (batch_size, seq_length),
1134
+ inputs_embeds,
1135
+ past_key_values_length,
1136
+ sliding_window=self.config.sliding_window if hasattr(self.config, "sliding_window") else None,
1137
+ )
1138
+
1139
+ hidden_states = inputs_embeds
1140
+
1141
+ if self.gradient_checkpointing and self.training:
1142
+ if use_cache:
1143
+ logger.warning_once(
1144
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1145
+ )
1146
+ use_cache = False
1147
+
1148
+ # decoder layers
1149
+ all_hidden_states = () if output_hidden_states else None
1150
+ all_self_attns = () if output_attentions else None
1151
+ next_decoder_cache = () if use_cache else None
1152
+
1153
+ for idx, decoder_layer in enumerate(self.layers):
1154
+ if output_hidden_states:
1155
+ all_hidden_states += (hidden_states,)
1156
+
1157
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
1158
+
1159
+ if self.gradient_checkpointing and self.training:
1160
+
1161
+ def create_custom_forward(module):
1162
+ def custom_forward(*inputs):
1163
+ # None for past_key_value
1164
+ return module(*inputs, past_key_value, output_attentions, padding_mask=padding_mask)
1165
+
1166
+ return custom_forward
1167
+
1168
+ layer_outputs = torch.utils.checkpoint.checkpoint(
1169
+ create_custom_forward(decoder_layer),
1170
+ hidden_states,
1171
+ attention_mask,
1172
+ position_ids,
1173
+ )
1174
+ else:
1175
+ layer_outputs = decoder_layer(
1176
+ hidden_states,
1177
+ attention_mask=attention_mask,
1178
+ position_ids=position_ids,
1179
+ past_key_value=past_key_value,
1180
+ output_attentions=output_attentions,
1181
+ use_cache=use_cache,
1182
+ padding_mask=padding_mask,
1183
+ )
1184
+
1185
+ hidden_states = layer_outputs[0]
1186
+
1187
+ if use_cache:
1188
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
1189
+
1190
+ if output_attentions:
1191
+ all_self_attns += (layer_outputs[1],)
1192
+
1193
+ hidden_states = self.norm(hidden_states)
1194
+
1195
+ # add hidden states from the last decoder layer
1196
+ if output_hidden_states:
1197
+ all_hidden_states += (hidden_states,)
1198
+
1199
+ next_cache = next_decoder_cache if use_cache else None
1200
+ if not return_dict:
1201
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1202
+ return BaseModelOutputWithPast(
1203
+ last_hidden_state=hidden_states,
1204
+ past_key_values=next_cache,
1205
+ hidden_states=all_hidden_states,
1206
+ attentions=all_self_attns,
1207
+ )
1208
+
1209
+
1210
+ class MistralForCausalLM(MistralPreTrainedModel):
1211
+ _tied_weights_keys = ["lm_head.weight"]
1212
+
1213
+ def __init__(self, config):
1214
+ super().__init__(config)
1215
+ self.model = MistralModel(config)
1216
+ self.vocab_size = config.vocab_size
1217
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1218
+
1219
+ # Initialize weights and apply final processing
1220
+ self.post_init()
1221
+
1222
+ def get_input_embeddings(self):
1223
+ return self.model.embed_tokens
1224
+
1225
+ def set_input_embeddings(self, value):
1226
+ self.model.embed_tokens = value
1227
+
1228
+ def get_output_embeddings(self):
1229
+ return self.lm_head
1230
+
1231
+ def set_output_embeddings(self, new_embeddings):
1232
+ self.lm_head = new_embeddings
1233
+
1234
+ def set_decoder(self, decoder):
1235
+ self.model = decoder
1236
+
1237
+ def get_decoder(self):
1238
+ return self.model
1239
+
1240
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1241
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1242
+ def forward(
1243
+ self,
1244
+ input_ids: torch.LongTensor = None,
1245
+ attention_mask: Optional[torch.Tensor] = None,
1246
+ position_ids: Optional[torch.LongTensor] = None,
1247
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1248
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1249
+ labels: Optional[torch.LongTensor] = None,
1250
+ use_cache: Optional[bool] = None,
1251
+ output_attentions: Optional[bool] = None,
1252
+ output_hidden_states: Optional[bool] = None,
1253
+ return_dict: Optional[bool] = None,
1254
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1255
+ r"""
1256
+ Args:
1257
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1258
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1259
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1260
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1261