x54-729 commited on
Commit
2e503ec
1 Parent(s): aafebc5

update modeling file to newest

Browse files
Files changed (2) hide show
  1. configuration_internlm2.py +42 -13
  2. modeling_internlm2.py +769 -352
configuration_internlm2.py CHANGED
@@ -44,9 +44,9 @@ class InternLM2Config(PretrainedConfig):
44
  intermediate_size (`int`, *optional*, defaults to 11008):
45
  Dimension of the MLP representations.
46
  num_hidden_layers (`int`, *optional*, defaults to 32):
47
- Number of hidden layers in the Transformer encoder.
48
  num_attention_heads (`int`, *optional*, defaults to 32):
49
- Number of attention heads for each attention layer in the Transformer encoder.
50
  num_key_value_heads (`int`, *optional*):
51
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
52
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
@@ -58,22 +58,42 @@ class InternLM2Config(PretrainedConfig):
58
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
59
  The non-linear activation function (function or string) in the decoder.
60
  max_position_embeddings (`int`, *optional*, defaults to 2048):
61
- The maximum sequence length that this model might ever be used with. Typically set this to something large
62
- just in case (e.g., 512 or 1024 or 2048).
63
  initializer_range (`float`, *optional*, defaults to 0.02):
64
  The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
65
- rms_norm_eps (`float`, *optional*, defaults to 1e-12):
66
  The epsilon used by the rms normalization layers.
67
  use_cache (`bool`, *optional*, defaults to `True`):
68
  Whether or not the model should return the last key/values attentions (not used by all models). Only
69
  relevant if `config.is_decoder=True`.
70
- tie_word_embeddings(`bool`, *optional*, defaults to `False`):
 
 
 
 
 
 
 
 
 
 
 
 
71
  Whether to tie weight embeddings
72
- Example:
73
-
 
 
 
 
 
 
 
 
74
  """
75
- model_type = "internlm2"
76
  _auto_class = "AutoConfig"
 
 
77
 
78
  def __init__( # pylint: disable=W0102
79
  self,
@@ -91,11 +111,12 @@ class InternLM2Config(PretrainedConfig):
91
  pad_token_id=0,
92
  bos_token_id=1,
93
  eos_token_id=2,
 
94
  tie_word_embeddings=False,
95
  bias=True,
96
  rope_theta=10000,
97
  rope_scaling=None,
98
- attn_implementation="eager",
99
  **kwargs,
100
  ):
101
  self.vocab_size = vocab_size
@@ -113,14 +134,15 @@ class InternLM2Config(PretrainedConfig):
113
  self.hidden_act = hidden_act
114
  self.initializer_range = initializer_range
115
  self.rms_norm_eps = rms_norm_eps
 
116
  self.use_cache = use_cache
117
  self.rope_theta = rope_theta
118
  self.rope_scaling = rope_scaling
119
  self._rope_scaling_validation()
120
-
121
  self.attn_implementation = attn_implementation
122
  if self.attn_implementation is None:
123
  self.attn_implementation = "eager"
 
124
  super().__init__(
125
  pad_token_id=pad_token_id,
126
  bos_token_id=bos_token_id,
@@ -147,5 +169,12 @@ class InternLM2Config(PretrainedConfig):
147
  raise ValueError(
148
  f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
149
  )
150
- if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor < 1.0:
151
- raise ValueError(f"`rope_scaling`'s factor field must be a float >= 1, got {rope_scaling_factor}")
 
 
 
 
 
 
 
 
44
  intermediate_size (`int`, *optional*, defaults to 11008):
45
  Dimension of the MLP representations.
46
  num_hidden_layers (`int`, *optional*, defaults to 32):
47
+ Number of hidden layers in the Transformer decoder.
48
  num_attention_heads (`int`, *optional*, defaults to 32):
49
+ Number of attention heads for each attention layer in the Transformer decoder.
50
  num_key_value_heads (`int`, *optional*):
51
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
52
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
 
58
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
59
  The non-linear activation function (function or string) in the decoder.
60
  max_position_embeddings (`int`, *optional*, defaults to 2048):
61
+ The maximum sequence length that this model might ever be used with. InternLM2 supports up to 32768 tokens.
 
62
  initializer_range (`float`, *optional*, defaults to 0.02):
63
  The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
64
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
65
  The epsilon used by the rms normalization layers.
66
  use_cache (`bool`, *optional*, defaults to `True`):
67
  Whether or not the model should return the last key/values attentions (not used by all models). Only
68
  relevant if `config.is_decoder=True`.
69
+ pad_token_id (`int`, *optional*):
70
+ Padding token id.
71
+ bos_token_id (`int`, *optional*, defaults to 1):
72
+ Beginning of stream token id.
73
+ eos_token_id (`int`, *optional*, defaults to 2):
74
+ End of stream token id.
75
+ pretraining_tp (`int`, *optional*, defaults to 1):
76
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
77
+ document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism)
78
+ to understand more about it. This value is necessary to ensure exact reproducibility
79
+ of the pretraining results. Please refer to [this
80
+ issue](https://github.com/pytorch/pytorch/issues/76232).
81
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
82
  Whether to tie weight embeddings
83
+ rope_theta (`float`, *optional*, defaults to 10000.0):
84
+ The base period of the RoPE embeddings.
85
+ rope_scaling (`Dict`, *optional*):
86
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
87
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
88
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
89
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
90
+ these scaling strategies behave:
91
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
92
+ experimental feature, subject to breaking API changes in future versions.
93
  """
 
94
  _auto_class = "AutoConfig"
95
+ model_type = "internlm2"
96
+ keys_to_ignore_at_inference = ["past_key_values"]
97
 
98
  def __init__( # pylint: disable=W0102
99
  self,
 
111
  pad_token_id=0,
112
  bos_token_id=1,
113
  eos_token_id=2,
114
+ pretraining_tp=1,
115
  tie_word_embeddings=False,
116
  bias=True,
117
  rope_theta=10000,
118
  rope_scaling=None,
119
+ attn_implementation=None,
120
  **kwargs,
121
  ):
122
  self.vocab_size = vocab_size
 
134
  self.hidden_act = hidden_act
135
  self.initializer_range = initializer_range
136
  self.rms_norm_eps = rms_norm_eps
137
+ self.pretraining_tp = pretraining_tp
138
  self.use_cache = use_cache
139
  self.rope_theta = rope_theta
140
  self.rope_scaling = rope_scaling
141
  self._rope_scaling_validation()
 
142
  self.attn_implementation = attn_implementation
143
  if self.attn_implementation is None:
144
  self.attn_implementation = "eager"
145
+
146
  super().__init__(
147
  pad_token_id=pad_token_id,
148
  bos_token_id=bos_token_id,
 
169
  raise ValueError(
170
  f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
171
  )
172
+ if (
173
+ rope_scaling_factor is None
174
+ or not isinstance(rope_scaling_factor, (float, int))
175
+ or rope_scaling_factor < 1.0
176
+ ):
177
+ raise ValueError(
178
+ f"`rope_scaling`'s factor field must be a number >= 1, got {rope_scaling_factor} "
179
+ f"of type {type(rope_scaling_factor)}"
180
+ )
modeling_internlm2.py CHANGED
@@ -13,11 +13,10 @@
13
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
  # See the License for the specific language governing permissions and
15
  # limitations under the License.
16
- """ PyTorch InternLM2 model."""
17
  import math
18
  import queue
19
  import threading
20
- import warnings
21
  from typing import List, Optional, Tuple, Union
22
 
23
  import torch
@@ -27,49 +26,54 @@ from einops import rearrange
27
  from torch import nn
28
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
29
  from transformers.activations import ACT2FN
 
 
30
  from transformers.modeling_outputs import (
31
  BaseModelOutputWithPast,
32
  CausalLMOutputWithPast,
 
33
  SequenceClassifierOutputWithPast,
 
34
  )
35
  from transformers.modeling_utils import PreTrainedModel
 
36
  from transformers.utils import (
37
  add_start_docstrings,
38
  add_start_docstrings_to_model_forward,
 
39
  logging,
40
  replace_return_docstrings,
41
  )
42
 
43
  try:
44
  from transformers.generation.streamers import BaseStreamer
45
- except: # noqa # pylint: disable=bare-except
46
  BaseStreamer = None
47
 
48
  from .configuration_internlm2 import InternLM2Config
49
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  logger = logging.get_logger(__name__)
51
 
52
  _CONFIG_FOR_DOC = "InternLM2Config"
53
 
54
- flash_attn_func, flash_attn_varlen_func = None, None
55
- pad_input, index_first_axis, unpad_input = None, None, None
56
- def _import_flash_attn():
57
- global flash_attn_func, flash_attn_varlen_func
58
- global pad_input, index_first_axis, unpad_input
59
- try:
60
- from flash_attn import flash_attn_func as _flash_attn_func, flash_attn_varlen_func as _flash_attn_varlen_func
61
- from flash_attn.bert_padding import pad_input as _pad_input, index_first_axis as _index_first_axis, unpad_input as _unpad_input
62
- flash_attn_func, flash_attn_varlen_func = _flash_attn_func, _flash_attn_varlen_func
63
- pad_input, index_first_axis, unpad_input = _pad_input, _index_first_axis, _unpad_input
64
- except ImportError:
65
- raise ImportError("flash_attn is not installed.")
66
-
67
- # Copied from transformers.models.llama.modeling_llama._get_unpad_data
68
  def _get_unpad_data(attention_mask):
69
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
70
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
71
  max_seqlen_in_batch = seqlens_in_batch.max().item()
72
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
73
  return (
74
  indices,
75
  cu_seqlens,
@@ -77,45 +81,10 @@ def _get_unpad_data(attention_mask):
77
  )
78
 
79
 
80
- # Copied from transformers.models.bart.modeling_bart._make_causal_mask
81
- def _make_causal_mask(
82
- input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
83
- ):
84
- """
85
- Make causal mask used for bi-directional self-attention.
86
- """
87
- bsz, tgt_len = input_ids_shape
88
- mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
89
- mask_cond = torch.arange(mask.size(-1), device=device)
90
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
91
- mask = mask.to(dtype)
92
-
93
- if past_key_values_length > 0:
94
- mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
95
- return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
96
-
97
-
98
- # Copied from transformers.models.bart.modeling_bart._expand_mask
99
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
100
- """
101
- Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
102
- """
103
- bsz, src_len = mask.size()
104
- tgt_len = tgt_len if tgt_len is not None else src_len
105
-
106
- expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
107
-
108
- inverted_mask = 1.0 - expanded_mask
109
-
110
- return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
111
-
112
-
113
- # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->InternLM2
114
  class InternLM2RMSNorm(nn.Module):
 
 
115
  def __init__(self, hidden_size, eps=1e-6):
116
- """
117
- InternLM2RMSNorm is equivalent to T5LayerNorm
118
- """
119
  super().__init__()
120
  self.weight = nn.Parameter(torch.ones(hidden_size))
121
  self.variance_epsilon = eps
@@ -128,93 +97,68 @@ class InternLM2RMSNorm(nn.Module):
128
  return self.weight * hidden_states.to(input_dtype)
129
 
130
 
131
- # Copied from transformers.model.llama.modeling_llama.LlamaRotaryEmbedding with Llama->InternLM2
 
 
132
  class InternLM2RotaryEmbedding(nn.Module):
133
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
134
- super().__init__()
135
 
 
 
 
136
  self.dim = dim
137
  self.max_position_embeddings = max_position_embeddings
138
  self.base = base
139
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
140
  self.register_buffer("inv_freq", inv_freq, persistent=False)
 
 
141
 
142
- # Build here to make `torch.jit.trace` work.
143
- self._set_cos_sin_cache(
144
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
145
- )
146
-
147
- def _set_cos_sin_cache(self, seq_len, device, dtype):
148
- self.max_seq_len_cached = seq_len
149
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
150
-
151
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
152
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
153
- emb = torch.cat((freqs, freqs), dim=-1)
154
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
155
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
156
-
157
- def forward(self, x, seq_len=None):
158
  # x: [bs, num_attention_heads, seq_len, head_size]
159
- if seq_len > self.max_seq_len_cached:
160
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=torch.float32)
161
-
162
- return (
163
- self.cos_cached[:seq_len].to(dtype=x.dtype),
164
- self.sin_cached[:seq_len].to(dtype=x.dtype),
165
- )
 
 
 
 
 
166
 
167
 
168
- # Copied from transformers.model.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->InternLM2
169
  class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
170
  """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
171
 
172
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
173
- self.scaling_factor = scaling_factor
174
- super().__init__(dim, max_position_embeddings, base, device)
175
-
176
- def _set_cos_sin_cache(self, seq_len, device, dtype):
177
- self.max_seq_len_cached = seq_len
178
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
179
- t = t / self.scaling_factor
180
 
181
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
182
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
183
- emb = torch.cat((freqs, freqs), dim=-1)
184
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
185
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
186
 
187
-
188
- # Copied from transformers.model.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->InternLM2
189
  class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
190
  """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
191
- Credits to the Reddit users /u/bloc97 and /u/emozilla.
192
- """
193
-
194
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
195
- self.scaling_factor = scaling_factor
196
- super().__init__(dim, max_position_embeddings, base, device)
197
-
198
- def _set_cos_sin_cache(self, seq_len, device, dtype):
199
- self.max_seq_len_cached = seq_len
200
 
 
 
 
201
  if seq_len > self.max_position_embeddings:
202
  base = self.base * (
203
  (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
204
  ) ** (self.dim / (self.dim - 2))
205
- inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
206
- self.register_buffer("inv_freq", inv_freq, persistent=False)
207
-
208
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
209
 
210
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
211
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
212
- emb = torch.cat((freqs, freqs), dim=-1)
213
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
214
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
215
 
216
 
217
- # Copied from transformers.model.llama.modeling_llama.rotate_half
218
  def rotate_half(x):
219
  """Rotates half the hidden dims of the input."""
220
  x1 = x[..., : x.shape[-1] // 2]
@@ -222,17 +166,36 @@ def rotate_half(x):
222
  return torch.cat((-x2, x1), dim=-1)
223
 
224
 
225
- # Copied from transformers.model.llama.modeling_llama.apply_rotary_pos_emb
226
- def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
227
- """Applies Rotary Position Embedding to the query and key tensors."""
228
- cos = cos[position_ids].unsqueeze(unsqueeze_dim)
229
- sin = sin[position_ids].unsqueeze(unsqueeze_dim)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
230
  q_embed = (q * cos) + (rotate_half(q) * sin)
231
  k_embed = (k * cos) + (rotate_half(k) * sin)
232
  return q_embed, k_embed
233
 
234
 
235
  class InternLM2MLP(nn.Module):
 
 
236
  def __init__(self, config):
237
  super().__init__()
238
  self.config = config
@@ -249,7 +212,6 @@ class InternLM2MLP(nn.Module):
249
  return down_proj
250
 
251
 
252
- # Copied from transformers.model.llama.modeling_llama.repeat_kv
253
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
254
  """
255
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -262,19 +224,27 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
262
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
263
 
264
 
265
- # Modified from transformers.model.llama.modeling_llama.LlamaAttention
266
  class InternLM2Attention(nn.Module):
267
  """Multi-headed attention from 'Attention Is All You Need' paper"""
268
 
269
- def __init__(self, config: InternLM2Config):
270
  super().__init__()
271
  self.config = config
 
 
 
 
 
 
 
 
272
  self.hidden_size = config.hidden_size
273
  self.num_heads = config.num_attention_heads
274
  self.head_dim = self.hidden_size // self.num_heads
275
  self.num_key_value_heads = config.num_key_value_heads
276
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
277
  self.max_position_embeddings = config.max_position_embeddings
 
278
  self.is_causal = True
279
 
280
  if (self.head_dim * self.num_heads) != self.hidden_size:
@@ -288,8 +258,8 @@ class InternLM2Attention(nn.Module):
288
  (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
289
  bias=config.bias,
290
  )
291
-
292
  self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
 
293
  self._init_rope()
294
 
295
  def _init_rope(self):
@@ -297,51 +267,49 @@ class InternLM2Attention(nn.Module):
297
  self.rotary_emb = InternLM2RotaryEmbedding(
298
  self.head_dim,
299
  max_position_embeddings=self.max_position_embeddings,
300
- base=self.config.rope_theta,
301
  )
302
  else:
303
  scaling_type = self.config.rope_scaling["type"]
304
  scaling_factor = self.config.rope_scaling["factor"]
305
- if scaling_type == "dynamic":
306
- self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
307
  self.head_dim,
308
  max_position_embeddings=self.max_position_embeddings,
309
- base=self.config.rope_theta,
310
  scaling_factor=scaling_factor,
 
311
  )
312
- elif scaling_type == "linear":
313
- self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
314
  self.head_dim,
315
  max_position_embeddings=self.max_position_embeddings,
316
- base=self.config.rope_theta,
317
  scaling_factor=scaling_factor,
 
318
  )
319
  else:
320
- raise ValueError("Currently we only support rotary embedding's type being 'dynamic' or 'linear'.")
321
- return self.rotary_emb
322
-
323
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
324
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
325
 
326
  def forward(
327
  self,
328
  hidden_states: torch.Tensor,
329
  attention_mask: Optional[torch.Tensor] = None,
330
  position_ids: Optional[torch.LongTensor] = None,
331
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
332
  output_attentions: bool = False,
333
- use_cache: bool = False,
334
- **kwargs,
335
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
336
- if "padding_mask" in kwargs:
337
- warnings.warn(
338
- "Passing `padding_mask` is deprecated and will be removed in v4.37. "
339
- "Please make sure use `attention_mask` instead.`"
340
- )
341
-
342
  bsz, q_len, _ = hidden_states.size()
343
 
344
- qkv_states = self.wqkv(hidden_states)
 
 
 
 
 
 
 
 
345
 
346
  qkv_states = rearrange(
347
  qkv_states,
@@ -351,44 +319,26 @@ class InternLM2Attention(nn.Module):
351
  )
352
 
353
  query_states = qkv_states[..., : self.num_key_value_groups, :]
354
- query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
355
- key_states = qkv_states[..., -2, :]
356
- value_states = qkv_states[..., -1, :]
357
 
358
- query_states = query_states.transpose(1, 2)
359
- key_states = key_states.transpose(1, 2)
360
- value_states = value_states.transpose(1, 2)
361
-
362
- kv_seq_len = key_states.shape[-2]
363
- if past_key_value is not None:
364
- kv_seq_len += past_key_value[0].shape[-2]
365
- cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
366
  query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
367
 
368
  if past_key_value is not None:
369
- # reuse k, v, self_attention
370
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
371
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
372
-
373
- past_key_value = (key_states, value_states) if use_cache else None
374
 
375
  key_states = repeat_kv(key_states, self.num_key_value_groups)
376
  value_states = repeat_kv(value_states, self.num_key_value_groups)
377
 
378
  attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
379
 
380
- if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
381
- raise ValueError(
382
- f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
383
- f" {attn_weights.size()}"
384
- )
385
-
386
- if attention_mask is not None:
387
- if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
388
- raise ValueError(
389
- f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
390
- )
391
- attn_weights = attn_weights + attention_mask
392
 
393
  # upcast attention to fp32
394
  attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
@@ -401,9 +351,20 @@ class InternLM2Attention(nn.Module):
401
  )
402
 
403
  attn_output = attn_output.transpose(1, 2).contiguous()
 
404
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
405
 
406
- attn_output = self.wo(attn_output)
 
 
 
 
 
 
 
 
 
 
407
 
408
  if not output_attentions:
409
  attn_weights = None
@@ -411,7 +372,6 @@ class InternLM2Attention(nn.Module):
411
  return attn_output, attn_weights, past_key_value
412
 
413
 
414
- # Modified from transformers.model.llama.modeling_llama.InternLM2FlashAttention2
415
  class InternLM2FlashAttention2(InternLM2Attention):
416
  """
417
  InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
@@ -419,26 +379,34 @@ class InternLM2FlashAttention2(InternLM2Attention):
419
  flash attention and deal with padding tokens in case the input contains any of them.
420
  """
421
 
 
 
 
 
 
 
 
 
 
 
 
422
  def forward(
423
  self,
424
  hidden_states: torch.Tensor,
425
  attention_mask: Optional[torch.LongTensor] = None,
426
  position_ids: Optional[torch.LongTensor] = None,
427
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
428
  output_attentions: bool = False,
429
  use_cache: bool = False,
430
- **kwargs,
431
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
432
- # InternLM2FlashAttention2 attention does not support output_attentions
433
- if "padding_mask" in kwargs:
434
- warnings.warn(
435
- "Passing `padding_mask` is deprecated and will be removed in v4.37. "
436
- "Please make sure use `attention_mask` instead.`"
437
  )
438
 
439
- # overwrite attention_mask with padding_mask
440
- attention_mask = kwargs.pop("padding_mask")
441
-
442
  output_attentions = False
443
 
444
  bsz, q_len, _ = hidden_states.size()
@@ -461,35 +429,61 @@ class InternLM2FlashAttention2(InternLM2Attention):
461
  key_states = key_states.transpose(1, 2)
462
  value_states = value_states.transpose(1, 2)
463
 
464
- kv_seq_len = key_states.shape[-2]
465
- if past_key_value is not None:
466
- kv_seq_len += past_key_value[0].shape[-2]
467
-
468
- cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
469
-
470
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
471
 
472
  if past_key_value is not None:
473
- # reuse k, v, self_attention
474
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
475
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
476
-
477
- past_key_value = (key_states, value_states) if use_cache else None
478
 
 
 
 
479
  query_states = query_states.transpose(1, 2)
480
  key_states = key_states.transpose(1, 2)
481
  value_states = value_states.transpose(1, 2)
482
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
483
  attn_output = self._flash_attention_forward(
484
- query_states, key_states, value_states, attention_mask, q_len
485
  )
 
486
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
487
  attn_output = self.wo(attn_output)
488
 
489
  if not output_attentions:
490
  attn_weights = None
491
 
492
- return attn_output, attn_weights, past_key_value
493
 
494
  def _flash_attention_forward(
495
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
@@ -508,23 +502,29 @@ class InternLM2FlashAttention2(InternLM2Attention):
508
  attention_mask (`torch.Tensor`):
509
  The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
510
  position of padding tokens and 1 for the position of non-padding tokens.
511
- dropout (`int`, *optional*):
512
  Attention dropout
513
  softmax_scale (`float`, *optional*):
514
  The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
515
  """
 
 
 
 
 
 
 
516
  # Contains at least one padding token in the sequence
517
- causal = self.is_causal and query_length != 1
518
  if attention_mask is not None:
519
  batch_size = query_states.shape[0]
520
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._unpad_input(
521
  query_states, key_states, value_states, attention_mask, query_length
522
  )
523
 
524
  cu_seqlens_q, cu_seqlens_k = cu_seq_lens
525
  max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
526
 
527
- attn_output_unpad = flash_attn_varlen_func(
528
  query_states,
529
  key_states,
530
  value_states,
@@ -537,27 +537,26 @@ class InternLM2FlashAttention2(InternLM2Attention):
537
  causal=causal,
538
  )
539
 
540
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
541
  else:
542
- attn_output = flash_attn_func(
543
  query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
544
  )
545
 
546
  return attn_output
547
 
548
- def _unpad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
549
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
550
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
551
 
552
- key_layer = index_first_axis(
553
  key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
554
  )
555
- value_layer = index_first_axis(
556
  value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
557
  )
558
-
559
  if query_length == kv_seq_len:
560
- query_layer = index_first_axis(
561
  query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
562
  )
563
  cu_seqlens_q = cu_seqlens_k
@@ -573,29 +572,139 @@ class InternLM2FlashAttention2(InternLM2Attention):
573
  else:
574
  # The -q_len: slice assumes left padding.
575
  attention_mask = attention_mask[:, -query_length:]
576
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
 
 
577
 
578
  return (
579
  query_layer,
580
  key_layer,
581
  value_layer,
582
- indices_q.to(torch.int64),
583
  (cu_seqlens_q, cu_seqlens_k),
584
  (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
585
  )
586
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
587
  INTERNLM2_ATTENTION_CLASSES = {
588
  "eager": InternLM2Attention,
589
  "flash_attention_2": InternLM2FlashAttention2,
 
590
  }
591
 
592
- # Modified from transformers.model.llama.modeling_llama.LlamaDecoderLayer
 
593
  class InternLM2DecoderLayer(nn.Module):
594
- def __init__(self, config: InternLM2Config):
 
 
595
  super().__init__()
596
  self.hidden_size = config.hidden_size
 
597
 
598
- self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config)
599
 
600
  self.feed_forward = InternLM2MLP(config)
601
  self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -606,10 +715,10 @@ class InternLM2DecoderLayer(nn.Module):
606
  hidden_states: torch.Tensor,
607
  attention_mask: Optional[torch.Tensor] = None,
608
  position_ids: Optional[torch.LongTensor] = None,
609
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
610
  output_attentions: Optional[bool] = False,
611
  use_cache: Optional[bool] = False,
612
- **kwargs,
613
  ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
614
  """
615
  Args:
@@ -625,12 +734,6 @@ class InternLM2DecoderLayer(nn.Module):
625
  (see `past_key_values`).
626
  past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
627
  """
628
- if "padding_mask" in kwargs:
629
- warnings.warn(
630
- "Passing `padding_mask` is deprecated and will be removed in v4.37. "
631
- "Please make sure use `attention_mask` instead.`"
632
- )
633
-
634
  residual = hidden_states
635
 
636
  hidden_states = self.attention_norm(hidden_states)
@@ -643,7 +746,7 @@ class InternLM2DecoderLayer(nn.Module):
643
  past_key_value=past_key_value,
644
  output_attentions=output_attentions,
645
  use_cache=use_cache,
646
- **kwargs,
647
  )
648
  hidden_states = residual + hidden_states
649
 
@@ -687,11 +790,20 @@ InternLM2_START_DOCSTRING = r"""
687
  InternLM2_START_DOCSTRING,
688
  )
689
  class InternLM2PreTrainedModel(PreTrainedModel):
 
 
 
 
690
  config_class = InternLM2Config
691
  base_model_prefix = "model"
692
  supports_gradient_checkpointing = True
693
  _no_split_modules = ["InternLM2DecoderLayer"]
694
- _skip_keys_device_placement = "past_key_values"
 
 
 
 
 
695
 
696
  def _init_weights(self, module):
697
  std = self.config.initializer_range
@@ -740,14 +852,19 @@ InternLM2_INPUTS_DOCSTRING = r"""
740
  config.n_positions - 1]`.
741
 
742
  [What are position IDs?](../glossary#position-ids)
743
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or
744
- when `config.use_cache=True`):
745
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
746
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
747
- `(batch_size, num_heads, decoder_sequence_length, embed_size_per_head)`.
 
 
 
 
 
748
 
749
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
750
- blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
751
 
752
  If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
753
  have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
@@ -767,10 +884,14 @@ InternLM2_INPUTS_DOCSTRING = r"""
767
  more detail.
768
  return_dict (`bool`, *optional*):
769
  Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
 
 
 
 
770
  """
771
 
772
 
773
- # Modified from transformers.model.llama.modeling_llama.LlamaModel
774
  @add_start_docstrings(
775
  "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
776
  InternLM2_START_DOCSTRING,
@@ -793,7 +914,9 @@ class InternLM2Model(InternLM2PreTrainedModel):
793
 
794
  self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
795
 
796
- self.layers = nn.ModuleList([InternLM2DecoderLayer(config) for _ in range(config.num_hidden_layers)])
 
 
797
  self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
798
 
799
  self.gradient_checkpointing = False
@@ -806,142 +929,96 @@ class InternLM2Model(InternLM2PreTrainedModel):
806
  def set_input_embeddings(self, value):
807
  self.tok_embeddings = value
808
 
809
- def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
810
- # create causal mask
811
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
812
- combined_attention_mask = None
813
- if input_shape[-1] > 1:
814
- combined_attention_mask = _make_causal_mask(
815
- input_shape,
816
- inputs_embeds.dtype,
817
- device=inputs_embeds.device,
818
- past_key_values_length=past_key_values_length,
819
- )
820
-
821
- if attention_mask is not None:
822
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
823
- expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
824
- inputs_embeds.device
825
- )
826
- combined_attention_mask = (
827
- expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
828
- )
829
-
830
- return combined_attention_mask
831
-
832
  @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
833
  def forward(
834
  self,
835
  input_ids: torch.LongTensor = None,
836
  attention_mask: Optional[torch.Tensor] = None,
837
  position_ids: Optional[torch.LongTensor] = None,
838
- past_key_values: Optional[List[torch.FloatTensor]] = None,
839
  inputs_embeds: Optional[torch.FloatTensor] = None,
840
  use_cache: Optional[bool] = None,
841
  output_attentions: Optional[bool] = None,
842
  output_hidden_states: Optional[bool] = None,
843
  return_dict: Optional[bool] = None,
 
844
  ) -> Union[Tuple, BaseModelOutputWithPast]:
845
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
846
  output_hidden_states = (
847
  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
848
  )
849
  use_cache = use_cache if use_cache is not None else self.config.use_cache
850
-
851
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
852
 
853
- if self.config.attn_implementation == "flash_attention_2":
854
- _import_flash_attn()
855
-
856
- # retrieve input_ids and inputs_embeds
857
- if input_ids is not None and inputs_embeds is not None:
858
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
859
- elif input_ids is not None:
860
- batch_size, seq_length = input_ids.shape[:2]
861
- elif inputs_embeds is not None:
862
- batch_size, seq_length = inputs_embeds.shape[:2]
863
- else:
864
- raise ValueError("You have to specify either input_ids or inputs_embeds")
865
-
866
- seq_length_with_past = seq_length
867
- past_key_values_length = 0
868
- if past_key_values is not None:
869
- past_key_values_length = past_key_values[0][0].shape[2]
870
- seq_length_with_past = seq_length_with_past + past_key_values_length
871
 
872
- if position_ids is None:
873
- device = input_ids.device if input_ids is not None else inputs_embeds.device
874
- position_ids = torch.arange(
875
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
876
  )
877
- position_ids = position_ids.unsqueeze(0)
878
 
879
  if inputs_embeds is None:
880
  inputs_embeds = self.tok_embeddings(input_ids)
881
 
882
- if self.config.attn_implementation == "flash_attention_2":
883
- # 2d mask is passed through the layers
884
- attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
885
- else:
886
- if attention_mask is None:
887
- attention_mask = torch.ones(
888
- (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
889
- )
890
- attention_mask = self._prepare_decoder_attention_mask(
891
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
892
  )
 
 
 
 
 
 
893
 
894
  # embed positions
895
  hidden_states = inputs_embeds
896
 
897
- if self.gradient_checkpointing and self.training:
898
- if use_cache:
899
- logger.warning_once(
900
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
901
- )
902
- use_cache = False
903
-
904
  # decoder layers
905
  all_hidden_states = () if output_hidden_states else None
906
  all_self_attns = () if output_attentions else None
907
- next_decoder_cache = () if use_cache else None
908
 
909
- for idx, decoder_layer in enumerate(self.layers):
910
  if output_hidden_states:
911
  all_hidden_states += (hidden_states,)
912
 
913
- past_key_value = past_key_values[idx] if past_key_values is not None else None
914
-
915
  if self.gradient_checkpointing and self.training:
916
-
917
- def create_custom_forward(module):
918
- def custom_forward(*inputs):
919
- # None for past_key_value
920
- return module(*inputs, output_attentions, None)
921
-
922
- return custom_forward
923
-
924
- layer_outputs = torch.utils.checkpoint.checkpoint(
925
- create_custom_forward(decoder_layer),
926
  hidden_states,
927
- attention_mask,
928
  position_ids,
929
- None,
 
 
 
930
  )
931
  else:
932
  layer_outputs = decoder_layer(
933
  hidden_states,
934
- attention_mask=attention_mask,
935
  position_ids=position_ids,
936
- past_key_value=past_key_value,
937
  output_attentions=output_attentions,
938
  use_cache=use_cache,
 
939
  )
940
 
941
  hidden_states = layer_outputs[0]
942
 
943
  if use_cache:
944
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
945
 
946
  if output_attentions:
947
  all_self_attns += (layer_outputs[1],)
@@ -953,6 +1030,9 @@ class InternLM2Model(InternLM2PreTrainedModel):
953
  all_hidden_states += (hidden_states,)
954
 
955
  next_cache = next_decoder_cache if use_cache else None
 
 
 
956
  if not return_dict:
957
  return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
958
  return BaseModelOutputWithPast(
@@ -962,11 +1042,95 @@ class InternLM2Model(InternLM2PreTrainedModel):
962
  attentions=all_self_attns,
963
  )
964
 
 
 
 
 
 
 
 
 
 
 
 
 
 
965
 
966
- # Modified from transformers.model.llama.modeling_llama.LlamaForCausalLM
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
967
  class InternLM2ForCausalLM(InternLM2PreTrainedModel):
968
- _auto_class = "AutoModelForCausalLM"
969
 
 
970
  _tied_weights_keys = ["output.weight"]
971
 
972
  def __init__(self, config):
@@ -1003,13 +1167,14 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1003
  input_ids: torch.LongTensor = None,
1004
  attention_mask: Optional[torch.Tensor] = None,
1005
  position_ids: Optional[torch.LongTensor] = None,
1006
- past_key_values: Optional[List[torch.FloatTensor]] = None,
1007
  inputs_embeds: Optional[torch.FloatTensor] = None,
1008
  labels: Optional[torch.LongTensor] = None,
1009
  use_cache: Optional[bool] = None,
1010
  output_attentions: Optional[bool] = None,
1011
  output_hidden_states: Optional[bool] = None,
1012
  return_dict: Optional[bool] = None,
 
1013
  ) -> Union[Tuple, CausalLMOutputWithPast]:
1014
  r"""
1015
  Args:
@@ -1025,8 +1190,8 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1025
  ```python
1026
  >>> from transformers import AutoTokenizer, InternLM2ForCausalLM
1027
 
1028
- >>> model = InternLM2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1029
- >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1030
 
1031
  >>> prompt = "Hey, are you conscious? Can you talk to me?"
1032
  >>> inputs = tokenizer(prompt, return_tensors="pt")
@@ -1054,10 +1219,19 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1054
  output_attentions=output_attentions,
1055
  output_hidden_states=output_hidden_states,
1056
  return_dict=return_dict,
 
1057
  )
1058
 
1059
  hidden_states = outputs[0]
1060
- logits = self.output(hidden_states)
 
 
 
 
 
 
 
 
1061
  logits = logits.float()
1062
 
1063
  loss = None
@@ -1086,19 +1260,48 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1086
  )
1087
 
1088
  def prepare_inputs_for_generation(
1089
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
 
 
 
 
 
 
 
1090
  ):
 
1091
  if past_key_values is not None:
1092
- past_length = past_key_values[0][0].shape[2]
1093
-
1094
- # Some generation methods already pass only the last input ID
1095
- if input_ids.shape[1] > past_length:
1096
- remove_prefix_length = past_length
 
 
 
 
1097
  else:
1098
- # Default to old behavior: keep only final ID
1099
- remove_prefix_length = input_ids.shape[1] - 1
1100
-
1101
- input_ids = input_ids[:, remove_prefix_length:]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1102
 
1103
  position_ids = kwargs.get("position_ids", None)
1104
  if attention_mask is not None and position_ids is None:
@@ -1112,13 +1315,24 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1112
  if inputs_embeds is not None and past_key_values is None:
1113
  model_inputs = {"inputs_embeds": inputs_embeds}
1114
  else:
1115
- model_inputs = {"input_ids": input_ids}
 
 
 
 
 
 
 
 
 
 
1116
 
1117
  model_inputs.update(
1118
  {
1119
  "position_ids": position_ids,
 
1120
  "past_key_values": past_key_values,
1121
- "use_cache": kwargs.get("use_cache"),
1122
  "attention_mask": attention_mask,
1123
  }
1124
  )
@@ -1133,7 +1347,9 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1133
  )
1134
  return reordered_past
1135
 
1136
- def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = [], meta_instruction=""):
 
 
1137
  if tokenizer.add_bos_token:
1138
  prompt = ""
1139
  else:
@@ -1150,17 +1366,21 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1150
  self,
1151
  tokenizer,
1152
  query: str,
1153
- history: List[Tuple[str, str]] = [],
1154
  streamer: Optional[BaseStreamer] = None,
1155
  max_new_tokens: int = 1024,
1156
  do_sample: bool = True,
1157
  temperature: float = 0.8,
1158
  top_p: float = 0.8,
1159
  meta_instruction: str = "You are an AI assistant whose name is InternLM (书生·浦语).\n"
1160
- "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n"
1161
- "- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文.",
 
 
1162
  **kwargs,
1163
  ):
 
 
1164
  inputs = self.build_inputs(tokenizer, query, history, meta_instruction)
1165
  inputs = {k: v.to(self.device) for k, v in inputs.items() if torch.is_tensor(v)}
1166
  # also add end-of-assistant token in eos token id to avoid unnecessary generation
@@ -1186,13 +1406,15 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1186
  self,
1187
  tokenizer,
1188
  query: str,
1189
- history: List[Tuple[str, str]] = [],
1190
  max_new_tokens: int = 1024,
1191
  do_sample: bool = True,
1192
  temperature: float = 0.8,
1193
  top_p: float = 0.8,
1194
  **kwargs,
1195
  ):
 
 
1196
  """
1197
  Return a generator in format: (response, history)
1198
  Eg.
@@ -1208,6 +1430,10 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1208
  response_queue = queue.Queue(maxsize=20)
1209
 
1210
  class ChatStreamer(BaseStreamer):
 
 
 
 
1211
  def __init__(self, tokenizer) -> None:
1212
  super().__init__()
1213
  self.tokenizer = tokenizer
@@ -1268,13 +1494,13 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1268
  return consumer()
1269
 
1270
 
1271
- # Copied from transformers.model.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
1272
  @add_start_docstrings(
1273
  """
1274
  The InternLM2 Model transformer with a sequence classification head on top (linear layer).
1275
 
1276
- [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification,
1277
- as other causal models (e.g. GPT-2) do.
1278
 
1279
  Since it does classification on the last token, it requires to know the position of the last token. If a
1280
  `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
@@ -1285,6 +1511,8 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1285
  InternLM2_START_DOCSTRING,
1286
  )
1287
  class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
 
 
1288
  def __init__(self, config):
1289
  super().__init__(config)
1290
  self.num_labels = config.num_labels
@@ -1306,7 +1534,7 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1306
  input_ids: torch.LongTensor = None,
1307
  attention_mask: Optional[torch.Tensor] = None,
1308
  position_ids: Optional[torch.LongTensor] = None,
1309
- past_key_values: Optional[List[torch.FloatTensor]] = None,
1310
  inputs_embeds: Optional[torch.FloatTensor] = None,
1311
  labels: Optional[torch.LongTensor] = None,
1312
  use_cache: Optional[bool] = None,
@@ -1347,9 +1575,10 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1347
  sequence_lengths = -1
1348
  else:
1349
  if input_ids is not None:
1350
- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
1351
- logits.device
1352
- )
 
1353
  else:
1354
  sequence_lengths = -1
1355
 
@@ -1361,7 +1590,7 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1361
  if self.config.problem_type is None:
1362
  if self.num_labels == 1:
1363
  self.config.problem_type = "regression"
1364
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1365
  self.config.problem_type = "single_label_classification"
1366
  else:
1367
  self.config.problem_type = "multi_label_classification"
@@ -1389,3 +1618,191 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1389
  hidden_states=transformer_outputs.hidden_states,
1390
  attentions=transformer_outputs.attentions,
1391
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
  # See the License for the specific language governing permissions and
15
  # limitations under the License.
16
+ """PyTorch InternLM2 model."""
17
  import math
18
  import queue
19
  import threading
 
20
  from typing import List, Optional, Tuple, Union
21
 
22
  import torch
 
26
  from torch import nn
27
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
28
  from transformers.activations import ACT2FN
29
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
30
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
31
  from transformers.modeling_outputs import (
32
  BaseModelOutputWithPast,
33
  CausalLMOutputWithPast,
34
+ QuestionAnsweringModelOutput,
35
  SequenceClassifierOutputWithPast,
36
+ TokenClassifierOutput,
37
  )
38
  from transformers.modeling_utils import PreTrainedModel
39
+ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
40
  from transformers.utils import (
41
  add_start_docstrings,
42
  add_start_docstrings_to_model_forward,
43
+ is_flash_attn_greater_or_equal_2_10,
44
  logging,
45
  replace_return_docstrings,
46
  )
47
 
48
  try:
49
  from transformers.generation.streamers import BaseStreamer
50
+ except Exception:
51
  BaseStreamer = None
52
 
53
  from .configuration_internlm2 import InternLM2Config
54
 
55
+
56
+ try:
57
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
58
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
59
+ except:
60
+ pass
61
+
62
+ try:
63
+ support_bf16_triu = torch.__version__ >= "2.1.0"
64
+ except Exception:
65
+ support_bf16_triu = False
66
+
67
  logger = logging.get_logger(__name__)
68
 
69
  _CONFIG_FOR_DOC = "InternLM2Config"
70
 
71
+
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  def _get_unpad_data(attention_mask):
73
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
74
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
75
  max_seqlen_in_batch = seqlens_in_batch.max().item()
76
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) # pylint: disable=E1102
77
  return (
78
  indices,
79
  cu_seqlens,
 
81
  )
82
 
83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  class InternLM2RMSNorm(nn.Module):
85
+ """InternLM2RMSNorm is equivalent to T5LayerNorm."""
86
+
87
  def __init__(self, hidden_size, eps=1e-6):
 
 
 
88
  super().__init__()
89
  self.weight = nn.Parameter(torch.ones(hidden_size))
90
  self.variance_epsilon = eps
 
97
  return self.weight * hidden_states.to(input_dtype)
98
 
99
 
100
+ ALL_LAYERNORM_LAYERS.append(InternLM2RMSNorm)
101
+
102
+
103
  class InternLM2RotaryEmbedding(nn.Module):
104
+ """Rotary Position Embedding for the InternLM2 model. Credits to the Reddit user /u/lucidrains."""
 
105
 
106
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
107
+ super().__init__()
108
+ self.scaling_factor = scaling_factor
109
  self.dim = dim
110
  self.max_position_embeddings = max_position_embeddings
111
  self.base = base
112
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
113
  self.register_buffer("inv_freq", inv_freq, persistent=False)
114
+ # For BC we register cos and sin cached
115
+ self.max_seq_len_cached = max_position_embeddings
116
 
117
+ @torch.no_grad()
118
+ def forward(self, x, position_ids):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
119
  # x: [bs, num_attention_heads, seq_len, head_size]
120
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
121
+ position_ids_expanded = position_ids[:, None, :].float()
122
+ # Force float32 since bfloat16 loses precision on long contexts
123
+ # See https://github.com/huggingface/transformers/pull/29285
124
+ device_type = x.device.type
125
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
126
+ with torch.autocast(device_type=device_type, enabled=False):
127
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
128
+ emb = torch.cat((freqs, freqs), dim=-1)
129
+ cos = emb.cos()
130
+ sin = emb.sin()
131
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
132
 
133
 
 
134
  class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
135
  """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
136
 
137
+ def forward(self, x, position_ids):
138
+ # difference to the original RoPE: a scaling factor is aplied to the position ids
139
+ position_ids = position_ids.float() / self.scaling_factor
140
+ cos, sin = super().forward(x, position_ids)
141
+ return cos, sin
 
 
 
142
 
 
 
 
 
 
143
 
 
 
144
  class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
145
  """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
146
+ Credits to the Reddit users /u/bloc97 and /u/emozilla"""
 
 
 
 
 
 
 
 
147
 
148
+ def forward(self, x, position_ids):
149
+ # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
150
+ seq_len = torch.max(position_ids) + 1
151
  if seq_len > self.max_position_embeddings:
152
  base = self.base * (
153
  (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
154
  ) ** (self.dim / (self.dim - 2))
155
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim))
156
+ self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation
 
 
157
 
158
+ cos, sin = super().forward(x, position_ids)
159
+ return cos, sin
 
 
 
160
 
161
 
 
162
  def rotate_half(x):
163
  """Rotates half the hidden dims of the input."""
164
  x1 = x[..., : x.shape[-1] // 2]
 
166
  return torch.cat((-x2, x1), dim=-1)
167
 
168
 
169
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): # pylint: disable=unused-argument
170
+ """Applies Rotary Position Embedding to the query and key tensors.
171
+
172
+ Args:
173
+ q (`torch.Tensor`): The query tensor.
174
+ k (`torch.Tensor`): The key tensor.
175
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
176
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
177
+ position_ids (`torch.Tensor`, *optional*):
178
+ Deprecated and unused.
179
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
180
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
181
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
182
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
183
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
184
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
185
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
186
+ Returns:
187
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
188
+ """
189
+ cos = cos.unsqueeze(unsqueeze_dim)
190
+ sin = sin.unsqueeze(unsqueeze_dim)
191
  q_embed = (q * cos) + (rotate_half(q) * sin)
192
  k_embed = (k * cos) + (rotate_half(k) * sin)
193
  return q_embed, k_embed
194
 
195
 
196
  class InternLM2MLP(nn.Module):
197
+ """MLP for InternLM2 model."""
198
+
199
  def __init__(self, config):
200
  super().__init__()
201
  self.config = config
 
212
  return down_proj
213
 
214
 
 
215
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
216
  """
217
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
 
224
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
225
 
226
 
 
227
  class InternLM2Attention(nn.Module):
228
  """Multi-headed attention from 'Attention Is All You Need' paper"""
229
 
230
+ def __init__(self, config: InternLM2Config, layer_idx: Optional[int] = None):
231
  super().__init__()
232
  self.config = config
233
+ self.layer_idx = layer_idx
234
+ if layer_idx is None:
235
+ logger.warning_once(
236
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
237
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
238
+ "when creating this class."
239
+ )
240
+
241
  self.hidden_size = config.hidden_size
242
  self.num_heads = config.num_attention_heads
243
  self.head_dim = self.hidden_size // self.num_heads
244
  self.num_key_value_heads = config.num_key_value_heads
245
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
246
  self.max_position_embeddings = config.max_position_embeddings
247
+ self.rope_theta = config.rope_theta
248
  self.is_causal = True
249
 
250
  if (self.head_dim * self.num_heads) != self.hidden_size:
 
258
  (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
259
  bias=config.bias,
260
  )
 
261
  self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
262
+
263
  self._init_rope()
264
 
265
  def _init_rope(self):
 
267
  self.rotary_emb = InternLM2RotaryEmbedding(
268
  self.head_dim,
269
  max_position_embeddings=self.max_position_embeddings,
270
+ base=self.rope_theta,
271
  )
272
  else:
273
  scaling_type = self.config.rope_scaling["type"]
274
  scaling_factor = self.config.rope_scaling["factor"]
275
+ if scaling_type == "linear":
276
+ self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
277
  self.head_dim,
278
  max_position_embeddings=self.max_position_embeddings,
 
279
  scaling_factor=scaling_factor,
280
+ base=self.rope_theta,
281
  )
282
+ elif scaling_type == "dynamic":
283
+ self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
284
  self.head_dim,
285
  max_position_embeddings=self.max_position_embeddings,
 
286
  scaling_factor=scaling_factor,
287
+ base=self.rope_theta,
288
  )
289
  else:
290
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
 
 
 
 
291
 
292
  def forward(
293
  self,
294
  hidden_states: torch.Tensor,
295
  attention_mask: Optional[torch.Tensor] = None,
296
  position_ids: Optional[torch.LongTensor] = None,
297
+ past_key_value: Optional[Cache] = None,
298
  output_attentions: bool = False,
299
+ use_cache: bool = False, # pylint: disable=unused-argument
300
+ cache_position: Optional[torch.LongTensor] = None,
301
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 
 
 
 
 
 
302
  bsz, q_len, _ = hidden_states.size()
303
 
304
+ if self.config.pretraining_tp > 1:
305
+ # split qkv_states by tp size
306
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
307
+ qkv_slices = self.wqkv.weight.split(key_value_slicing, dim=0)
308
+ qkv_states = torch.cat(
309
+ [F.linear(hidden_states, qkv_slice) for qkv_slice in qkv_slices], dim=-1 # pylint: disable=E1102
310
+ )
311
+ else:
312
+ qkv_states = self.wqkv(hidden_states)
313
 
314
  qkv_states = rearrange(
315
  qkv_states,
 
319
  )
320
 
321
  query_states = qkv_states[..., : self.num_key_value_groups, :]
322
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d").transpose(1, 2)
323
+ key_states = qkv_states[..., -2, :].transpose(1, 2)
324
+ value_states = qkv_states[..., -1, :].transpose(1, 2)
325
 
326
+ cos, sin = self.rotary_emb(value_states, position_ids)
 
 
 
 
 
 
 
327
  query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
328
 
329
  if past_key_value is not None:
330
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
331
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
332
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
333
 
334
  key_states = repeat_kv(key_states, self.num_key_value_groups)
335
  value_states = repeat_kv(value_states, self.num_key_value_groups)
336
 
337
  attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
338
 
339
+ if attention_mask is not None: # no matter the length, we just slice it
340
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
341
+ attn_weights = attn_weights + causal_mask
 
 
 
 
 
 
 
 
 
342
 
343
  # upcast attention to fp32
344
  attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
 
351
  )
352
 
353
  attn_output = attn_output.transpose(1, 2).contiguous()
354
+
355
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
356
 
357
+ if self.config.pretraining_tp > 1:
358
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
359
+ o_proj_slices = self.wo.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
360
+ attn_output = sum(
361
+ [
362
+ F.linear(attn_output[i], o_proj_slices[i]) # pylint: disable=E1102
363
+ for i in range(self.config.pretraining_tp)
364
+ ]
365
+ )
366
+ else:
367
+ attn_output = self.wo(attn_output)
368
 
369
  if not output_attentions:
370
  attn_weights = None
 
372
  return attn_output, attn_weights, past_key_value
373
 
374
 
 
375
  class InternLM2FlashAttention2(InternLM2Attention):
376
  """
377
  InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
 
379
  flash attention and deal with padding tokens in case the input contains any of them.
380
  """
381
 
382
+ def __init__(self, *args, **kwargs):
383
+ super().__init__(*args, **kwargs)
384
+
385
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
386
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement,
387
+ # that was made default for flash_attn>=2.1. This attribute is used to handle this difference.
388
+ # Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
389
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1)
390
+ # produces a wrong mask (top-left).
391
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
392
+
393
  def forward(
394
  self,
395
  hidden_states: torch.Tensor,
396
  attention_mask: Optional[torch.LongTensor] = None,
397
  position_ids: Optional[torch.LongTensor] = None,
398
+ past_key_value: Optional[Cache] = None,
399
  output_attentions: bool = False,
400
  use_cache: bool = False,
401
+ cache_position: Optional[torch.LongTensor] = None,
402
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
403
+ if isinstance(past_key_value, StaticCache):
404
+ raise ValueError(
405
+ "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
406
+ "make sure to use `sdpa` in the mean time, and open an issue at "
407
+ "https://github.com/huggingface/transformers"
408
  )
409
 
 
 
 
410
  output_attentions = False
411
 
412
  bsz, q_len, _ = hidden_states.size()
 
429
  key_states = key_states.transpose(1, 2)
430
  value_states = value_states.transpose(1, 2)
431
 
432
+ cos, sin = self.rotary_emb(value_states, position_ids)
433
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
 
 
 
 
434
 
435
  if past_key_value is not None:
436
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
437
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
438
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
439
 
440
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout
441
+ # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
442
+ # to be able to avoid many of these transpose/reshape/view.
443
  query_states = query_states.transpose(1, 2)
444
  key_states = key_states.transpose(1, 2)
445
  value_states = value_states.transpose(1, 2)
446
 
447
+ # dropout_rate = self.attention_dropout if self.training else 0.0
448
+ dropout_rate = 0.0
449
+
450
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
451
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
452
+ # cast them back in the correct dtype just to be sure everything works as expected.
453
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
454
+ # in fp32. (InternLM2RMSNorm handles it correctly)
455
+
456
+ input_dtype = query_states.dtype
457
+ if input_dtype == torch.float32:
458
+ if torch.is_autocast_enabled():
459
+ target_dtype = torch.get_autocast_gpu_dtype()
460
+ # Handle the case where the model is quantized
461
+ elif hasattr(self.config, "_pre_quantization_dtype"):
462
+ target_dtype = self.config._pre_quantization_dtype
463
+ else:
464
+ target_dtype = self.wqkv.weight.dtype
465
+
466
+ logger.warning_once(
467
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
468
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
469
+ f" {target_dtype}."
470
+ )
471
+
472
+ query_states = query_states.to(target_dtype)
473
+ key_states = key_states.to(target_dtype)
474
+ value_states = value_states.to(target_dtype)
475
+
476
  attn_output = self._flash_attention_forward(
477
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
478
  )
479
+
480
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
481
  attn_output = self.wo(attn_output)
482
 
483
  if not output_attentions:
484
  attn_weights = None
485
 
486
+ return attn_output, attn_weights, past_key_value # pylint: disable=E0606
487
 
488
  def _flash_attention_forward(
489
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
 
502
  attention_mask (`torch.Tensor`):
503
  The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
504
  position of padding tokens and 1 for the position of non-padding tokens.
505
+ dropout (`float`):
506
  Attention dropout
507
  softmax_scale (`float`, *optional*):
508
  The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
509
  """
510
+ if not self._flash_attn_uses_top_left_mask:
511
+ causal = self.is_causal
512
+ else:
513
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1.
514
+ # For details, please see the comment in InternLM2FlashAttention2 __init__.
515
+ causal = self.is_causal and query_length != 1
516
+
517
  # Contains at least one padding token in the sequence
 
518
  if attention_mask is not None:
519
  batch_size = query_states.shape[0]
520
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
521
  query_states, key_states, value_states, attention_mask, query_length
522
  )
523
 
524
  cu_seqlens_q, cu_seqlens_k = cu_seq_lens
525
  max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
526
 
527
+ attn_output_unpad = flash_attn_varlen_func( # pylint: disable=E0606
528
  query_states,
529
  key_states,
530
  value_states,
 
537
  causal=causal,
538
  )
539
 
540
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) # pylint: disable=E0606
541
  else:
542
+ attn_output = flash_attn_func( # pylint: disable=E0606
543
  query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
544
  )
545
 
546
  return attn_output
547
 
548
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
549
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
550
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
551
 
552
+ key_layer = index_first_axis( # pylint: disable=E0606
553
  key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
554
  )
555
+ value_layer = index_first_axis( # pylint: disable=E0606
556
  value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
557
  )
 
558
  if query_length == kv_seq_len:
559
+ query_layer = index_first_axis( # pylint: disable=E0606
560
  query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
561
  )
562
  cu_seqlens_q = cu_seqlens_k
 
572
  else:
573
  # The -q_len: slice assumes left padding.
574
  attention_mask = attention_mask[:, -query_length:]
575
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( # pylint: disable=E0606
576
+ query_layer, attention_mask
577
+ )
578
 
579
  return (
580
  query_layer,
581
  key_layer,
582
  value_layer,
583
+ indices_q,
584
  (cu_seqlens_q, cu_seqlens_k),
585
  (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
586
  )
587
 
588
+
589
+ # Copied from transformers.models.llama.modeling_llama.LllamaSdpaAttention with Llama->InternLM2
590
+ class InternLM2SdpaAttention(InternLM2Attention):
591
+ """
592
+ InternLM2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
593
+ `InternLM2Attention` as the weights of the module stays untouched. The only changes are on the forward pass
594
+ to adapt to SDPA API.
595
+ """
596
+
597
+ # Adapted from InternLM2Attention.forward
598
+ def forward(
599
+ self,
600
+ hidden_states: torch.Tensor,
601
+ attention_mask: Optional[torch.Tensor] = None,
602
+ position_ids: Optional[torch.LongTensor] = None,
603
+ past_key_value: Optional[Cache] = None,
604
+ output_attentions: bool = False,
605
+ use_cache: bool = False,
606
+ cache_position: Optional[torch.LongTensor] = None,
607
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
608
+ if output_attentions:
609
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
610
+ # once this is implemented.
611
+ logger.warning_once(
612
+ "InternLM2Model uses InternLM2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
613
+ "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
614
+ "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
615
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
616
+ )
617
+ return super().forward(
618
+ hidden_states=hidden_states,
619
+ attention_mask=attention_mask,
620
+ position_ids=position_ids,
621
+ past_key_value=past_key_value,
622
+ output_attentions=output_attentions,
623
+ use_cache=use_cache,
624
+ cache_position=cache_position,
625
+ )
626
+
627
+ bsz, q_len, _ = hidden_states.size()
628
+
629
+ qkv_states = self.wqkv(hidden_states)
630
+
631
+ qkv_states = rearrange(
632
+ qkv_states,
633
+ "b q (h gs d) -> b q h gs d",
634
+ gs=2 + self.num_key_value_groups,
635
+ d=self.head_dim,
636
+ )
637
+
638
+ query_states = qkv_states[..., : self.num_key_value_groups, :]
639
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
640
+ key_states = qkv_states[..., -2, :]
641
+ value_states = qkv_states[..., -1, :]
642
+
643
+ query_states = query_states.transpose(1, 2)
644
+ key_states = key_states.transpose(1, 2)
645
+ value_states = value_states.transpose(1, 2)
646
+
647
+ cos, sin = self.rotary_emb(value_states, position_ids)
648
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
649
+
650
+ if past_key_value is not None:
651
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
652
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
653
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
654
+
655
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
656
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
657
+
658
+ causal_mask = attention_mask
659
+ if attention_mask is not None:
660
+ causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
661
+
662
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
663
+ # custom attn_mask, Reference: https://github.com/pytorch/pytorch/issues/112577.
664
+ if query_states.device.type == "cuda" and causal_mask is not None:
665
+ query_states = query_states.contiguous()
666
+ key_states = key_states.contiguous()
667
+ value_states = value_states.contiguous()
668
+
669
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of
670
+ # an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph
671
+ # options. An inline conditional prevents dynamic shapes from compiling.
672
+ is_causal = bool(causal_mask is None and q_len > 1)
673
+
674
+ attn_output = torch.nn.functional.scaled_dot_product_attention( # pylint: disable=E1102
675
+ query_states,
676
+ key_states,
677
+ value_states,
678
+ attn_mask=causal_mask,
679
+ dropout_p=0.0,
680
+ is_causal=is_causal,
681
+ )
682
+
683
+ attn_output = attn_output.transpose(1, 2).contiguous()
684
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
685
+
686
+ attn_output = self.wo(attn_output)
687
+
688
+ return attn_output, None, past_key_value
689
+
690
+
691
  INTERNLM2_ATTENTION_CLASSES = {
692
  "eager": InternLM2Attention,
693
  "flash_attention_2": InternLM2FlashAttention2,
694
+ "sdpa": InternLM2SdpaAttention,
695
  }
696
 
697
+
698
+ # Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer with Llama->InternLM2
699
  class InternLM2DecoderLayer(nn.Module):
700
+ """InternLM2 Decoder Layer. This module is a single layer of the InternLM2 model."""
701
+
702
+ def __init__(self, config: InternLM2Config, layer_idx: int):
703
  super().__init__()
704
  self.hidden_size = config.hidden_size
705
+ self.layer_idx = layer_idx
706
 
707
+ self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx)
708
 
709
  self.feed_forward = InternLM2MLP(config)
710
  self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
715
  hidden_states: torch.Tensor,
716
  attention_mask: Optional[torch.Tensor] = None,
717
  position_ids: Optional[torch.LongTensor] = None,
718
+ past_key_value: Optional[Cache] = None,
719
  output_attentions: Optional[bool] = False,
720
  use_cache: Optional[bool] = False,
721
+ cache_position: Optional[torch.LongTensor] = None,
722
  ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
723
  """
724
  Args:
 
734
  (see `past_key_values`).
735
  past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
736
  """
 
 
 
 
 
 
737
  residual = hidden_states
738
 
739
  hidden_states = self.attention_norm(hidden_states)
 
746
  past_key_value=past_key_value,
747
  output_attentions=output_attentions,
748
  use_cache=use_cache,
749
+ cache_position=cache_position,
750
  )
751
  hidden_states = residual + hidden_states
752
 
 
790
  InternLM2_START_DOCSTRING,
791
  )
792
  class InternLM2PreTrainedModel(PreTrainedModel):
793
+ """
794
+ InternLM2 pretraiend model's base class.
795
+ """
796
+
797
  config_class = InternLM2Config
798
  base_model_prefix = "model"
799
  supports_gradient_checkpointing = True
800
  _no_split_modules = ["InternLM2DecoderLayer"]
801
+ _skip_keys_device_placement = ["past_key_values"]
802
+ _supports_flash_attn_2 = True
803
+ _supports_sdpa = True
804
+ _supports_cache_class = True
805
+ _supports_quantized_cache = True
806
+ _supports_static_cache = True
807
 
808
  def _init_weights(self, module):
809
  std = self.config.initializer_range
 
852
  config.n_positions - 1]`.
853
 
854
  [What are position IDs?](../glossary#position-ids)
855
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
856
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
857
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
858
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
859
+
860
+ Two formats are allowed:
861
+ - a [`~cache_utils.Cache`] instance;
862
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
863
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
864
+ cache format.
865
 
866
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
867
+ legacy cache format will be returned.
868
 
869
  If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
870
  have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
 
884
  more detail.
885
  return_dict (`bool`, *optional*):
886
  Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
887
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
888
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
889
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
890
+ the complete sequence length.
891
  """
892
 
893
 
894
+ # Modified from transformers.models.llama.modeling_llama.LlamaModel with Llama->InternLM2
895
  @add_start_docstrings(
896
  "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
897
  InternLM2_START_DOCSTRING,
 
914
 
915
  self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
916
 
917
+ self.layers = nn.ModuleList(
918
+ [InternLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
919
+ )
920
  self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
921
 
922
  self.gradient_checkpointing = False
 
929
  def set_input_embeddings(self, value):
930
  self.tok_embeddings = value
931
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
932
  @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
933
  def forward(
934
  self,
935
  input_ids: torch.LongTensor = None,
936
  attention_mask: Optional[torch.Tensor] = None,
937
  position_ids: Optional[torch.LongTensor] = None,
938
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
939
  inputs_embeds: Optional[torch.FloatTensor] = None,
940
  use_cache: Optional[bool] = None,
941
  output_attentions: Optional[bool] = None,
942
  output_hidden_states: Optional[bool] = None,
943
  return_dict: Optional[bool] = None,
944
+ cache_position: Optional[torch.LongTensor] = None,
945
  ) -> Union[Tuple, BaseModelOutputWithPast]:
946
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
947
  output_hidden_states = (
948
  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
949
  )
950
  use_cache = use_cache if use_cache is not None else self.config.use_cache
 
951
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
952
 
953
+ if (input_ids is None) ^ (inputs_embeds is not None):
954
+ raise ValueError(
955
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
956
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
957
 
958
+ if self.gradient_checkpointing and self.training and use_cache:
959
+ logger.warning_once(
960
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
 
961
  )
962
+ use_cache = False
963
 
964
  if inputs_embeds is None:
965
  inputs_embeds = self.tok_embeddings(input_ids)
966
 
967
+ return_legacy_cache = False
968
+ if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs)
969
+ return_legacy_cache = True
970
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
971
+
972
+ if cache_position is None:
973
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
974
+ cache_position = torch.arange(
975
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
 
976
  )
977
+ if position_ids is None:
978
+ position_ids = cache_position.unsqueeze(0)
979
+
980
+ causal_mask = self._update_causal_mask(
981
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
982
+ )
983
 
984
  # embed positions
985
  hidden_states = inputs_embeds
986
 
 
 
 
 
 
 
 
987
  # decoder layers
988
  all_hidden_states = () if output_hidden_states else None
989
  all_self_attns = () if output_attentions else None
990
+ next_decoder_cache = None
991
 
992
+ for decoder_layer in self.layers:
993
  if output_hidden_states:
994
  all_hidden_states += (hidden_states,)
995
 
 
 
996
  if self.gradient_checkpointing and self.training:
997
+ layer_outputs = self._gradient_checkpointing_func(
998
+ decoder_layer.__call__,
 
 
 
 
 
 
 
 
999
  hidden_states,
1000
+ causal_mask,
1001
  position_ids,
1002
+ past_key_values,
1003
+ output_attentions,
1004
+ use_cache,
1005
+ cache_position,
1006
  )
1007
  else:
1008
  layer_outputs = decoder_layer(
1009
  hidden_states,
1010
+ attention_mask=causal_mask,
1011
  position_ids=position_ids,
1012
+ past_key_value=past_key_values,
1013
  output_attentions=output_attentions,
1014
  use_cache=use_cache,
1015
+ cache_position=cache_position,
1016
  )
1017
 
1018
  hidden_states = layer_outputs[0]
1019
 
1020
  if use_cache:
1021
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1022
 
1023
  if output_attentions:
1024
  all_self_attns += (layer_outputs[1],)
 
1030
  all_hidden_states += (hidden_states,)
1031
 
1032
  next_cache = next_decoder_cache if use_cache else None
1033
+ if return_legacy_cache:
1034
+ next_cache = next_cache.to_legacy_cache()
1035
+
1036
  if not return_dict:
1037
  return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1038
  return BaseModelOutputWithPast(
 
1042
  attentions=all_self_attns,
1043
  )
1044
 
1045
+ def _update_causal_mask(
1046
+ self,
1047
+ attention_mask: torch.Tensor,
1048
+ input_tensor: torch.Tensor,
1049
+ cache_position: torch.Tensor,
1050
+ past_key_values: Cache,
1051
+ output_attentions: bool,
1052
+ ):
1053
+ # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length
1054
+ # even when the static KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at
1055
+ # each decode steps due to the dynamic shapes. (`recording cudagraph tree for symint key 13`, etc.), which is
1056
+ # VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using `fullgraph=True`.
1057
+ # See more context in https://github.com/huggingface/transformers/pull/29114
1058
 
1059
+ if self.config.attn_implementation == "flash_attention_2":
1060
+ if attention_mask is not None and 0.0 in attention_mask:
1061
+ return attention_mask
1062
+ return None
1063
+
1064
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
1065
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
1066
+ # to infer the attention mask.
1067
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1068
+ using_static_cache = isinstance(past_key_values, StaticCache)
1069
+
1070
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
1071
+ if self.config.attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
1072
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
1073
+ attention_mask,
1074
+ inputs_embeds=input_tensor,
1075
+ past_key_values_length=past_seen_tokens,
1076
+ is_training=self.training,
1077
+ ):
1078
+ return None
1079
+
1080
+ dtype, device = input_tensor.dtype, input_tensor.device
1081
+ min_dtype = torch.finfo(dtype).min
1082
+ sequence_length = input_tensor.shape[1]
1083
+ if using_static_cache:
1084
+ target_length = past_key_values.get_max_length()
1085
+ else:
1086
+ target_length = (
1087
+ attention_mask.shape[-1]
1088
+ if isinstance(attention_mask, torch.Tensor)
1089
+ else past_seen_tokens + sequence_length + 1
1090
+ )
1091
+
1092
+ if attention_mask is not None and attention_mask.dim() == 4:
1093
+ # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
1094
+ if attention_mask.max() != 0:
1095
+ raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
1096
+ causal_mask = attention_mask
1097
+ else:
1098
+ causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
1099
+ if sequence_length != 1:
1100
+ if support_bf16_triu or dtype == torch.float32:
1101
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1102
+ else:
1103
+ triu_mask = torch.triu(torch.ones(causal_mask.size(), device=device), diagonal=1).bool()
1104
+ causal_mask.masked_fill_(~triu_mask, 0)
1105
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
1106
+ causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
1107
+ if attention_mask is not None:
1108
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
1109
+ mask_length = attention_mask.shape[-1]
1110
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
1111
+ padding_mask = padding_mask == 0
1112
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
1113
+ padding_mask, min_dtype
1114
+ )
1115
+ if (
1116
+ self.config.attn_implementation == "sdpa"
1117
+ and attention_mask is not None
1118
+ and attention_mask.device.type == "cuda"
1119
+ and not output_attentions
1120
+ ):
1121
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1122
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1123
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1124
+