BlackSamorez commited on
Commit
0a62b81
1 Parent(s): e44bce2

new dispatch

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