Crystalcareai commited on
Commit
55e9861
1 Parent(s): a00a1ac

Update config.json

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