Iker commited on
Commit
b0e43b6
1 Parent(s): 347611c

Upload modeling_flash_llama.py

Browse files
Files changed (1) hide show
  1. modeling_flash_llama.py +1055 -0
modeling_flash_llama.py ADDED
@@ -0,0 +1,1055 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # From https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/modeling_flash_llama.py
3
+ # With seqlen fix from Alex Birch: https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/17
4
+ # With dtype Fix by Oscar Sainz
5
+ # With Beam Search Fix by Iker García-Ferrero
6
+
7
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
+ #
9
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
+ # and OPT implementations in this library. It has been modified from its
11
+ # original forms to accommodate minor architectural differences compared
12
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
+ #
14
+ # Licensed under the Apache License, Version 2.0 (the "License");
15
+ # you may not use this file except in compliance with the License.
16
+ # You may obtain a copy of the License at
17
+ #
18
+ # http://www.apache.org/licenses/LICENSE-2.0
19
+ #
20
+ # Unless required by applicable law or agreed to in writing, software
21
+ # distributed under the License is distributed on an "AS IS" BASIS,
22
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
+ # See the License for the specific language governing permissions and
24
+ # limitations under the License.
25
+ """ PyTorch LLaMA model."""
26
+ from typing import List, Optional, Tuple, Union
27
+
28
+ import torch
29
+ import torch.nn.functional as F
30
+ import torch.utils.checkpoint
31
+ from torch import nn
32
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
33
+
34
+ from transformers.activations import ACT2FN
35
+ from transformers.modeling_outputs import (
36
+ BaseModelOutputWithPast,
37
+ CausalLMOutputWithPast,
38
+ SequenceClassifierOutputWithPast,
39
+ )
40
+ from transformers.modeling_utils import PreTrainedModel
41
+ from transformers.models.llama.configuration_llama import LlamaConfig
42
+ from transformers.utils import (
43
+ add_start_docstrings,
44
+ add_start_docstrings_to_model_forward,
45
+ logging,
46
+ replace_return_docstrings,
47
+ )
48
+
49
+
50
+ try:
51
+ from flash_attn.bert_padding import pad_input, unpad_input
52
+ from flash_attn.flash_attn_interface import (
53
+ flash_attn_kvpacked_func,
54
+ flash_attn_varlen_kvpacked_func,
55
+ )
56
+
57
+ flash_attn_v2_installed = True
58
+ print(">>>> Flash Attention installed")
59
+ except ImportError:
60
+ flash_attn_v2_installed = False
61
+ raise ImportError("Please install Flash Attention: `pip install flash-attn --no-build-isolation`")
62
+
63
+ try:
64
+ from flash_attn.layers.rotary import apply_rotary_emb_func
65
+
66
+ flash_rope_installed = True
67
+ print(">>>> Flash RoPE installed")
68
+ except ImportError:
69
+ flash_rope_installed = False
70
+ raise ImportError(
71
+ "Please install RoPE kernels: `pip install"
72
+ " git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary`"
73
+ )
74
+
75
+ logger = logging.get_logger(__name__)
76
+
77
+ _CONFIG_FOR_DOC = "LlamaConfig"
78
+
79
+
80
+ # @torch.jit.script
81
+ def rmsnorm_func(hidden_states, weight, variance_epsilon):
82
+ input_dtype = hidden_states.dtype
83
+ hidden_states = hidden_states.to(torch.float32)
84
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
85
+ hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
86
+ return (weight * hidden_states).to(input_dtype)
87
+
88
+
89
+ class LlamaRMSNorm(nn.Module):
90
+ def __init__(self, hidden_size, eps=1e-6):
91
+ """
92
+ LlamaRMSNorm is equivalent to T5LayerNorm
93
+ """
94
+ super().__init__()
95
+ self.weight = nn.Parameter(torch.ones(hidden_size))
96
+ self.register_buffer(
97
+ "variance_epsilon",
98
+ torch.tensor(eps),
99
+ persistent=False,
100
+ )
101
+
102
+ def forward(self, hidden_states):
103
+ return rmsnorm_func(hidden_states, self.weight, self.variance_epsilon)
104
+
105
+
106
+ class FlashRotaryEmbedding(torch.nn.Module):
107
+ """
108
+ The rotary position embeddings from RoFormer_ (Su et. al).
109
+ A crucial insight from the method is that the query and keys are
110
+ transformed by rotation matrices which depend on the relative positions.
111
+ Other implementations are available in the Rotary Transformer repo_ and in
112
+ GPT-NeoX_, GPT-NeoX was an inspiration
113
+ .. _RoFormer: https://arxiv.org/abs/2104.09864
114
+ .. _repo: https://github.com/ZhuiyiTechnology/roformer
115
+ .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
116
+ If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
117
+ A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96
118
+ Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py
119
+ """
120
+
121
+ def __init__(
122
+ self,
123
+ dim: int,
124
+ base=10000.0,
125
+ interleaved=False,
126
+ scale_base=None,
127
+ scaling_factor=1.0,
128
+ pos_idx_in_fp32=True,
129
+ device=None,
130
+ ):
131
+ """
132
+ interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
133
+ of 1st half and 2nd half (GPT-NeoX style).
134
+ pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
135
+ otherwise they might be in lower precision.
136
+ This option was added because previously (before 2023-07-02), when we construct
137
+ the position indices, we use the dtype of self.inv_freq. In most cases this would
138
+ be fp32, but if the model is trained in pure bf16 (not mixed precision), then
139
+ self.inv_freq would be bf16, and the position indices are also in bf16.
140
+ Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
141
+ embeddings for some positions will coincide.
142
+ To maintain compatibility with models previously trained in pure bf16,
143
+ we add this option.
144
+ scaling_factor: RotaryEmbedding extended with linear scaling.
145
+ """
146
+ super().__init__()
147
+ self.dim = dim
148
+ self.base = float(base)
149
+ self.pos_idx_in_fp32 = pos_idx_in_fp32
150
+ # Generate and save the inverse frequency buffer (non trainable)
151
+ inv_freq = self._compute_inv_freq(device)
152
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
153
+ self.interleaved = interleaved
154
+ self.scale_base = scale_base
155
+ self.scaling_factor = scaling_factor
156
+ scale = (
157
+ (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
158
+ if scale_base is not None
159
+ else None
160
+ )
161
+ self.register_buffer("scale", scale)
162
+
163
+ self._seq_len_cached = 0
164
+ self._cos_cached = None
165
+ self._sin_cached = None
166
+ self._cos_k_cached = None
167
+ self._sin_k_cached = None
168
+
169
+ def _compute_inv_freq(self, device=None):
170
+ return 1 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
171
+
172
+ def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
173
+ # Reset the tables if the sequence length has changed,
174
+ # if we're on a new device (possibly due to tracing for instance),
175
+ # or if we're switching from inference mode to training
176
+ if (
177
+ seqlen > self._seq_len_cached
178
+ or self._cos_cached.device != device
179
+ or self._cos_cached.dtype != dtype
180
+ or (self.training and self._cos_cached.is_inference())
181
+ ):
182
+ self._seq_len_cached = seqlen
183
+ # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
184
+ # And the output of arange can be quite large, so bf16 would lose a lot of precision.
185
+ # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
186
+ if self.pos_idx_in_fp32:
187
+ t = torch.arange(seqlen, device=device, dtype=torch.float32)
188
+ t /= self.scaling_factor
189
+ # We want fp32 here as well since inv_freq will be multiplied with t, and the output
190
+ # will be large. Having it in bf16 will lose a lot of precision and cause the
191
+ # cos & sin output to change significantly.
192
+ # We want to recompute self.inv_freq if it was not loaded in fp32
193
+ if self.inv_freq.dtype != torch.float32:
194
+ inv_freq = self.inv_freq.to(torch.float32)
195
+ else:
196
+ inv_freq = self.inv_freq
197
+ else:
198
+ t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
199
+ t /= self.scaling_factor
200
+ inv_freq = self.inv_freq
201
+ # Don't do einsum, it converts fp32 to fp16 under AMP
202
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
203
+ freqs = torch.outer(t, inv_freq)
204
+ if self.scale is None:
205
+ self._cos_cached = torch.cos(freqs).to(dtype)
206
+ self._sin_cached = torch.sin(freqs).to(dtype)
207
+ else:
208
+ power = (
209
+ torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
210
+ ) / self.scale_base
211
+ scale = self.scale.to(device=power.device) ** power.unsqueeze(-1)
212
+ # We want the multiplication by scale to happen in fp32
213
+ self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
214
+ self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
215
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
216
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
217
+
218
+ def forward(self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
219
+ """
220
+ q: (batch, seqlen, nheads, headdim)
221
+ k: (batch, seqlen, nheads, headdim)
222
+ seqlen_offset: can be used in generation where the qkv being passed in is only the last
223
+ token in the batch.
224
+ """
225
+ self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype)
226
+ if self.scale is None:
227
+ return apply_rotary_emb_func(
228
+ q,
229
+ self._cos_cached[seqlen_offset:],
230
+ self._sin_cached[seqlen_offset:],
231
+ self.interleaved,
232
+ True, # inplace=True
233
+ ), apply_rotary_emb_func(
234
+ k,
235
+ self._cos_cached[seqlen_offset:],
236
+ self._sin_cached[seqlen_offset:],
237
+ self.interleaved,
238
+ True, # inplace=True
239
+ )
240
+ else:
241
+ assert False
242
+
243
+
244
+ class LlamaMLP(nn.Module):
245
+ def __init__(self, config):
246
+ super().__init__()
247
+ self.config = config
248
+ self.hidden_size = config.hidden_size
249
+ self.intermediate_size = config.intermediate_size
250
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
251
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
252
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
253
+ self.act_fn = ACT2FN[config.hidden_act]
254
+
255
+ def forward(self, x):
256
+ if self.config.pretraining_tp > 1:
257
+ slice = self.intermediate_size // self.config.pretraining_tp
258
+ gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
259
+ up_proj_slices = self.up_proj.weight.split(slice, dim=0)
260
+ down_proj_slices = self.down_proj.weight.split(slice, dim=1)
261
+
262
+ gate_proj = torch.cat(
263
+ [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
264
+ )
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
+ @torch.jit.script
279
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
280
+ """
281
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
282
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
283
+ """
284
+ batch, slen, _, num_key_value_heads, head_dim = hidden_states.shape
285
+ if n_rep == 1:
286
+ return hidden_states
287
+ hidden_states = hidden_states[:, :, :, :, None, :].expand(batch, slen, 2, num_key_value_heads, n_rep, head_dim)
288
+ return hidden_states.reshape(batch, slen, 2, num_key_value_heads * n_rep, head_dim)
289
+
290
+
291
+ class LlamaAttention(nn.Module):
292
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
293
+
294
+ def __init__(self, config: LlamaConfig):
295
+ super().__init__()
296
+ self.config = config
297
+ self.hidden_size = config.hidden_size
298
+ self.num_heads = config.num_attention_heads
299
+ self.head_dim = self.hidden_size // self.num_heads
300
+ self.num_key_value_heads = config.num_key_value_heads
301
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
302
+ self.max_position_embeddings = config.max_position_embeddings
303
+
304
+ if (self.head_dim * self.num_heads) != self.hidden_size:
305
+ raise ValueError(
306
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
307
+ f" and `num_heads`: {self.num_heads})."
308
+ )
309
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
310
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
311
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
312
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
313
+
314
+ self.register_buffer(
315
+ "norm_factor",
316
+ torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
317
+ persistent=False,
318
+ )
319
+
320
+ if self.config.rope_scaling is None:
321
+ scaling_factor = 1
322
+ else:
323
+ scaling_type = self.config.rope_scaling["type"]
324
+ scaling_factor = self.config.rope_scaling["factor"]
325
+ assert scaling_type == "linear"
326
+
327
+ self.rotary_emb = FlashRotaryEmbedding(
328
+ self.head_dim,
329
+ base=10000,
330
+ interleaved=False,
331
+ scaling_factor=scaling_factor,
332
+ )
333
+
334
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
335
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
336
+
337
+ def forward(
338
+ self,
339
+ hidden_states: torch.Tensor,
340
+ attention_mask: Optional[torch.Tensor] = None,
341
+ position_ids: Optional[torch.LongTensor] = None,
342
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
343
+ output_attentions: bool = False,
344
+ use_cache: bool = False,
345
+ is_padded_inputs: Optional[bool] = False,
346
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
347
+ bsz, q_len, h_size = hidden_states.size()
348
+
349
+ has_layer_past = past_key_value is not None
350
+
351
+ if has_layer_past:
352
+ past_kv = past_key_value[0]
353
+ past_len = past_key_value[1]
354
+ else:
355
+ past_len = 0
356
+
357
+ if self.config.pretraining_tp > 1:
358
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
359
+ query_slices = self.q_proj.weight.split(
360
+ (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
361
+ )
362
+ key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
363
+ value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
364
+
365
+ q = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
366
+ q = torch.cat(q, dim=-1)
367
+
368
+ k = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
369
+ k = torch.cat(k, dim=-1)
370
+
371
+ v = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
372
+ v = torch.cat(v, dim=-1)
373
+
374
+ else:
375
+ q = self.q_proj(hidden_states)
376
+ k = self.k_proj(hidden_states)
377
+ v = self.v_proj(hidden_states)
378
+
379
+ q = q.view(bsz, q_len, self.num_heads, self.head_dim)
380
+ k = k.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
381
+ v = v.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
382
+
383
+ q, k = self.rotary_emb(q, k, past_len)
384
+
385
+ kv = torch.stack([k, v], 2)
386
+ kv = repeat_kv(kv, self.num_key_value_groups)
387
+
388
+ # Make sure both are same dtype
389
+ if q.dtype != kv.dtype:
390
+ kv = kv.to(q.dtype)
391
+
392
+ # Cache QKV values
393
+ if has_layer_past:
394
+ new_len = past_len + q.size(1)
395
+ if new_len > past_kv.size(1):
396
+ past_kv = torch.cat(
397
+ [past_kv, torch.empty(bsz, 256, 2, kv.size(3), kv.size(4), dtype=kv.dtype, device=kv.device)], 1
398
+ )
399
+ past_kv[:, past_len:new_len] = kv
400
+ kv = past_kv[:, :new_len]
401
+ else:
402
+ past_kv = kv
403
+
404
+ past_key_value = (past_kv, past_len + q.size(1)) if use_cache else None
405
+
406
+ if is_padded_inputs:
407
+ # varlen, ignore padding tokens, efficient for large batch with many paddings
408
+
409
+ assert attention_mask is not None
410
+
411
+ unpadded_kv, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(kv, attention_mask)
412
+ unpadded_q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask[:, -q.size(1) :])
413
+
414
+ # Make sure both are same dtype
415
+ if unpadded_q.dtype != unpadded_kv.dtype:
416
+ unpadded_kv = unpadded_kv.to(unpadded_q.dtype)
417
+
418
+ attn_outputs = flash_attn_varlen_kvpacked_func(
419
+ unpadded_q,
420
+ unpadded_kv,
421
+ cu_seqlens_q,
422
+ cu_seqlens_k,
423
+ max_seqlen_q,
424
+ max_seqlen_k,
425
+ dropout_p=0.0,
426
+ softmax_scale=1.0 / self.norm_factor,
427
+ causal=(not has_layer_past),
428
+ return_attn_probs=output_attentions,
429
+ )
430
+
431
+ attn_output = attn_outputs[0] if output_attentions else attn_outputs
432
+ attn_output = pad_input(attn_output, indices_q, bsz, q_len).reshape(bsz, q_len, h_size)
433
+ attn_weights = attn_outputs[2] if output_attentions else None
434
+
435
+ else:
436
+ # Make sure both are same dtype
437
+ if q.dtype != kv.dtype:
438
+ kv = kv.to(q.dtype)
439
+
440
+ # no padding tokens, more efficient
441
+ attn_outputs = flash_attn_kvpacked_func(
442
+ q,
443
+ kv,
444
+ dropout_p=0.0,
445
+ softmax_scale=1.0 / self.norm_factor,
446
+ causal=(not has_layer_past),
447
+ return_attn_probs=output_attentions,
448
+ )
449
+
450
+ attn_output = attn_outputs[0] if output_attentions else attn_outputs
451
+ attn_output = attn_output.reshape(bsz, q_len, h_size)
452
+ attn_weights = attn_outputs[2] if output_attentions else None
453
+
454
+ if self.config.pretraining_tp > 1:
455
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
456
+ o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
457
+ attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
458
+ else:
459
+ attn_output = self.o_proj(attn_output)
460
+
461
+ if not output_attentions:
462
+ attn_weights = None
463
+
464
+ return attn_output, attn_weights, past_key_value
465
+
466
+
467
+ class LlamaDecoderLayer(nn.Module):
468
+ def __init__(self, config: LlamaConfig):
469
+ super().__init__()
470
+ self.hidden_size = config.hidden_size
471
+ self.self_attn = LlamaAttention(config=config)
472
+ self.mlp = LlamaMLP(config)
473
+ self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
474
+ self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
475
+
476
+ def forward(
477
+ self,
478
+ hidden_states: torch.Tensor,
479
+ attention_mask: Optional[torch.Tensor] = None,
480
+ position_ids: Optional[torch.LongTensor] = None,
481
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
482
+ is_padded_inputs: Optional[bool] = False,
483
+ output_attentions: Optional[bool] = False,
484
+ use_cache: Optional[bool] = False,
485
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
486
+ """
487
+ Args:
488
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
489
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
490
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
491
+ output_attentions (`bool`, *optional*):
492
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
493
+ returned tensors for more detail.
494
+ use_cache (`bool`, *optional*):
495
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
496
+ (see `past_key_values`).
497
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
498
+ """
499
+
500
+ residual = hidden_states
501
+
502
+ hidden_states = self.input_layernorm(hidden_states)
503
+
504
+ # Self Attention
505
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
506
+ hidden_states=hidden_states,
507
+ attention_mask=attention_mask,
508
+ position_ids=position_ids,
509
+ past_key_value=past_key_value,
510
+ output_attentions=output_attentions,
511
+ use_cache=use_cache,
512
+ is_padded_inputs=is_padded_inputs,
513
+ )
514
+ hidden_states = residual + hidden_states
515
+
516
+ # Fully Connected
517
+ residual = hidden_states
518
+ hidden_states = self.post_attention_layernorm(hidden_states)
519
+ hidden_states = self.mlp(hidden_states)
520
+ hidden_states = residual + hidden_states
521
+
522
+ outputs = (hidden_states,)
523
+
524
+ if output_attentions:
525
+ outputs += (self_attn_weights,)
526
+
527
+ if use_cache:
528
+ outputs += (present_key_value,)
529
+
530
+ return outputs
531
+
532
+
533
+ LLAMA_START_DOCSTRING = r"""
534
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
535
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
536
+ etc.)
537
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
538
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
539
+ and behavior.
540
+ Parameters:
541
+ config ([`LlamaConfig`]):
542
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
543
+ load the weights associated with the model, only the configuration. Check out the
544
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
545
+ """
546
+
547
+
548
+ @add_start_docstrings(
549
+ "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
550
+ LLAMA_START_DOCSTRING,
551
+ )
552
+ class LlamaPreTrainedModel(PreTrainedModel):
553
+ config_class = LlamaConfig
554
+ base_model_prefix = "model"
555
+ supports_gradient_checkpointing = True
556
+ _no_split_modules = ["LlamaDecoderLayer"]
557
+ _skip_keys_device_placement = "past_key_values"
558
+
559
+ def _init_weights(self, module):
560
+ std = self.config.initializer_range
561
+ if isinstance(module, nn.Linear):
562
+ module.weight.data.normal_(mean=0.0, std=std)
563
+ if module.bias is not None:
564
+ module.bias.data.zero_()
565
+ elif isinstance(module, nn.Embedding):
566
+ module.weight.data.normal_(mean=0.0, std=std)
567
+ if module.padding_idx is not None:
568
+ module.weight.data[module.padding_idx].zero_()
569
+
570
+ def _set_gradient_checkpointing(self, module, value=False):
571
+ if isinstance(module, LlamaModel):
572
+ module.gradient_checkpointing = value
573
+
574
+
575
+ LLAMA_INPUTS_DOCSTRING = r"""
576
+ Args:
577
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
578
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
579
+ it.
580
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
581
+ [`PreTrainedTokenizer.__call__`] for details.
582
+ [What are input IDs?](../glossary#input-ids)
583
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
584
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
585
+ - 1 for tokens that are **not masked**,
586
+ - 0 for tokens that are **masked**.
587
+ [What are attention masks?](../glossary#attention-mask)
588
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
589
+ [`PreTrainedTokenizer.__call__`] for details.
590
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
591
+ `past_key_values`).
592
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
593
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
594
+ information on the default strategy.
595
+ - 1 indicates the head is **not masked**,
596
+ - 0 indicates the head is **masked**.
597
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
598
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
599
+ config.n_positions - 1]`.
600
+ [What are position IDs?](../glossary#position-ids)
601
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
602
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
603
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
604
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
605
+ Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
606
+ blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
607
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
608
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
609
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
610
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
611
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
612
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
613
+ model's internal embedding lookup matrix.
614
+ use_cache (`bool`, *optional*):
615
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
616
+ `past_key_values`).
617
+ output_attentions (`bool`, *optional*):
618
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
619
+ tensors for more detail.
620
+ output_hidden_states (`bool`, *optional*):
621
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
622
+ more detail.
623
+ return_dict (`bool`, *optional*):
624
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
625
+ """
626
+
627
+
628
+ @add_start_docstrings(
629
+ "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
630
+ LLAMA_START_DOCSTRING,
631
+ )
632
+ class LlamaModel(LlamaPreTrainedModel):
633
+ """
634
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
635
+ Args:
636
+ config: LlamaConfig
637
+ """
638
+
639
+ def __init__(self, config: LlamaConfig):
640
+ super().__init__(config)
641
+ self.padding_idx = config.pad_token_id
642
+ self.vocab_size = config.vocab_size
643
+
644
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
645
+ self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
646
+ self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
647
+
648
+ self.gradient_checkpointing = False
649
+ # Initialize weights and apply final processing
650
+ self.post_init()
651
+
652
+ def get_input_embeddings(self):
653
+ return self.embed_tokens
654
+
655
+ def set_input_embeddings(self, value):
656
+ self.embed_tokens = value
657
+
658
+ @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
659
+ def forward(
660
+ self,
661
+ input_ids: torch.LongTensor = None,
662
+ attention_mask: Optional[torch.Tensor] = None,
663
+ position_ids: Optional[torch.LongTensor] = None,
664
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
665
+ inputs_embeds: Optional[torch.FloatTensor] = None,
666
+ use_cache: Optional[bool] = None,
667
+ output_attentions: Optional[bool] = None,
668
+ output_hidden_states: Optional[bool] = None,
669
+ return_dict: Optional[bool] = None,
670
+ is_padded_inputs: Optional[bool] = False,
671
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
672
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
673
+ output_hidden_states = (
674
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
675
+ )
676
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
677
+
678
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
679
+
680
+ # retrieve input_ids and inputs_embeds
681
+ if input_ids is not None and inputs_embeds is not None:
682
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
683
+ elif input_ids is not None:
684
+ batch_size, seq_length = input_ids.shape
685
+ elif inputs_embeds is not None:
686
+ batch_size, seq_length, _ = inputs_embeds.shape
687
+ else:
688
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
689
+
690
+ seq_length_with_past = seq_length
691
+ past_key_values_length = 0
692
+
693
+ if past_key_values is not None:
694
+ past_key_values_length = past_key_values[0][0].shape[2]
695
+ seq_length_with_past = seq_length_with_past + past_key_values_length
696
+
697
+ position_ids = None
698
+
699
+ if inputs_embeds is None:
700
+ inputs_embeds = self.embed_tokens(input_ids)
701
+
702
+ hidden_states = inputs_embeds
703
+
704
+ if self.gradient_checkpointing and self.training:
705
+ if use_cache:
706
+ logger.warning_once(
707
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
708
+ )
709
+ use_cache = False
710
+
711
+ # decoder layers
712
+ all_hidden_states = () if output_hidden_states else None
713
+ all_self_attns = () if output_attentions else None
714
+ next_decoder_cache = () if use_cache else None
715
+
716
+ for idx, decoder_layer in enumerate(self.layers):
717
+ if output_hidden_states:
718
+ all_hidden_states += (hidden_states,)
719
+
720
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
721
+
722
+ if self.gradient_checkpointing and self.training:
723
+
724
+ def create_custom_forward(module):
725
+ def custom_forward(*inputs):
726
+ # None for past_key_value
727
+ return module(*inputs, output_attentions, None)
728
+
729
+ return custom_forward
730
+
731
+ layer_outputs = torch.utils.checkpoint.checkpoint(
732
+ create_custom_forward(decoder_layer),
733
+ hidden_states,
734
+ attention_mask,
735
+ position_ids,
736
+ None,
737
+ is_padded_inputs,
738
+ )
739
+ else:
740
+ layer_outputs = decoder_layer(
741
+ hidden_states,
742
+ attention_mask=attention_mask,
743
+ position_ids=position_ids,
744
+ past_key_value=past_key_value,
745
+ output_attentions=output_attentions,
746
+ use_cache=use_cache,
747
+ is_padded_inputs=is_padded_inputs,
748
+ )
749
+
750
+ hidden_states = layer_outputs[0]
751
+
752
+ if use_cache:
753
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
754
+
755
+ if output_attentions:
756
+ all_self_attns += (layer_outputs[1],)
757
+
758
+ hidden_states = self.norm(hidden_states)
759
+
760
+ # add hidden states from the last decoder layer
761
+ if output_hidden_states:
762
+ all_hidden_states += (hidden_states,)
763
+
764
+ next_cache = next_decoder_cache if use_cache else None
765
+ if not return_dict:
766
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
767
+ return BaseModelOutputWithPast(
768
+ last_hidden_state=hidden_states,
769
+ past_key_values=next_cache,
770
+ hidden_states=all_hidden_states,
771
+ attentions=all_self_attns,
772
+ )
773
+
774
+
775
+ class LlamaForCausalLM(LlamaPreTrainedModel):
776
+ _tied_weights_keys = ["lm_head.weight"]
777
+
778
+ def __init__(self, config):
779
+ super().__init__(config)
780
+ self.model = LlamaModel(config)
781
+ self.vocab_size = config.vocab_size
782
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
783
+
784
+ # Initialize weights and apply final processing
785
+ self.post_init()
786
+
787
+ def get_input_embeddings(self):
788
+ return self.model.embed_tokens
789
+
790
+ def set_input_embeddings(self, value):
791
+ self.model.embed_tokens = value
792
+
793
+ def get_output_embeddings(self):
794
+ return self.lm_head
795
+
796
+ def set_output_embeddings(self, new_embeddings):
797
+ self.lm_head = new_embeddings
798
+
799
+ def set_decoder(self, decoder):
800
+ self.model = decoder
801
+
802
+ def get_decoder(self):
803
+ return self.model
804
+
805
+ @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
806
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
807
+ def forward(
808
+ self,
809
+ input_ids: torch.LongTensor = None,
810
+ attention_mask: Optional[torch.Tensor] = None,
811
+ position_ids: Optional[torch.LongTensor] = None,
812
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
813
+ inputs_embeds: Optional[torch.FloatTensor] = None,
814
+ labels: Optional[torch.LongTensor] = None,
815
+ use_cache: Optional[bool] = None,
816
+ output_attentions: Optional[bool] = None,
817
+ output_hidden_states: Optional[bool] = None,
818
+ return_dict: Optional[bool] = None,
819
+ is_padded_inputs: Optional[bool] = None,
820
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
821
+ r"""
822
+ Args:
823
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
824
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
825
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
826
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
827
+ Returns:
828
+ Example:
829
+ ```python
830
+ from transformers import AutoTokenizer, LlamaForCausalLM
831
+ model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
832
+ tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
833
+ prompt = "Hey, are you conscious? Can you talk to me?"
834
+ inputs = tokenizer(prompt, return_tensors="pt")
835
+ # Generate
836
+ generate_ids = model.generate(inputs.input_ids, max_length=30)
837
+ tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
838
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
839
+ ```"""
840
+
841
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
842
+ output_hidden_states = (
843
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
844
+ )
845
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
846
+
847
+ is_padded_inputs = (attention_mask is not None) and (not attention_mask.all().item())
848
+
849
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
850
+ outputs = self.model(
851
+ input_ids=input_ids,
852
+ attention_mask=attention_mask,
853
+ position_ids=position_ids,
854
+ past_key_values=past_key_values,
855
+ inputs_embeds=inputs_embeds,
856
+ use_cache=use_cache,
857
+ output_attentions=output_attentions,
858
+ output_hidden_states=output_hidden_states,
859
+ return_dict=return_dict,
860
+ is_padded_inputs=is_padded_inputs,
861
+ )
862
+
863
+ hidden_states = outputs[0]
864
+ if self.config.pretraining_tp > 1:
865
+ lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
866
+ logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
867
+ logits = torch.cat(logits, dim=-1)
868
+ else:
869
+ logits = self.lm_head(hidden_states)
870
+ logits = logits.float()
871
+
872
+ loss = None
873
+ if labels is not None:
874
+ # Shift so that tokens < n predict n
875
+ shift_logits = logits[..., :-1, :].contiguous()
876
+ shift_labels = labels[..., 1:].contiguous()
877
+ # Flatten the tokens
878
+ loss_fct = CrossEntropyLoss()
879
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
880
+ shift_labels = shift_labels.view(-1)
881
+ # Enable model parallelism
882
+ shift_labels = shift_labels.to(shift_logits.device)
883
+ loss = loss_fct(shift_logits, shift_labels)
884
+
885
+ if not return_dict:
886
+ output = (logits,) + outputs[1:]
887
+ return (loss,) + output if loss is not None else output
888
+
889
+ return CausalLMOutputWithPast(
890
+ loss=loss,
891
+ logits=logits,
892
+ past_key_values=outputs.past_key_values,
893
+ hidden_states=outputs.hidden_states,
894
+ attentions=outputs.attentions,
895
+ )
896
+
897
+ def prepare_inputs_for_generation(
898
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
899
+ ):
900
+ if past_key_values:
901
+ input_ids = input_ids[:, -1:]
902
+
903
+ position_ids = kwargs.get("position_ids", None)
904
+
905
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
906
+ if inputs_embeds is not None and past_key_values is None:
907
+ model_inputs = {"inputs_embeds": inputs_embeds}
908
+ else:
909
+ model_inputs = {"input_ids": input_ids}
910
+
911
+ model_inputs.update(
912
+ {
913
+ "position_ids": position_ids,
914
+ "past_key_values": past_key_values,
915
+ "use_cache": kwargs.get("use_cache"),
916
+ "attention_mask": attention_mask,
917
+ "is_padded_inputs": (attention_mask is not None) and (not attention_mask.all().item()),
918
+ }
919
+ )
920
+ return model_inputs
921
+
922
+ @staticmethod
923
+ def _reorder_cache(past_key_values, beam_idx):
924
+ reordered_past = ()
925
+ for layer_past in past_key_values:
926
+ reordered_past += (
927
+ tuple(
928
+ (
929
+ past_state.index_select(0, beam_idx.to(past_state.device))
930
+ if type(past_state) == torch.Tensor
931
+ else past_state # There is an int in the last layer, it is not supposed to be there,
932
+ # but this hack works to deal with it.
933
+ )
934
+ for past_state in layer_past
935
+ ),
936
+ )
937
+ return reordered_past
938
+
939
+
940
+ @add_start_docstrings(
941
+ """
942
+ The LLaMa Model transformer with a sequence classification head on top (linear layer).
943
+ [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
944
+ (e.g. GPT-2) do.
945
+ Since it does classification on the last token, it requires to know the position of the last token. If a
946
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
947
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
948
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
949
+ each row of the batch).
950
+ """,
951
+ LLAMA_START_DOCSTRING,
952
+ )
953
+ class LlamaForSequenceClassification(LlamaPreTrainedModel):
954
+ def __init__(self, config):
955
+ super().__init__(config)
956
+ self.num_labels = config.num_labels
957
+ self.model = LlamaModel(config)
958
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
959
+
960
+ # Initialize weights and apply final processing
961
+ self.post_init()
962
+
963
+ def get_input_embeddings(self):
964
+ return self.model.embed_tokens
965
+
966
+ def set_input_embeddings(self, value):
967
+ self.model.embed_tokens = value
968
+
969
+ @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
970
+ def forward(
971
+ self,
972
+ input_ids: torch.LongTensor = None,
973
+ attention_mask: Optional[torch.Tensor] = None,
974
+ position_ids: Optional[torch.LongTensor] = None,
975
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
976
+ inputs_embeds: Optional[torch.FloatTensor] = None,
977
+ labels: Optional[torch.LongTensor] = None,
978
+ use_cache: Optional[bool] = None,
979
+ output_attentions: Optional[bool] = None,
980
+ output_hidden_states: Optional[bool] = None,
981
+ return_dict: Optional[bool] = None,
982
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
983
+ r"""
984
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
985
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
986
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
987
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
988
+ """
989
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
990
+
991
+ transformer_outputs = self.model(
992
+ input_ids,
993
+ attention_mask=attention_mask,
994
+ position_ids=position_ids,
995
+ past_key_values=past_key_values,
996
+ inputs_embeds=inputs_embeds,
997
+ use_cache=use_cache,
998
+ output_attentions=output_attentions,
999
+ output_hidden_states=output_hidden_states,
1000
+ return_dict=return_dict,
1001
+ )
1002
+ hidden_states = transformer_outputs[0]
1003
+ logits = self.score(hidden_states)
1004
+
1005
+ if input_ids is not None:
1006
+ batch_size = input_ids.shape[0]
1007
+ else:
1008
+ batch_size = inputs_embeds.shape[0]
1009
+
1010
+ if self.config.pad_token_id is None and batch_size != 1:
1011
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1012
+ if self.config.pad_token_id is None:
1013
+ sequence_lengths = -1
1014
+ else:
1015
+ if input_ids is not None:
1016
+ sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
1017
+ else:
1018
+ sequence_lengths = -1
1019
+
1020
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1021
+
1022
+ loss = None
1023
+ if labels is not None:
1024
+ labels = labels.to(logits.device)
1025
+ if self.config.problem_type is None:
1026
+ if self.num_labels == 1:
1027
+ self.config.problem_type = "regression"
1028
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1029
+ self.config.problem_type = "single_label_classification"
1030
+ else:
1031
+ self.config.problem_type = "multi_label_classification"
1032
+
1033
+ if self.config.problem_type == "regression":
1034
+ loss_fct = MSELoss()
1035
+ if self.num_labels == 1:
1036
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1037
+ else:
1038
+ loss = loss_fct(pooled_logits, labels)
1039
+ elif self.config.problem_type == "single_label_classification":
1040
+ loss_fct = CrossEntropyLoss()
1041
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1042
+ elif self.config.problem_type == "multi_label_classification":
1043
+ loss_fct = BCEWithLogitsLoss()
1044
+ loss = loss_fct(pooled_logits, labels)
1045
+ if not return_dict:
1046
+ output = (pooled_logits,) + transformer_outputs[1:]
1047
+ return ((loss,) + output) if loss is not None else output
1048
+
1049
+ return SequenceClassifierOutputWithPast(
1050
+ loss=loss,
1051
+ logits=pooled_logits,
1052
+ past_key_values=transformer_outputs.past_key_values,
1053
+ hidden_states=transformer_outputs.hidden_states,
1054
+ attentions=transformer_outputs.attentions,
1055
+ )