x54-729 commited on
Commit
a8ad04e
1 Parent(s): c24f301

update for new version

Browse files
Files changed (3) hide show
  1. config.json +2 -1
  2. configuration_internlm2.py +33 -11
  3. modeling_internlm2.py +759 -352
config.json CHANGED
@@ -28,5 +28,6 @@
28
  "torch_dtype": "bfloat16",
29
  "transformers_version": "4.33.0",
30
  "use_cache": true,
31
- "vocab_size": 92544
 
32
  }
 
28
  "torch_dtype": "bfloat16",
29
  "transformers_version": "4.33.0",
30
  "use_cache": true,
31
+ "vocab_size": 92544,
32
+ "pretraining_tp": 1
33
  }
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,
 
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,
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,48 @@ 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 +75,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 +91,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 +160,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 +206,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 +218,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 +252,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 +261,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 +313,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 +345,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 +366,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 +373,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 +423,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 +496,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 +531,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 +566,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 +709,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 +728,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 +740,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 +784,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 +846,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 +878,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 +908,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 +923,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 +1024,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 +1036,91 @@ 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 +1157,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 +1180,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 +1209,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 +1250,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 +1305,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 +1337,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 +1356,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 +1396,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 +1420,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 +1484,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 +1501,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 +1524,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 +1565,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 +1580,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 +1608,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_2_available,
44
+ is_flash_attn_greater_or_equal_2_10,
45
  logging,
46
  replace_return_docstrings,
47
  )
48
 
49
  try:
50
  from transformers.generation.streamers import BaseStreamer
51
+ except Exception:
52
  BaseStreamer = None
53
 
54
  from .configuration_internlm2 import InternLM2Config
55
 
56
+ if is_flash_attn_2_available():
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
+
60
+
61
  logger = logging.get_logger(__name__)
62
 
63
  _CONFIG_FOR_DOC = "InternLM2Config"
64
 
65
+
 
 
 
 
 
 
 
 
 
 
 
 
 
66
  def _get_unpad_data(attention_mask):
67
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
68
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
69
  max_seqlen_in_batch = seqlens_in_batch.max().item()
70
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) # pylint: disable=E1102
71
  return (
72
  indices,
73
  cu_seqlens,
 
75
  )
76
 
77
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
78
  class InternLM2RMSNorm(nn.Module):
79
+ """InternLM2RMSNorm is equivalent to T5LayerNorm."""
80
+
81
  def __init__(self, hidden_size, eps=1e-6):
 
 
 
82
  super().__init__()
83
  self.weight = nn.Parameter(torch.ones(hidden_size))
84
  self.variance_epsilon = eps
 
91
  return self.weight * hidden_states.to(input_dtype)
92
 
93
 
94
+ ALL_LAYERNORM_LAYERS.append(InternLM2RMSNorm)
95
+
96
+
97
  class InternLM2RotaryEmbedding(nn.Module):
98
+ """Rotary Position Embedding for the InternLM2 model. Credits to the Reddit user /u/lucidrains."""
 
99
 
100
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
101
+ super().__init__()
102
+ self.scaling_factor = scaling_factor
103
  self.dim = dim
104
  self.max_position_embeddings = max_position_embeddings
105
  self.base = base
106
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
107
  self.register_buffer("inv_freq", inv_freq, persistent=False)
108
+ # For BC we register cos and sin cached
109
+ self.max_seq_len_cached = max_position_embeddings
110
 
111
+ @torch.no_grad()
112
+ def forward(self, x, position_ids):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
113
  # x: [bs, num_attention_heads, seq_len, head_size]
114
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
115
+ position_ids_expanded = position_ids[:, None, :].float()
116
+ # Force float32 since bfloat16 loses precision on long contexts
117
+ # See https://github.com/huggingface/transformers/pull/29285
118
+ device_type = x.device.type
119
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
120
+ with torch.autocast(device_type=device_type, enabled=False):
121
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
122
+ emb = torch.cat((freqs, freqs), dim=-1)
123
+ cos = emb.cos()
124
+ sin = emb.sin()
125
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
126
 
127
 
 
128
  class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
129
  """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
130
 
131
+ def forward(self, x, position_ids):
132
+ # difference to the original RoPE: a scaling factor is aplied to the position ids
133
+ position_ids = position_ids.float() / self.scaling_factor
134
+ cos, sin = super().forward(x, position_ids)
135
+ return cos, sin
 
 
 
136
 
 
 
 
 
 
137
 
 
 
138
  class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
139
  """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
140
+ Credits to the Reddit users /u/bloc97 and /u/emozilla"""
 
 
 
 
 
 
 
 
141
 
142
+ def forward(self, x, position_ids):
143
+ # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
144
+ seq_len = torch.max(position_ids) + 1
145
  if seq_len > self.max_position_embeddings:
146
  base = self.base * (
147
  (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
148
  ) ** (self.dim / (self.dim - 2))
149
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim))
150
+ self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation
 
 
151
 
152
+ cos, sin = super().forward(x, position_ids)
153
+ return cos, sin
 
 
 
154
 
155
 
 
156
  def rotate_half(x):
157
  """Rotates half the hidden dims of the input."""
158
  x1 = x[..., : x.shape[-1] // 2]
 
160
  return torch.cat((-x2, x1), dim=-1)
161
 
162
 
163
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): # pylint: disable=unused-argument
164
+ """Applies Rotary Position Embedding to the query and key tensors.
165
+
166
+ Args:
167
+ q (`torch.Tensor`): The query tensor.
168
+ k (`torch.Tensor`): The key tensor.
169
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
170
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
171
+ position_ids (`torch.Tensor`, *optional*):
172
+ Deprecated and unused.
173
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
174
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
175
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
176
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
177
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
178
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
179
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
180
+ Returns:
181
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
182
+ """
183
+ cos = cos.unsqueeze(unsqueeze_dim)
184
+ sin = sin.unsqueeze(unsqueeze_dim)
185
  q_embed = (q * cos) + (rotate_half(q) * sin)
186
  k_embed = (k * cos) + (rotate_half(k) * sin)
187
  return q_embed, k_embed
188
 
189
 
190
  class InternLM2MLP(nn.Module):
191
+ """MLP for InternLM2 model."""
192
+
193
  def __init__(self, config):
194
  super().__init__()
195
  self.config = config
 
206
  return down_proj
207
 
208
 
 
209
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
210
  """
211
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
 
218
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
219
 
220
 
 
221
  class InternLM2Attention(nn.Module):
222
  """Multi-headed attention from 'Attention Is All You Need' paper"""
223
 
224
+ def __init__(self, config: InternLM2Config, layer_idx: Optional[int] = None):
225
  super().__init__()
226
  self.config = config
227
+ self.layer_idx = layer_idx
228
+ if layer_idx is None:
229
+ logger.warning_once(
230
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
231
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
232
+ "when creating this class."
233
+ )
234
+
235
  self.hidden_size = config.hidden_size
236
  self.num_heads = config.num_attention_heads
237
  self.head_dim = self.hidden_size // self.num_heads
238
  self.num_key_value_heads = config.num_key_value_heads
239
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
240
  self.max_position_embeddings = config.max_position_embeddings
241
+ self.rope_theta = config.rope_theta
242
  self.is_causal = True
243
 
244
  if (self.head_dim * self.num_heads) != self.hidden_size:
 
252
  (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
253
  bias=config.bias,
254
  )
 
255
  self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
256
+
257
  self._init_rope()
258
 
259
  def _init_rope(self):
 
261
  self.rotary_emb = InternLM2RotaryEmbedding(
262
  self.head_dim,
263
  max_position_embeddings=self.max_position_embeddings,
264
+ base=self.rope_theta,
265
  )
266
  else:
267
  scaling_type = self.config.rope_scaling["type"]
268
  scaling_factor = self.config.rope_scaling["factor"]
269
+ if scaling_type == "linear":
270
+ self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
271
  self.head_dim,
272
  max_position_embeddings=self.max_position_embeddings,
 
273
  scaling_factor=scaling_factor,
274
+ base=self.rope_theta,
275
  )
276
+ elif scaling_type == "dynamic":
277
+ self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
278
  self.head_dim,
279
  max_position_embeddings=self.max_position_embeddings,
 
280
  scaling_factor=scaling_factor,
281
+ base=self.rope_theta,
282
  )
283
  else:
284
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
 
 
 
 
285
 
286
  def forward(
287
  self,
288
  hidden_states: torch.Tensor,
289
  attention_mask: Optional[torch.Tensor] = None,
290
  position_ids: Optional[torch.LongTensor] = None,
291
+ past_key_value: Optional[Cache] = None,
292
  output_attentions: bool = False,
293
+ use_cache: bool = False, # pylint: disable=unused-argument
294
+ cache_position: Optional[torch.LongTensor] = None,
295
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 
 
 
 
 
 
296
  bsz, q_len, _ = hidden_states.size()
297
 
298
+ if self.config.pretraining_tp > 1:
299
+ # split qkv_states by tp size
300
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
301
+ qkv_slices = self.wqkv.weight.split(key_value_slicing, dim=0)
302
+ qkv_states = torch.cat(
303
+ [F.linear(hidden_states, qkv_slice) for qkv_slice in qkv_slices], dim=-1 # pylint: disable=E1102
304
+ )
305
+ else:
306
+ qkv_states = self.wqkv(hidden_states)
307
 
308
  qkv_states = rearrange(
309
  qkv_states,
 
313
  )
314
 
315
  query_states = qkv_states[..., : self.num_key_value_groups, :]
316
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d").transpose(1, 2)
317
+ key_states = qkv_states[..., -2, :].transpose(1, 2)
318
+ value_states = qkv_states[..., -1, :].transpose(1, 2)
319
 
320
+ cos, sin = self.rotary_emb(value_states, position_ids)
 
 
 
 
 
 
 
321
  query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
322
 
323
  if past_key_value is not None:
324
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
325
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
326
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
327
 
328
  key_states = repeat_kv(key_states, self.num_key_value_groups)
329
  value_states = repeat_kv(value_states, self.num_key_value_groups)
330
 
331
  attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
332
 
333
+ if attention_mask is not None: # no matter the length, we just slice it
334
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
335
+ attn_weights = attn_weights + causal_mask
 
 
 
 
 
 
 
 
 
336
 
337
  # upcast attention to fp32
338
  attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
 
345
  )
346
 
347
  attn_output = attn_output.transpose(1, 2).contiguous()
348
+
349
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
350
 
351
+ if self.config.pretraining_tp > 1:
352
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
353
+ o_proj_slices = self.wo.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
354
+ attn_output = sum(
355
+ [
356
+ F.linear(attn_output[i], o_proj_slices[i]) # pylint: disable=E1102
357
+ for i in range(self.config.pretraining_tp)
358
+ ]
359
+ )
360
+ else:
361
+ attn_output = self.wo(attn_output)
362
 
363
  if not output_attentions:
364
  attn_weights = None
 
366
  return attn_output, attn_weights, past_key_value
367
 
368
 
 
369
  class InternLM2FlashAttention2(InternLM2Attention):
370
  """
371
  InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
 
373
  flash attention and deal with padding tokens in case the input contains any of them.
374
  """
375
 
376
+ def __init__(self, *args, **kwargs):
377
+ super().__init__(*args, **kwargs)
378
+
379
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
380
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement,
381
+ # that was made default for flash_attn>=2.1. This attribute is used to handle this difference.
382
+ # Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
383
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1)
384
+ # produces a wrong mask (top-left).
385
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
386
+
387
  def forward(
388
  self,
389
  hidden_states: torch.Tensor,
390
  attention_mask: Optional[torch.LongTensor] = None,
391
  position_ids: Optional[torch.LongTensor] = None,
392
+ past_key_value: Optional[Cache] = None,
393
  output_attentions: bool = False,
394
  use_cache: bool = False,
395
+ cache_position: Optional[torch.LongTensor] = None,
396
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
397
+ if isinstance(past_key_value, StaticCache):
398
+ raise ValueError(
399
+ "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
400
+ "make sure to use `sdpa` in the mean time, and open an issue at "
401
+ "https://github.com/huggingface/transformers"
402
  )
403
 
 
 
 
404
  output_attentions = False
405
 
406
  bsz, q_len, _ = hidden_states.size()
 
423
  key_states = key_states.transpose(1, 2)
424
  value_states = value_states.transpose(1, 2)
425
 
426
+ cos, sin = self.rotary_emb(value_states, position_ids)
427
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
 
 
 
 
428
 
429
  if past_key_value is not None:
430
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
431
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
432
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
433
 
434
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout
435
+ # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
436
+ # to be able to avoid many of these transpose/reshape/view.
437
  query_states = query_states.transpose(1, 2)
438
  key_states = key_states.transpose(1, 2)
439
  value_states = value_states.transpose(1, 2)
440
 
441
+ # dropout_rate = self.attention_dropout if self.training else 0.0
442
+ dropout_rate = 0.0
443
+
444
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
445
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
446
+ # cast them back in the correct dtype just to be sure everything works as expected.
447
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
448
+ # in fp32. (InternLM2RMSNorm handles it correctly)
449
+
450
+ input_dtype = query_states.dtype
451
+ if input_dtype == torch.float32:
452
+ if torch.is_autocast_enabled():
453
+ target_dtype = torch.get_autocast_gpu_dtype()
454
+ # Handle the case where the model is quantized
455
+ elif hasattr(self.config, "_pre_quantization_dtype"):
456
+ target_dtype = self.config._pre_quantization_dtype
457
+ else:
458
+ target_dtype = self.wqkv.weight.dtype
459
+
460
+ logger.warning_once(
461
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
462
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
463
+ f" {target_dtype}."
464
+ )
465
+
466
+ query_states = query_states.to(target_dtype)
467
+ key_states = key_states.to(target_dtype)
468
+ value_states = value_states.to(target_dtype)
469
+
470
  attn_output = self._flash_attention_forward(
471
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
472
  )
473
+
474
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
475
  attn_output = self.wo(attn_output)
476
 
477
  if not output_attentions:
478
  attn_weights = None
479
 
480
+ return attn_output, attn_weights, past_key_value # pylint: disable=E0606
481
 
482
  def _flash_attention_forward(
483
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
 
496
  attention_mask (`torch.Tensor`):
497
  The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
498
  position of padding tokens and 1 for the position of non-padding tokens.
499
+ dropout (`float`):
500
  Attention dropout
501
  softmax_scale (`float`, *optional*):
502
  The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
503
  """
504
+ if not self._flash_attn_uses_top_left_mask:
505
+ causal = self.is_causal
506
+ else:
507
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1.
508
+ # For details, please see the comment in InternLM2FlashAttention2 __init__.
509
+ causal = self.is_causal and query_length != 1
510
+
511
  # Contains at least one padding token in the sequence
 
512
  if attention_mask is not None:
513
  batch_size = query_states.shape[0]
514
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
515
  query_states, key_states, value_states, attention_mask, query_length
516
  )
517
 
518
  cu_seqlens_q, cu_seqlens_k = cu_seq_lens
519
  max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
520
 
521
+ attn_output_unpad = flash_attn_varlen_func( # pylint: disable=E0606
522
  query_states,
523
  key_states,
524
  value_states,
 
531
  causal=causal,
532
  )
533
 
534
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) # pylint: disable=E0606
535
  else:
536
+ attn_output = flash_attn_func( # pylint: disable=E0606
537
  query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
538
  )
539
 
540
  return attn_output
541
 
542
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
543
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
544
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
545
 
546
+ key_layer = index_first_axis( # pylint: disable=E0606
547
  key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
548
  )
549
+ value_layer = index_first_axis( # pylint: disable=E0606
550
  value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
551
  )
 
552
  if query_length == kv_seq_len:
553
+ query_layer = index_first_axis( # pylint: disable=E0606
554
  query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
555
  )
556
  cu_seqlens_q = cu_seqlens_k
 
566
  else:
567
  # The -q_len: slice assumes left padding.
568
  attention_mask = attention_mask[:, -query_length:]
569
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( # pylint: disable=E0606
570
+ query_layer, attention_mask
571
+ )
572
 
573
  return (
574
  query_layer,
575
  key_layer,
576
  value_layer,
577
+ indices_q,
578
  (cu_seqlens_q, cu_seqlens_k),
579
  (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
580
  )
581
 
582
+
583
+ # Copied from transformers.models.llama.modeling_llama.LllamaSdpaAttention with Llama->InternLM2
584
+ class InternLM2SdpaAttention(InternLM2Attention):
585
+ """
586
+ InternLM2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
587
+ `InternLM2Attention` as the weights of the module stays untouched. The only changes are on the forward pass
588
+ to adapt to SDPA API.
589
+ """
590
+
591
+ # Adapted from InternLM2Attention.forward
592
+ def forward(
593
+ self,
594
+ hidden_states: torch.Tensor,
595
+ attention_mask: Optional[torch.Tensor] = None,
596
+ position_ids: Optional[torch.LongTensor] = None,
597
+ past_key_value: Optional[Cache] = None,
598
+ output_attentions: bool = False,
599
+ use_cache: bool = False,
600
+ cache_position: Optional[torch.LongTensor] = None,
601
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
602
+ if output_attentions:
603
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
604
+ # once this is implemented.
605
+ logger.warning_once(
606
+ "InternLM2Model uses InternLM2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
607
+ "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
608
+ "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
609
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
610
+ )
611
+ return super().forward(
612
+ hidden_states=hidden_states,
613
+ attention_mask=attention_mask,
614
+ position_ids=position_ids,
615
+ past_key_value=past_key_value,
616
+ output_attentions=output_attentions,
617
+ use_cache=use_cache,
618
+ cache_position=cache_position,
619
+ )
620
+
621
+ bsz, q_len, _ = hidden_states.size()
622
+
623
+ qkv_states = self.wqkv(hidden_states)
624
+
625
+ qkv_states = rearrange(
626
+ qkv_states,
627
+ "b q (h gs d) -> b q h gs d",
628
+ gs=2 + self.num_key_value_groups,
629
+ d=self.head_dim,
630
+ )
631
+
632
+ query_states = qkv_states[..., : self.num_key_value_groups, :]
633
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
634
+ key_states = qkv_states[..., -2, :]
635
+ value_states = qkv_states[..., -1, :]
636
+
637
+ query_states = query_states.transpose(1, 2)
638
+ key_states = key_states.transpose(1, 2)
639
+ value_states = value_states.transpose(1, 2)
640
+
641
+ cos, sin = self.rotary_emb(value_states, position_ids)
642
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
643
+
644
+ if past_key_value is not None:
645
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
646
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
647
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
648
+
649
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
650
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
651
+
652
+ causal_mask = attention_mask
653
+ if attention_mask is not None:
654
+ causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
655
+
656
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
657
+ # custom attn_mask, Reference: https://github.com/pytorch/pytorch/issues/112577.
658
+ if query_states.device.type == "cuda" and causal_mask is not None:
659
+ query_states = query_states.contiguous()
660
+ key_states = key_states.contiguous()
661
+ value_states = value_states.contiguous()
662
+
663
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of
664
+ # an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph
665
+ # options. An inline conditional prevents dynamic shapes from compiling.
666
+ is_causal = bool(causal_mask is None and q_len > 1)
667
+
668
+ attn_output = torch.nn.functional.scaled_dot_product_attention( # pylint: disable=E1102
669
+ query_states,
670
+ key_states,
671
+ value_states,
672
+ attn_mask=causal_mask,
673
+ dropout_p=0.0,
674
+ is_causal=is_causal,
675
+ )
676
+
677
+ attn_output = attn_output.transpose(1, 2).contiguous()
678
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
679
+
680
+ attn_output = self.wo(attn_output)
681
+
682
+ return attn_output, None, past_key_value
683
+
684
+
685
  INTERNLM2_ATTENTION_CLASSES = {
686
  "eager": InternLM2Attention,
687
  "flash_attention_2": InternLM2FlashAttention2,
688
+ "sdpa": InternLM2SdpaAttention,
689
  }
690
 
691
+
692
+ # Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer with Llama->InternLM2
693
  class InternLM2DecoderLayer(nn.Module):
694
+ """InternLM2 Decoder Layer. This module is a single layer of the InternLM2 model."""
695
+
696
+ def __init__(self, config: InternLM2Config, layer_idx: int):
697
  super().__init__()
698
  self.hidden_size = config.hidden_size
699
+ self.layer_idx = layer_idx
700
 
701
+ self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx)
702
 
703
  self.feed_forward = InternLM2MLP(config)
704
  self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
709
  hidden_states: torch.Tensor,
710
  attention_mask: Optional[torch.Tensor] = None,
711
  position_ids: Optional[torch.LongTensor] = None,
712
+ past_key_value: Optional[Cache] = None,
713
  output_attentions: Optional[bool] = False,
714
  use_cache: Optional[bool] = False,
715
+ cache_position: Optional[torch.LongTensor] = None,
716
  ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
717
  """
718
  Args:
 
728
  (see `past_key_values`).
729
  past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
730
  """
 
 
 
 
 
 
731
  residual = hidden_states
732
 
733
  hidden_states = self.attention_norm(hidden_states)
 
740
  past_key_value=past_key_value,
741
  output_attentions=output_attentions,
742
  use_cache=use_cache,
743
+ cache_position=cache_position,
744
  )
745
  hidden_states = residual + hidden_states
746
 
 
784
  InternLM2_START_DOCSTRING,
785
  )
786
  class InternLM2PreTrainedModel(PreTrainedModel):
787
+ """
788
+ InternLM2 pretraiend model's base class.
789
+ """
790
+
791
  config_class = InternLM2Config
792
  base_model_prefix = "model"
793
  supports_gradient_checkpointing = True
794
  _no_split_modules = ["InternLM2DecoderLayer"]
795
+ _skip_keys_device_placement = ["past_key_values"]
796
+ _supports_flash_attn_2 = True
797
+ _supports_sdpa = True
798
+ _supports_cache_class = True
799
+ _supports_quantized_cache = True
800
+ _supports_static_cache = True
801
 
802
  def _init_weights(self, module):
803
  std = self.config.initializer_range
 
846
  config.n_positions - 1]`.
847
 
848
  [What are position IDs?](../glossary#position-ids)
849
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
850
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
851
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
852
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
853
+
854
+ Two formats are allowed:
855
+ - a [`~cache_utils.Cache`] instance;
856
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
857
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
858
+ cache format.
859
 
860
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
861
+ legacy cache format will be returned.
862
 
863
  If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
864
  have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
 
878
  more detail.
879
  return_dict (`bool`, *optional*):
880
  Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
881
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
882
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
883
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
884
+ the complete sequence length.
885
  """
886
 
887
 
888
+ # Modified from transformers.models.llama.modeling_llama.LlamaModel with Llama->InternLM2
889
  @add_start_docstrings(
890
  "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
891
  InternLM2_START_DOCSTRING,
 
908
 
909
  self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
910
 
911
+ self.layers = nn.ModuleList(
912
+ [InternLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
913
+ )
914
  self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
915
 
916
  self.gradient_checkpointing = False
 
923
  def set_input_embeddings(self, value):
924
  self.tok_embeddings = value
925
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
926
  @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
927
  def forward(
928
  self,
929
  input_ids: torch.LongTensor = None,
930
  attention_mask: Optional[torch.Tensor] = None,
931
  position_ids: Optional[torch.LongTensor] = None,
932
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
933
  inputs_embeds: Optional[torch.FloatTensor] = None,
934
  use_cache: Optional[bool] = None,
935
  output_attentions: Optional[bool] = None,
936
  output_hidden_states: Optional[bool] = None,
937
  return_dict: Optional[bool] = None,
938
+ cache_position: Optional[torch.LongTensor] = None,
939
  ) -> Union[Tuple, BaseModelOutputWithPast]:
940
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
941
  output_hidden_states = (
942
  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
943
  )
944
  use_cache = use_cache if use_cache is not None else self.config.use_cache
 
945
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
946
 
947
+ if (input_ids is None) ^ (inputs_embeds is not None):
948
+ raise ValueError(
949
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
950
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
951
 
952
+ if self.gradient_checkpointing and self.training and use_cache:
953
+ logger.warning_once(
954
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
 
955
  )
956
+ use_cache = False
957
 
958
  if inputs_embeds is None:
959
  inputs_embeds = self.tok_embeddings(input_ids)
960
 
961
+ return_legacy_cache = False
962
+ if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs)
963
+ return_legacy_cache = True
964
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
965
+
966
+ if cache_position is None:
967
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
968
+ cache_position = torch.arange(
969
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
 
970
  )
971
+ if position_ids is None:
972
+ position_ids = cache_position.unsqueeze(0)
973
+
974
+ causal_mask = self._update_causal_mask(
975
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
976
+ )
977
 
978
  # embed positions
979
  hidden_states = inputs_embeds
980
 
 
 
 
 
 
 
 
981
  # decoder layers
982
  all_hidden_states = () if output_hidden_states else None
983
  all_self_attns = () if output_attentions else None
984
+ next_decoder_cache = None
985
 
986
+ for decoder_layer in self.layers:
987
  if output_hidden_states:
988
  all_hidden_states += (hidden_states,)
989
 
 
 
990
  if self.gradient_checkpointing and self.training:
991
+ layer_outputs = self._gradient_checkpointing_func(
992
+ decoder_layer.__call__,
 
 
 
 
 
 
 
 
993
  hidden_states,
994
+ causal_mask,
995
  position_ids,
996
+ past_key_values,
997
+ output_attentions,
998
+ use_cache,
999
+ cache_position,
1000
  )
1001
  else:
1002
  layer_outputs = decoder_layer(
1003
  hidden_states,
1004
+ attention_mask=causal_mask,
1005
  position_ids=position_ids,
1006
+ past_key_value=past_key_values,
1007
  output_attentions=output_attentions,
1008
  use_cache=use_cache,
1009
+ cache_position=cache_position,
1010
  )
1011
 
1012
  hidden_states = layer_outputs[0]
1013
 
1014
  if use_cache:
1015
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1016
 
1017
  if output_attentions:
1018
  all_self_attns += (layer_outputs[1],)
 
1024
  all_hidden_states += (hidden_states,)
1025
 
1026
  next_cache = next_decoder_cache if use_cache else None
1027
+ if return_legacy_cache:
1028
+ next_cache = next_cache.to_legacy_cache()
1029
+
1030
  if not return_dict:
1031
  return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1032
  return BaseModelOutputWithPast(
 
1036
  attentions=all_self_attns,
1037
  )
1038
 
1039
+ def _update_causal_mask(
1040
+ self,
1041
+ attention_mask: torch.Tensor,
1042
+ input_tensor: torch.Tensor,
1043
+ cache_position: torch.Tensor,
1044
+ past_key_values: Cache,
1045
+ output_attentions: bool,
1046
+ ):
1047
+ # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length
1048
+ # even when the static KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at
1049
+ # each decode steps due to the dynamic shapes. (`recording cudagraph tree for symint key 13`, etc.), which is
1050
+ # VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using `fullgraph=True`.
1051
+ # See more context in https://github.com/huggingface/transformers/pull/29114
1052
 
1053
+ if self.config.attn_implementation == "flash_attention_2":
1054
+ if attention_mask is not None and 0.0 in attention_mask:
1055
+ return attention_mask
1056
+ return None
1057
+
1058
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
1059
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
1060
+ # to infer the attention mask.
1061
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1062
+ using_static_cache = isinstance(past_key_values, StaticCache)
1063
+
1064
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
1065
+ if self.config.attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
1066
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
1067
+ attention_mask,
1068
+ inputs_embeds=input_tensor,
1069
+ past_key_values_length=past_seen_tokens,
1070
+ is_training=self.training,
1071
+ ):
1072
+ return None
1073
+
1074
+ dtype, device = input_tensor.dtype, input_tensor.device
1075
+ min_dtype = torch.finfo(dtype).min
1076
+ sequence_length = input_tensor.shape[1]
1077
+ if using_static_cache:
1078
+ target_length = past_key_values.get_max_length()
1079
+ else:
1080
+ target_length = (
1081
+ attention_mask.shape[-1]
1082
+ if isinstance(attention_mask, torch.Tensor)
1083
+ else past_seen_tokens + sequence_length + 1
1084
+ )
1085
+
1086
+ if attention_mask is not None and attention_mask.dim() == 4:
1087
+ # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
1088
+ if attention_mask.max() != 0:
1089
+ raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
1090
+ causal_mask = attention_mask
1091
+ else:
1092
+ causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
1093
+ if sequence_length != 1:
1094
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1095
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
1096
+ causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
1097
+ if attention_mask is not None:
1098
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
1099
+ mask_length = attention_mask.shape[-1]
1100
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
1101
+ padding_mask = padding_mask == 0
1102
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
1103
+ padding_mask, min_dtype
1104
+ )
1105
+ if (
1106
+ self.config.attn_implementation == "sdpa"
1107
+ and attention_mask is not None
1108
+ and attention_mask.device.type == "cuda"
1109
+ and not output_attentions
1110
+ ):
1111
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1112
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1113
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1114
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) # pylint: disable=E1120
1115
+
1116
+ return causal_mask
1117
+
1118
+
1119
+ # Modified from transformers.models.llama.modeling_llama.LlamaForCausalLM
1120
  class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1121
+ """Causal language model (CLM) for InternLM2."""
1122
 
1123
+ _auto_class = "AutoModelForCausalLM"
1124
  _tied_weights_keys = ["output.weight"]
1125
 
1126
  def __init__(self, config):
 
1157
  input_ids: torch.LongTensor = None,
1158
  attention_mask: Optional[torch.Tensor] = None,
1159
  position_ids: Optional[torch.LongTensor] = None,
1160
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1161
  inputs_embeds: Optional[torch.FloatTensor] = None,
1162
  labels: Optional[torch.LongTensor] = None,
1163
  use_cache: Optional[bool] = None,
1164
  output_attentions: Optional[bool] = None,
1165
  output_hidden_states: Optional[bool] = None,
1166
  return_dict: Optional[bool] = None,
1167
+ cache_position: Optional[torch.LongTensor] = None,
1168
  ) -> Union[Tuple, CausalLMOutputWithPast]:
1169
  r"""
1170
  Args:
 
1180
  ```python
1181
  >>> from transformers import AutoTokenizer, InternLM2ForCausalLM
1182
 
1183
+ >>> model = InternLM2ForCausalLM.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
1184
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
1185
 
1186
  >>> prompt = "Hey, are you conscious? Can you talk to me?"
1187
  >>> inputs = tokenizer(prompt, return_tensors="pt")
 
1209
  output_attentions=output_attentions,
1210
  output_hidden_states=output_hidden_states,
1211
  return_dict=return_dict,
1212
+ cache_position=cache_position,
1213
  )
1214
 
1215
  hidden_states = outputs[0]
1216
+ if self.config.pretraining_tp > 1:
1217
+ output_slices = self.output.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
1218
+ logits = [
1219
+ F.linear(hidden_states, output_slices[i]) # pylint: disable=not-callable
1220
+ for i in range(self.config.pretraining_tp)
1221
+ ]
1222
+ logits = torch.cat(logits, dim=-1)
1223
+ else:
1224
+ logits = self.output(hidden_states)
1225
  logits = logits.float()
1226
 
1227
  loss = None
 
1250
  )
1251
 
1252
  def prepare_inputs_for_generation(
1253
+ self,
1254
+ input_ids,
1255
+ past_key_values=None,
1256
+ attention_mask=None,
1257
+ inputs_embeds=None,
1258
+ cache_position=None,
1259
+ use_cache=True,
1260
+ **kwargs,
1261
  ):
1262
+ past_length = 0
1263
  if past_key_values is not None:
1264
+ if isinstance(past_key_values, Cache):
1265
+ past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
1266
+ max_cache_length = (
1267
+ torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
1268
+ if past_key_values.get_max_length() is not None
1269
+ else None
1270
+ )
1271
+ cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
1272
+ # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
1273
  else:
1274
+ cache_length = past_length = past_key_values[0][0].shape[2]
1275
+ max_cache_length = None
1276
+
1277
+ # Keep only the unprocessed tokens:
1278
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1279
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input)
1280
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1281
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1282
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1283
+ # input_ids based on the past_length.
1284
+ elif past_length < input_ids.shape[1]:
1285
+ input_ids = input_ids[:, past_length:]
1286
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1287
+
1288
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1289
+ if (
1290
+ max_cache_length is not None
1291
+ and attention_mask is not None
1292
+ and cache_length + input_ids.shape[1] > max_cache_length
1293
+ ):
1294
+ attention_mask = attention_mask[:, -max_cache_length:] # pylint: disable=E1130
1295
 
1296
  position_ids = kwargs.get("position_ids", None)
1297
  if attention_mask is not None and position_ids is None:
 
1305
  if inputs_embeds is not None and past_key_values is None:
1306
  model_inputs = {"inputs_embeds": inputs_embeds}
1307
  else:
1308
+ # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
1309
+ # recompiles graphs as the stride of the inputs is a guard.
1310
+ # Ref: https://github.com/huggingface/transformers/pull/29114
1311
+ # TODO: use `next_tokens` directly instead.
1312
+ model_inputs = {"input_ids": input_ids.contiguous()}
1313
+
1314
+ input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
1315
+ if cache_position is None:
1316
+ cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
1317
+ elif use_cache:
1318
+ cache_position = cache_position[-input_length:]
1319
 
1320
  model_inputs.update(
1321
  {
1322
  "position_ids": position_ids,
1323
+ "cache_position": cache_position,
1324
  "past_key_values": past_key_values,
1325
+ "use_cache": use_cache,
1326
  "attention_mask": attention_mask,
1327
  }
1328
  )
 
1337
  )
1338
  return reordered_past
1339
 
1340
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, meta_instruction=""):
1341
+ if history is None:
1342
+ history = []
1343
  if tokenizer.add_bos_token:
1344
  prompt = ""
1345
  else:
 
1356
  self,
1357
  tokenizer,
1358
  query: str,
1359
+ history: Optional[List[Tuple[str, str]]] = None,
1360
  streamer: Optional[BaseStreamer] = None,
1361
  max_new_tokens: int = 1024,
1362
  do_sample: bool = True,
1363
  temperature: float = 0.8,
1364
  top_p: float = 0.8,
1365
  meta_instruction: str = "You are an AI assistant whose name is InternLM (书生·浦语).\n"
1366
+ "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory "
1367
+ "(上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n"
1368
+ "- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such "
1369
+ "as English and 中文.",
1370
  **kwargs,
1371
  ):
1372
+ if history is None:
1373
+ history = []
1374
  inputs = self.build_inputs(tokenizer, query, history, meta_instruction)
1375
  inputs = {k: v.to(self.device) for k, v in inputs.items() if torch.is_tensor(v)}
1376
  # also add end-of-assistant token in eos token id to avoid unnecessary generation
 
1396
  self,
1397
  tokenizer,
1398
  query: str,
1399
+ history: List[Tuple[str, str]] = None,
1400
  max_new_tokens: int = 1024,
1401
  do_sample: bool = True,
1402
  temperature: float = 0.8,
1403
  top_p: float = 0.8,
1404
  **kwargs,
1405
  ):
1406
+ if history is None:
1407
+ history = []
1408
  """
1409
  Return a generator in format: (response, history)
1410
  Eg.
 
1420
  response_queue = queue.Queue(maxsize=20)
1421
 
1422
  class ChatStreamer(BaseStreamer):
1423
+ """
1424
+ Streamer used in generate to print words one by one.
1425
+ """
1426
+
1427
  def __init__(self, tokenizer) -> None:
1428
  super().__init__()
1429
  self.tokenizer = tokenizer
 
1484
  return consumer()
1485
 
1486
 
1487
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
1488
  @add_start_docstrings(
1489
  """
1490
  The InternLM2 Model transformer with a sequence classification head on top (linear layer).
1491
 
1492
+ [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1493
+ (e.g. GPT-2) do.
1494
 
1495
  Since it does classification on the last token, it requires to know the position of the last token. If a
1496
  `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
 
1501
  InternLM2_START_DOCSTRING,
1502
  )
1503
  class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1504
+ """Sequence Classification Head for InternLM2 Model."""
1505
+
1506
  def __init__(self, config):
1507
  super().__init__(config)
1508
  self.num_labels = config.num_labels
 
1524
  input_ids: torch.LongTensor = None,
1525
  attention_mask: Optional[torch.Tensor] = None,
1526
  position_ids: Optional[torch.LongTensor] = None,
1527
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1528
  inputs_embeds: Optional[torch.FloatTensor] = None,
1529
  labels: Optional[torch.LongTensor] = None,
1530
  use_cache: Optional[bool] = None,
 
1565
  sequence_lengths = -1
1566
  else:
1567
  if input_ids is not None:
1568
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1569
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1570
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1571
+ sequence_lengths = sequence_lengths.to(logits.device)
1572
  else:
1573
  sequence_lengths = -1
1574
 
 
1580
  if self.config.problem_type is None:
1581
  if self.num_labels == 1:
1582
  self.config.problem_type = "regression"
1583
+ elif self.num_labels > 1 and (labels.dtype in (torch.long, torch.int)):
1584
  self.config.problem_type = "single_label_classification"
1585
  else:
1586
  self.config.problem_type = "multi_label_classification"
 
1608
  hidden_states=transformer_outputs.hidden_states,
1609
  attentions=transformer_outputs.attentions,
1610
  )
1611
+
1612
+
1613
+ # Copied from transformers.models.llama.modeling_llama.LlamaForQuestionAnswering with Llama->InternLM2
1614
+ @add_start_docstrings(
1615
+ """
1616
+ The InternLM2 Model transformer with a span classification head on top for extractive question-answering tasks like
1617
+ SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
1618
+ """,
1619
+ InternLM2_START_DOCSTRING,
1620
+ )
1621
+ class InternLM2ForQuestionAnswering(InternLM2PreTrainedModel):
1622
+ """Question Answering model for InternLM2."""
1623
+
1624
+ base_model_prefix = "transformer"
1625
+
1626
+ def __init__(self, config):
1627
+ super().__init__(config)
1628
+ self.transformer = InternLM2Model(config)
1629
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
1630
+
1631
+ # Initialize weights and apply final processing
1632
+ self.post_init()
1633
+
1634
+ def get_input_embeddings(self):
1635
+ return self.transformer.embed_tokens
1636
+
1637
+ def set_input_embeddings(self, value):
1638
+ self.transformer.embed_tokens = value
1639
+
1640
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1641
+ def forward(
1642
+ self,
1643
+ input_ids: Optional[torch.LongTensor] = None,
1644
+ attention_mask: Optional[torch.FloatTensor] = None,
1645
+ position_ids: Optional[torch.LongTensor] = None,
1646
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1647
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1648
+ start_positions: Optional[torch.LongTensor] = None,
1649
+ end_positions: Optional[torch.LongTensor] = None,
1650
+ output_attentions: Optional[bool] = None,
1651
+ output_hidden_states: Optional[bool] = None,
1652
+ return_dict: Optional[bool] = None,
1653
+ ) -> Union[Tuple, QuestionAnsweringModelOutput]:
1654
+ r"""
1655
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1656
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1657
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1658
+ are not taken into account for computing the loss.
1659
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1660
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1661
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1662
+ are not taken into account for computing the loss.
1663
+ """
1664
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1665
+
1666
+ outputs = self.transformer(
1667
+ input_ids,
1668
+ attention_mask=attention_mask,
1669
+ position_ids=position_ids,
1670
+ past_key_values=past_key_values,
1671
+ inputs_embeds=inputs_embeds,
1672
+ output_attentions=output_attentions,
1673
+ output_hidden_states=output_hidden_states,
1674
+ return_dict=return_dict,
1675
+ )
1676
+
1677
+ sequence_output = outputs[0]
1678
+
1679
+ logits = self.qa_outputs(sequence_output)
1680
+ start_logits, end_logits = logits.split(1, dim=-1)
1681
+ start_logits = start_logits.squeeze(-1).contiguous()
1682
+ end_logits = end_logits.squeeze(-1).contiguous()
1683
+
1684
+ total_loss = None
1685
+ if start_positions is not None and end_positions is not None:
1686
+ # If we are on multi-GPU, split add a dimension
1687
+ if len(start_positions.size()) > 1:
1688
+ start_positions = start_positions.squeeze(-1).to(start_logits.device)
1689
+ if len(end_positions.size()) > 1:
1690
+ end_positions = end_positions.squeeze(-1).to(end_logits.device)
1691
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1692
+ ignored_index = start_logits.size(1)
1693
+ start_positions = start_positions.clamp(0, ignored_index)
1694
+ end_positions = end_positions.clamp(0, ignored_index)
1695
+
1696
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1697
+ start_loss = loss_fct(start_logits, start_positions)
1698
+ end_loss = loss_fct(end_logits, end_positions)
1699
+ total_loss = (start_loss + end_loss) / 2
1700
+
1701
+ if not return_dict:
1702
+ output = (start_logits, end_logits) + outputs[2:]
1703
+ return ((total_loss,) + output) if total_loss is not None else output
1704
+
1705
+ return QuestionAnsweringModelOutput(
1706
+ loss=total_loss,
1707
+ start_logits=start_logits,
1708
+ end_logits=end_logits,
1709
+ hidden_states=outputs.hidden_states,
1710
+ attentions=outputs.attentions,
1711
+ )
1712
+
1713
+
1714
+ # Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->InternLM2
1715
+ @add_start_docstrings(
1716
+ """
1717
+ The InternLM2 Model transformer with a token classification head on top (a linear layer on top of the hidden-states
1718
+ output) e.g. for Named-Entity-Recognition (NER) tasks.
1719
+ """,
1720
+ InternLM2_START_DOCSTRING,
1721
+ )
1722
+ class InternLM2ForTokenClassification(InternLM2PreTrainedModel):
1723
+ """Token classification model for InternLM2."""
1724
+
1725
+ def __init__(self, config):
1726
+ super().__init__(config)
1727
+ self.num_labels = config.num_labels
1728
+ self.model = InternLM2Model(config)
1729
+ if getattr(config, "classifier_dropout", None) is not None:
1730
+ classifier_dropout = config.classifier_dropout
1731
+ elif getattr(config, "hidden_dropout", None) is not None:
1732
+ classifier_dropout = config.hidden_dropout
1733
+ else:
1734
+ classifier_dropout = 0.1
1735
+ self.dropout = nn.Dropout(classifier_dropout)
1736
+ self.score = nn.Linear(config.hidden_size, config.num_labels)
1737
+
1738
+ # Initialize weights and apply final processing
1739
+ self.post_init()
1740
+
1741
+ def get_input_embeddings(self):
1742
+ return self.model.embed_tokens
1743
+
1744
+ def set_input_embeddings(self, value):
1745
+ self.model.embed_tokens = value
1746
+
1747
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1748
+ def forward(
1749
+ self,
1750
+ input_ids: torch.LongTensor = None,
1751
+ attention_mask: Optional[torch.Tensor] = None,
1752
+ position_ids: Optional[torch.LongTensor] = None,
1753
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1754
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1755
+ labels: Optional[torch.LongTensor] = None,
1756
+ use_cache: Optional[bool] = None,
1757
+ output_attentions: Optional[bool] = None,
1758
+ output_hidden_states: Optional[bool] = None,
1759
+ return_dict: Optional[bool] = None,
1760
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1761
+ r"""
1762
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1763
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1764
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1765
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1766
+ """
1767
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1768
+
1769
+ outputs = self.model(
1770
+ input_ids,
1771
+ attention_mask=attention_mask,
1772
+ position_ids=position_ids,
1773
+ past_key_values=past_key_values,
1774
+ inputs_embeds=inputs_embeds,
1775
+ use_cache=use_cache,
1776
+ output_attentions=output_attentions,
1777
+ output_hidden_states=output_hidden_states,
1778
+ return_dict=return_dict,
1779
+ )
1780
+ sequence_output = outputs[0]
1781
+ sequence_output = self.dropout(sequence_output)
1782
+ logits = self.score(sequence_output)
1783
+
1784
+ loss = None
1785
+ if labels is not None:
1786
+ loss_fct = CrossEntropyLoss()
1787
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1788
+
1789
+ if not return_dict:
1790
+ output = (logits,) + outputs[2:]
1791
+ return ((loss,) + output) if loss is not None else output
1792
+
1793
+ return TokenClassifierOutput(
1794
+ loss=loss,
1795
+ logits=logits,
1796
+ hidden_states=outputs.hidden_states,
1797
+ attentions=outputs.attentions,
1798
+ )