jon-tow commited on
Commit
efd6eab
1 Parent(s): 64c3d6a

merge: upload `transformers` implementation

Browse files
config.json CHANGED
@@ -1,11 +1,7 @@
1
  {
2
  "architectures": [
3
- "StableLMEpochForCausalLM"
4
  ],
5
- "auto_map": {
6
- "AutoConfig": "configuration_stablelm_epoch.StableLMEpochConfig",
7
- "AutoModelForCausalLM": "modeling_stablelm_epoch.StableLMEpochForCausalLM"
8
- },
9
  "bos_token_id": 100257,
10
  "eos_token_id": 100257,
11
  "hidden_act": "silu",
@@ -13,18 +9,16 @@
13
  "initializer_range": 0.02,
14
  "intermediate_size": 5632,
15
  "max_position_embeddings": 4096,
16
- "model_type": "stablelm_epoch",
17
- "norm_eps": 1e-05,
18
  "num_attention_heads": 32,
19
- "num_heads": 32,
20
  "num_hidden_layers": 24,
21
  "num_key_value_heads": 32,
22
- "rope_pct": 0.25,
23
  "rope_theta": 10000,
24
- "rotary_scaling_factor": 1.0,
25
  "tie_word_embeddings": false,
26
  "torch_dtype": "bfloat16",
27
- "transformers_version": "4.36.2",
28
  "use_cache": true,
29
  "use_qkv_bias": true,
30
  "vocab_size": 100352
 
1
  {
2
  "architectures": [
3
+ "StableLmForCausalLM"
4
  ],
 
 
 
 
5
  "bos_token_id": 100257,
6
  "eos_token_id": 100257,
7
  "hidden_act": "silu",
 
9
  "initializer_range": 0.02,
10
  "intermediate_size": 5632,
11
  "max_position_embeddings": 4096,
12
+ "model_type": "stablelm",
13
+ "layer_norm_eps": 1e-05,
14
  "num_attention_heads": 32,
 
15
  "num_hidden_layers": 24,
16
  "num_key_value_heads": 32,
17
+ "partial_rotary_factor": 0.25,
18
  "rope_theta": 10000,
 
19
  "tie_word_embeddings": false,
20
  "torch_dtype": "bfloat16",
21
+ "transformers_version": "4.38.0",
22
  "use_cache": true,
23
  "use_qkv_bias": true,
24
  "vocab_size": 100352
configuration_stablelm_epoch.py → configuration_stablelm.py RENAMED
@@ -1,4 +1,5 @@
1
- # Copyright 2023 Stability and The HuggingFace Inc. team. All rights reserved.
 
2
  #
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
  # you may not use this file except in compliance with the License.
@@ -11,23 +12,36 @@
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
  # See the License for the specific language governing permissions and
13
  # limitations under the License.
14
- """ StableLM Epoch model configuration"""
15
- from transformers import PretrainedConfig
 
16
  from transformers.utils import logging
17
 
18
 
19
  logger = logging.get_logger(__name__)
20
 
 
 
 
 
 
21
 
22
- class StableLMEpochConfig(PretrainedConfig):
23
  r"""
24
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
25
- documentation from [`PretrainedConfig`] for more information.
 
 
 
 
 
 
 
26
 
27
  Args:
28
- vocab_size (`int`, *optional*, defaults to 50_304):
29
  Vocabulary size of the StableLM model. Defines the number of different tokens that
30
- can be represented by the `inputs_ids` passed when calling [`StableLMEpochModel`].
31
  intermediate_size (`int`, *optional*, defaults to 6912):
32
  Dimension of the MLP representations.
33
  hidden_size (`int`, *optional*, defaults to 2560):
@@ -36,7 +50,7 @@ class StableLMEpochConfig(PretrainedConfig):
36
  Number of hidden layers in the Transformer decoder.
37
  num_attention_heads (`int`, *optional*, defaults to 32):
38
  Number of attention heads for each attention layer in the Transformer encoder.
39
- num_key_value_heads (`int`, *optional*):
40
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
41
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
42
  `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
@@ -46,72 +60,122 @@ class StableLMEpochConfig(PretrainedConfig):
46
  `num_attention_heads`.
47
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
48
  The non-linear activation function (function or string).
49
- rope_pct (`float`, *optional*, defaults to 1.0):
50
- Percentage of hidden dimensions to allocate to rotary embeddings.
51
- rope_theta (`float`, *optional*, defaults to 10000.0):
52
- The base period of the RoPE embeddings.
53
- max_position_embeddings (`int`, *optional*, defaults to 2048):
54
  The maximum sequence length that this model might ever be used with.
55
  Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
56
- initializer_range (`float`, *optional*, defaults to 1e-5):
57
  The standard deviation of the truncated_normal_initializer for initializing
58
  all weight matrices.
59
- norm_eps (`float`, *optional*, defaults to 1e-8):
60
  The epsilon used by the normalization layers.
61
  use_cache (`bool`, *optional*, defaults to `True`):
62
  Whether or not the model should return the last key/values attentions
63
  (not used by all models). Only relevant if `config.is_decoder=True`.
64
- use_qkv_bias (`bool`, *optional*, defaults to `True`):
 
 
 
 
 
 
 
 
 
 
 
 
65
  Whether or not the model should use bias for qkv layers.
66
- tie_word_embeddings(`bool`, *optional*, defaults to `False`):
67
- Whether to tie weight embeddings
68
  attention_dropout (`float`, *optional*, defaults to 0.0):
69
  The dropout ratio for the attention probabilities.
70
- """
71
- model_type = "stablelm_epoch"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  keys_to_ignore_at_inference = ["past_key_values"]
73
 
74
  def __init__(
75
  self,
76
- vocab_size=50_304,
77
  intermediate_size=6912,
78
  hidden_size=2560,
79
  num_hidden_layers=32,
80
  num_attention_heads=32,
81
  num_key_value_heads=32,
82
  hidden_act="silu",
83
- rope_pct=0.25,
84
- rope_theta=10_000,
85
  max_position_embeddings=4096,
86
  initializer_range=0.02,
87
- norm_eps=1.0e-5,
88
  use_cache=True,
89
- use_qkv_bias=True,
90
- bos_token_id=0,
91
- eos_token_id=2,
92
  tie_word_embeddings=False,
93
- attention_dropout: float = 0.0,
 
 
 
 
 
 
 
94
  **kwargs,
95
  ):
96
  self.vocab_size = vocab_size
97
  self.max_position_embeddings = max_position_embeddings
98
- self.intermediate_size = intermediate_size
99
  self.hidden_size = hidden_size
 
100
  self.num_hidden_layers = num_hidden_layers
101
  self.num_attention_heads = num_attention_heads
102
  self.num_key_value_heads = num_key_value_heads
103
  self.hidden_act = hidden_act
104
- self.rope_pct = rope_pct
105
- self.rope_theta = rope_theta
106
  self.initializer_range = initializer_range
107
- self.norm_eps = norm_eps
108
  self.use_cache = use_cache
 
 
109
  self.use_qkv_bias = use_qkv_bias
110
- self.tie_word_embeddings = tie_word_embeddings
111
  self.attention_dropout = attention_dropout
 
 
 
112
  super().__init__(
113
  bos_token_id=bos_token_id,
114
  eos_token_id=eos_token_id,
115
  tie_word_embeddings=tie_word_embeddings,
116
  **kwargs,
117
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 Stability AI and The HuggingFace Inc. team. All rights reserved.
3
  #
4
  # Licensed under the Apache License, Version 2.0 (the "License");
5
  # you may not use this file except in compliance with the License.
 
12
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
+ """ StableLM model configuration """
16
+
17
+ from transformers.configuration_utils import PretrainedConfig
18
  from transformers.utils import logging
19
 
20
 
21
  logger = logging.get_logger(__name__)
22
 
23
+ STABLELM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
24
+ "stabilityai/stablelm-3b-4e1t": "https://huggingface.co/stabilityai/stablelm-3b-4e1t/resolve/main/config.json",
25
+ # See all StableLM models at https://huggingface.co/models?filter=stablelm
26
+ }
27
+
28
 
29
+ class StableLmConfig(PretrainedConfig):
30
  r"""
31
+ This is the configuration class to store the configuration of a [`~StableLmModel`].
32
+ It is used to instantiate an StableLM model according to the specified arguments, defining the model
33
+ architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of
34
+ the StableLM [stabilityai/stablelm-3b-4e1t](https://huggingface.co/stabilityai/stablelm-3b-4e1t) architecture.
35
+
36
+ Configuration objects inherit from [`PretrainedConfig`] and can be used
37
+ to control the model outputs. Read the documentation from [`PretrainedConfig`]
38
+ for more information.
39
+
40
 
41
  Args:
42
+ vocab_size (`int`, *optional*, defaults to 50304):
43
  Vocabulary size of the StableLM model. Defines the number of different tokens that
44
+ can be represented by the `inputs_ids` passed when calling [`StableLmModel`].
45
  intermediate_size (`int`, *optional*, defaults to 6912):
46
  Dimension of the MLP representations.
47
  hidden_size (`int`, *optional*, defaults to 2560):
 
50
  Number of hidden layers in the Transformer decoder.
51
  num_attention_heads (`int`, *optional*, defaults to 32):
52
  Number of attention heads for each attention layer in the Transformer encoder.
53
+ num_key_value_heads (`int`, *optional*, defaults to 32):
54
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
55
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
56
  `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
 
60
  `num_attention_heads`.
61
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
62
  The non-linear activation function (function or string).
63
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
 
 
 
 
64
  The maximum sequence length that this model might ever be used with.
65
  Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
66
+ initializer_range (`float`, *optional*, defaults to 0.02):
67
  The standard deviation of the truncated_normal_initializer for initializing
68
  all weight matrices.
69
+ layer_norm_eps (`float`, *optional*, defaults to 1e-05):
70
  The epsilon used by the normalization layers.
71
  use_cache (`bool`, *optional*, defaults to `True`):
72
  Whether or not the model should return the last key/values attentions
73
  (not used by all models). Only relevant if `config.is_decoder=True`.
74
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
75
+ Whether the model's input and output word embeddings should be tied.
76
+ rope_theta (`float`, *optional*, defaults to `10000.0`):
77
+ The base period of the RoPE embeddings.
78
+ rope_scaling (`Dict`, *optional*):
79
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
80
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
81
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
82
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
83
+ these scaling strategies behave:
84
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This
85
+ is an experimental feature, subject to breaking API changes in future versions.
86
+ use_qkv_bias (`bool`, *optional*, defaults to `False`):
87
  Whether or not the model should use bias for qkv layers.
88
+ hidden_dropout (`float`, *optional*, defaults to 0.0):
89
+ The dropout ratio after applying the MLP to the hidden states.
90
  attention_dropout (`float`, *optional*, defaults to 0.0):
91
  The dropout ratio for the attention probabilities.
92
+ partial_rotary_factor (`float`, *optional*, defaults to 0.25):
93
+ Percentage of the query and keys which will have rotary embedding.
94
+ bos_token_id (int, *optional*, defaults to 0):
95
+ The id of the `BOS` token in the vocabulary.
96
+ eos_token_id (int, *optional*, defaults to 0):
97
+ The id of the `EOS` token in the vocabulary.
98
+
99
+ Example:
100
+
101
+ ```python
102
+ >>> from transformers import StableLmModel, StableLmConfig
103
+
104
+ >>> # Initializing a StableLM stablelm-3b style configuration
105
+ >>> configuration = StableLmConfig()
106
+ ```"""
107
+
108
+ model_type = "stablelm"
109
  keys_to_ignore_at_inference = ["past_key_values"]
110
 
111
  def __init__(
112
  self,
113
+ vocab_size=50304,
114
  intermediate_size=6912,
115
  hidden_size=2560,
116
  num_hidden_layers=32,
117
  num_attention_heads=32,
118
  num_key_value_heads=32,
119
  hidden_act="silu",
 
 
120
  max_position_embeddings=4096,
121
  initializer_range=0.02,
122
+ layer_norm_eps=1.0e-5,
123
  use_cache=True,
 
 
 
124
  tie_word_embeddings=False,
125
+ rope_theta=10_000,
126
+ rope_scaling=None,
127
+ use_qkv_bias=False,
128
+ hidden_dropout=0.0,
129
+ attention_dropout=0.0,
130
+ partial_rotary_factor=0.25,
131
+ bos_token_id=0,
132
+ eos_token_id=0,
133
  **kwargs,
134
  ):
135
  self.vocab_size = vocab_size
136
  self.max_position_embeddings = max_position_embeddings
 
137
  self.hidden_size = hidden_size
138
+ self.intermediate_size = intermediate_size
139
  self.num_hidden_layers = num_hidden_layers
140
  self.num_attention_heads = num_attention_heads
141
  self.num_key_value_heads = num_key_value_heads
142
  self.hidden_act = hidden_act
 
 
143
  self.initializer_range = initializer_range
144
+ self.layer_norm_eps = layer_norm_eps
145
  self.use_cache = use_cache
146
+ self.rope_theta = rope_theta
147
+ self.rope_scaling = rope_scaling
148
  self.use_qkv_bias = use_qkv_bias
149
+ self.hidden_dropout = hidden_dropout
150
  self.attention_dropout = attention_dropout
151
+ self.partial_rotary_factor = partial_rotary_factor
152
+ self._rope_scaling_validation()
153
+
154
  super().__init__(
155
  bos_token_id=bos_token_id,
156
  eos_token_id=eos_token_id,
157
  tie_word_embeddings=tie_word_embeddings,
158
  **kwargs,
159
  )
160
+
161
+ # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
162
+ def _rope_scaling_validation(self):
163
+ """
164
+ Validate the `rope_scaling` configuration.
165
+ """
166
+ if self.rope_scaling is None:
167
+ return
168
+
169
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
170
+ raise ValueError(
171
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
172
+ f"got {self.rope_scaling}"
173
+ )
174
+ rope_scaling_type = self.rope_scaling.get("type", None)
175
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
176
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
177
+ raise ValueError(
178
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
179
+ )
180
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
181
+ raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
generation_config.json CHANGED
@@ -2,5 +2,5 @@
2
  "_from_model_config": true,
3
  "bos_token_id": 100257,
4
  "eos_token_id": 100257,
5
- "transformers_version": "4.36.2"
6
  }
 
2
  "_from_model_config": true,
3
  "bos_token_id": 100257,
4
  "eos_token_id": 100257,
5
+ "transformers_version": "4.38.0"
6
  }
modeling_stablelm_epoch.py → modeling_stablelm.py RENAMED
@@ -1,5 +1,10 @@
1
  # coding=utf-8
2
- # Copyright 2023 Stability AI, EleutherAI, and The HuggingFace Inc. team. All rights reserved.
 
 
 
 
 
3
  #
4
  # Licensed under the Apache License, Version 2.0 (the "License");
5
  # you may not use this file except in compliance with the License.
@@ -12,48 +17,48 @@
12
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
- #
16
- # This code is based off the following work:
17
- # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
18
- # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
19
- """ PyTorch StableLM Epoch model. """
20
- from typing import Optional, Tuple, Union
21
  import math
22
- import warnings
23
 
24
  import torch
25
  import torch.nn.functional as F
26
  import torch.utils.checkpoint
27
  from torch import nn
28
- from torch.nn import CrossEntropyLoss
29
 
30
- from transformers.cache_utils import Cache
31
- from transformers.modeling_outputs import (
32
- BaseModelOutputWithPast,
33
- CausalLMOutputWithPast,
34
- )
35
  from transformers.modeling_utils import PreTrainedModel
36
- from transformers.utils import logging, is_flash_attn_greater_or_equal_2_10
 
 
 
 
 
 
 
 
37
 
38
- from .configuration_stablelm_epoch import StableLMEpochConfig
39
 
40
- try:
41
  from flash_attn import flash_attn_func, flash_attn_varlen_func
42
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
43
- except:
44
- flash_attn_func, flash_attn_varlen_func = None, None
45
- index_first_axis, pad_input, unpad_input = None, None, None
46
 
47
 
48
  logger = logging.get_logger(__name__)
49
 
 
 
50
 
51
  # Copied from transformers.models.llama.modeling_llama._get_unpad_data
52
  def _get_unpad_data(attention_mask):
53
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
54
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
55
  max_seqlen_in_batch = seqlens_in_batch.max().item()
56
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
57
  return (
58
  indices,
59
  cu_seqlens,
@@ -61,113 +66,144 @@ def _get_unpad_data(attention_mask):
61
  )
62
 
63
 
64
- # Copied from transformers.models.bart.modeling_bart._make_causal_mask
65
- def _make_causal_mask(
66
- input_ids_shape: torch.Size,
67
- dtype: torch.dtype,
68
- device: torch.device,
69
- past_key_values_length: int = 0,
70
- ):
71
- """Make causal mask used for bi-directional self-attention."""
72
- batch_size, tgt_len = input_ids_shape
73
- mask = torch.full((tgt_len, tgt_len), torch.finfo(torch.float16).min, device=device)
74
- mask_cond = torch.arange(mask.size(-1), device=device)
75
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
76
- mask = mask.to(dtype)
77
- if past_key_values_length > 0:
78
- mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
79
- return mask[None, None, :, :].expand(batch_size, 1, tgt_len, tgt_len + past_key_values_length)
80
-
81
-
82
- # Copied from transformers.models.bart.modeling_bart._expand_mask
83
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
84
- """Expands attention_mask from `[batch_size, seq_len]` to `[batch_size, 1, tgt_seq_len, src_seq_len]`."""
85
- batch_size, src_len = mask.size()
86
- tgt_len = tgt_len if tgt_len is not None else src_len
87
-
88
- expanded_mask = mask[:, None, None, :].expand(batch_size, 1, tgt_len, src_len).to(dtype)
89
- inverted_mask = 1.0 - expanded_mask
90
-
91
- return inverted_mask.masked_fill(
92
- inverted_mask.to(torch.bool), torch.finfo(dtype).min
93
- )
94
-
95
-
96
- class RotaryEmbedding(nn.Module):
97
- def __init__(
98
- self,
99
- dim: int,
100
- max_position_embeddings: int,
101
- base: int = 10_000,
102
- device: Optional[torch.device] = None,
103
- ):
104
  super().__init__()
105
 
106
  self.dim = dim
107
  self.max_position_embeddings = max_position_embeddings
108
  self.base = base
109
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
110
  self.register_buffer("inv_freq", inv_freq, persistent=False)
111
 
112
  # Build here to make `torch.jit.trace` work.
113
  self._set_cos_sin_cache(
114
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype(),
115
  )
116
 
117
- def _set_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype):
118
  self.max_seq_len_cached = seq_len
119
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
120
 
121
- # Don't do einsum, it converts fp32 to fp16 under AMP
122
- # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
123
  freqs = torch.outer(t, self.inv_freq)
124
  # Different from paper, but it uses a different permutation in order to obtain the same calculation
125
  emb = torch.cat((freqs, freqs), dim=-1)
126
- self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
127
- self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
128
 
129
- def forward(self, x: torch.Tensor, seq_len: Optional[int] = None):
130
- # x: [batch_size, num_heads, seq_len, head_size]
131
  if seq_len > self.max_seq_len_cached:
132
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=torch.get_default_dtype())
 
133
  return (
134
- self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
135
- self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
136
  )
137
 
138
 
139
- def rotate_half(x: torch.Tensor):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
140
  """Rotates half the hidden dims of the input."""
141
- x1, x2 = torch.chunk(x, 2, dim=-1)
 
142
  return torch.cat((-x2, x1), dim=-1)
143
 
144
 
145
- def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
146
- # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
147
- cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
148
- sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
149
- cos = cos[position_ids].unsqueeze(1) # [batch_size, 1, seq_len, dim]
150
- sin = sin[position_ids].unsqueeze(1) # [batch_size, 1, seq_len, dim]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
151
  q_embed = (q * cos) + (rotate_half(q) * sin)
152
  k_embed = (k * cos) + (rotate_half(k) * sin)
153
  return q_embed, k_embed
154
 
155
 
156
- class MLP(nn.Module):
157
- def __init__(self, config: StableLMEpochConfig):
 
158
  super().__init__()
159
  self.config = config
160
  self.hidden_size = config.hidden_size
161
  self.intermediate_size = config.intermediate_size
162
- self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
163
- self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
164
- self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
165
- self.act_fn = nn.SiLU()
166
 
167
- def forward(self, x: torch.Tensor) -> torch.Tensor:
168
  return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
169
 
170
 
 
171
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
172
  """
173
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -180,48 +216,79 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
180
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
181
 
182
 
183
- class Attention(nn.Module):
184
- def __init__(self, config: StableLMEpochConfig):
 
 
185
  super().__init__()
186
  self.config = config
 
 
 
 
 
 
 
 
187
  self.hidden_size = config.hidden_size
188
  self.num_heads = config.num_attention_heads
189
  self.head_dim = self.hidden_size // self.num_heads
190
  self.num_key_value_heads = config.num_key_value_heads
191
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
192
  self.max_position_embeddings = config.max_position_embeddings
 
 
193
  self.is_causal = True
194
- self.attention_dropout = config.attention_dropout
195
 
196
  if (self.head_dim * self.num_heads) != self.hidden_size:
197
  raise ValueError(
198
  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
199
  f" and `num_heads`: {self.num_heads})."
200
  )
201
-
202
  self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.use_qkv_bias)
203
  self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_qkv_bias)
204
  self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_qkv_bias)
205
  self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
206
 
 
207
  self._init_rope()
208
 
 
209
  def _init_rope(self):
210
- self.rotary_ndims = int(self.head_dim * self.config.rope_pct)
211
- self.rotary_emb = RotaryEmbedding(
212
- self.rotary_ndims,
213
- max_position_embeddings=self.config.max_position_embeddings,
214
- base=self.config.rope_theta,
215
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
216
 
217
  def forward(
218
  self,
219
- hidden_states: torch.FloatTensor,
220
- attention_mask: torch.FloatTensor,
221
- position_ids: torch.LongTensor,
222
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
223
- output_attentions: Optional[bool] = False,
224
- use_cache: Optional[bool] = False,
225
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
226
  bsz, q_len, _ = hidden_states.size()
227
 
@@ -233,27 +300,37 @@ class Attention(nn.Module):
233
  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
234
  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
235
 
236
- query_rot = query_states[..., : self.rotary_ndims]
237
- query_pass = query_states[..., self.rotary_ndims :]
238
- key_rot = key_states[..., : self.rotary_ndims]
239
- key_pass = key_states[..., self.rotary_ndims :]
240
-
241
  kv_seq_len = key_states.shape[-2]
242
  if past_key_value is not None:
243
- kv_seq_len += past_key_value[0].shape[-2]
 
 
 
 
 
 
244
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
245
- query_states, key_states = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
246
 
247
- # [batch_size, num_heads, seq_len, head_dim]
248
- query_states = torch.cat((query_states, query_pass), dim=-1)
249
- key_states = torch.cat((key_states, key_pass), dim=-1)
 
 
 
 
 
 
 
 
250
 
251
- if past_key_value is not None:
252
- # Reuse k, v, self_attention
253
- key_states = torch.cat((past_key_value[0], key_states), dim=2)
254
- value_states = torch.cat((past_key_value[1], value_states), dim=2)
255
 
256
- past_key_value = (key_states, value_states) if use_cache else None
 
 
 
257
 
258
  # Repeat k/v heads if n_kv_heads < n_heads
259
  key_states = repeat_kv(key_states, self.num_key_value_groups)
@@ -274,9 +351,10 @@ class Attention(nn.Module):
274
  )
275
  attn_weights = attn_weights + attention_mask
276
 
277
- # Upcast attention to fp32
278
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
279
- attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
 
280
  attn_output = torch.matmul(attn_weights, value_states)
281
 
282
  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
@@ -285,11 +363,9 @@ class Attention(nn.Module):
285
  f" {attn_output.size()}"
286
  )
287
 
288
- # Merge heads
289
  attn_output = attn_output.transpose(1, 2).contiguous()
290
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
291
 
292
- # Final linear projection
293
  attn_output = self.o_proj(attn_output)
294
 
295
  if not output_attentions:
@@ -298,11 +374,14 @@ class Attention(nn.Module):
298
  return attn_output, attn_weights, past_key_value
299
 
300
 
301
- class FlashAttention2(Attention):
302
  """
303
- Reference: https://github.com/huggingface/transformers/blob/5d36025ca13d05151b7a0c761e90d429c4644a30/src/transformers/models/llama/modeling_llama.py#L456
 
 
304
  """
305
 
 
306
  def __init__(self, *args, **kwargs):
307
  super().__init__(*args, **kwargs)
308
 
@@ -321,14 +400,7 @@ class FlashAttention2(Attention):
321
  use_cache: bool = False,
322
  **kwargs,
323
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
324
- # FlashAttention2 attention does not support output_attentions
325
- if "padding_mask" in kwargs:
326
- warnings.warn(
327
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
328
- )
329
-
330
- # overwrite attention_mask with padding_mask
331
- attention_mask = kwargs.pop("padding_mask")
332
 
333
  output_attentions = False
334
 
@@ -345,27 +417,35 @@ class FlashAttention2(Attention):
345
  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
346
  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
347
 
348
- query_rot = query_states[..., : self.rotary_ndims]
349
- query_pass = query_states[..., self.rotary_ndims :]
350
- key_rot = key_states[..., : self.rotary_ndims]
351
- key_pass = key_states[..., self.rotary_ndims :]
352
-
353
  kv_seq_len = key_states.shape[-2]
354
  if past_key_value is not None:
355
- kv_seq_len += past_key_value[0].shape[-2]
 
 
 
 
 
 
356
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
357
- query_states, key_states = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
358
 
359
- # [batch_size, num_heads, seq_len, head_dim]
360
- query_states = torch.cat((query_states, query_pass), dim=-1)
361
- key_states = torch.cat((key_states, key_pass), dim=-1)
 
 
 
 
 
 
 
 
 
 
 
362
 
363
  if past_key_value is not None:
364
- # Reuse k, v, self_attention
365
- key_states = torch.cat((past_key_value[0], key_states), dim=2)
366
- value_states = torch.cat((past_key_value[1], value_states), dim=2)
367
-
368
- past_key_value = (key_states, value_states) if use_cache else None
369
 
370
  # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
371
  # to be able to avoid many of these transpose/reshape/view.
@@ -376,8 +456,14 @@ class FlashAttention2(Attention):
376
  dropout_rate = self.attention_dropout if self.training else 0.0
377
 
378
  attn_output = self._flash_attention_forward(
379
- query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
 
 
 
 
 
380
  )
 
381
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
382
  attn_output = self.o_proj(attn_output)
383
 
@@ -386,6 +472,7 @@ class FlashAttention2(Attention):
386
 
387
  return attn_output, attn_weights, past_key_value
388
 
 
389
  def _flash_attention_forward(
390
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
391
  ):
@@ -411,7 +498,7 @@ class FlashAttention2(Attention):
411
  if not self._flash_attn_uses_top_left_mask:
412
  causal = self.is_causal
413
  else:
414
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in FlashAttention2 __init__.
415
  causal = self.is_causal and query_length != 1
416
 
417
  # Contains at least one padding token in the sequence
@@ -445,6 +532,7 @@ class FlashAttention2(Attention):
445
 
446
  return attn_output
447
 
 
448
  def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
449
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
450
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
@@ -485,28 +573,50 @@ class FlashAttention2(Attention):
485
 
486
 
487
  ATTENTION_CLASSES = {
488
- "eager": Attention,
489
- "flash_attention_2": FlashAttention2,
490
  }
491
 
492
 
493
- class DecoderLayer(nn.Module):
494
- def __init__(self, config: StableLMEpochConfig):
495
  super().__init__()
496
- self.self_attn = ATTENTION_CLASSES[config._attn_implementation](config=config)
497
- self.mlp = MLP(config)
498
- self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
499
- self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
 
 
500
 
501
  def forward(
502
  self,
503
- hidden_states: Optional[torch.FloatTensor],
504
- attention_mask: Optional[torch.FloatTensor] = None,
505
  position_ids: Optional[torch.LongTensor] = None,
506
  past_key_value: Optional[Tuple[torch.Tensor]] = None,
507
  output_attentions: Optional[bool] = False,
508
  use_cache: Optional[bool] = False,
509
- ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
510
  residual = hidden_states
511
 
512
  hidden_states = self.input_layernorm(hidden_states)
@@ -526,7 +636,9 @@ class DecoderLayer(nn.Module):
526
  residual = hidden_states
527
  hidden_states = self.post_attention_layernorm(hidden_states)
528
  hidden_states = self.mlp(hidden_states)
529
- hidden_states = residual + hidden_states
 
 
530
 
531
  outputs = (hidden_states,)
532
 
@@ -539,45 +651,142 @@ class DecoderLayer(nn.Module):
539
  return outputs
540
 
541
 
542
- class StableLMEpochPreTrainedModel(PreTrainedModel):
543
- """An abstract class to handle weights initialization and a simple interface
544
- for downloading and loading pretrained models.
545
- """
 
 
 
 
 
 
 
 
 
 
 
 
546
 
547
- config_class = StableLMEpochConfig
 
 
 
 
 
548
  base_model_prefix = "model"
549
  supports_gradient_checkpointing = True
550
- _no_split_modules = ["DecoderLayer"]
551
  _skip_keys_device_placement = "past_key_values"
552
  _supports_flash_attn_2 = True
 
553
 
554
- def _init_weights(self, module: nn.Module):
555
- """Initialize the weights"""
556
  if isinstance(module, nn.Linear):
557
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
558
  if module.bias is not None:
559
  module.bias.data.zero_()
560
  elif isinstance(module, nn.Embedding):
561
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
562
  if module.padding_idx is not None:
563
  module.weight.data[module.padding_idx].zero_()
564
- elif isinstance(module, nn.LayerNorm):
565
- module.bias.data.zero_()
566
- module.weight.data.fill_(1.0)
567
 
568
- def _set_gradient_checkpointing(self, module: nn.Module, value=False):
569
- if isinstance(module, StableLMEpochModel):
570
- module.gradient_checkpointing = value
571
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
572
 
573
- class StableLMEpochModel(StableLMEpochPreTrainedModel):
574
- def __init__(self, config: StableLMEpochConfig):
 
 
 
575
  super().__init__(config)
576
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
577
- self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
578
- self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
 
 
 
 
 
579
 
580
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
581
  self.gradient_checkpointing = False
582
  # Initialize weights and apply final processing
583
  self.post_init()
@@ -585,43 +794,16 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
585
  def get_input_embeddings(self):
586
  return self.embed_tokens
587
 
588
- def set_input_embeddings(self, value: nn.Module):
589
  self.embed_tokens = value
590
 
591
- # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
592
- def _prepare_decoder_attention_mask(
593
- self,
594
- attention_mask: torch.Tensor,
595
- input_shape: torch.Size,
596
- inputs_embeds: torch.Tensor,
597
- past_key_values_length: int,
598
- ):
599
- # Create causal mask
600
- # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
601
- combined_attention_mask = None
602
- if input_shape[-1] > 1:
603
- combined_attention_mask = _make_causal_mask(
604
- input_shape,
605
- inputs_embeds.dtype,
606
- device=inputs_embeds.device,
607
- past_key_values_length=past_key_values_length,
608
- )
609
-
610
- if attention_mask is not None:
611
- # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
612
- expanded_attn_mask = _expand_mask(
613
- attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
614
- ).to(inputs_embeds.device)
615
- combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
616
-
617
- return combined_attention_mask
618
-
619
  def forward(
620
  self,
621
- input_ids: Optional[torch.LongTensor] = None,
622
- attention_mask: Optional[torch.FloatTensor] = None,
623
  position_ids: Optional[torch.LongTensor] = None,
624
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
625
  inputs_embeds: Optional[torch.FloatTensor] = None,
626
  use_cache: Optional[bool] = None,
627
  output_attentions: Optional[bool] = None,
@@ -629,103 +811,85 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
629
  return_dict: Optional[bool] = None,
630
  ) -> Union[Tuple, BaseModelOutputWithPast]:
631
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
632
- output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 
 
633
  use_cache = use_cache if use_cache is not None else self.config.use_cache
634
 
635
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
636
 
637
- # Retrieve input_ids and inputs_embeds
638
  if input_ids is not None and inputs_embeds is not None:
639
- raise ValueError(
640
- "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
641
- )
642
  elif input_ids is not None:
643
  batch_size, seq_length = input_ids.shape
644
  elif inputs_embeds is not None:
645
  batch_size, seq_length, _ = inputs_embeds.shape
646
  else:
647
- raise ValueError(
648
- "You have to specify either decoder_input_ids or decoder_inputs_embeds"
649
- )
650
 
651
  seq_length_with_past = seq_length
652
  past_key_values_length = 0
653
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
654
  if position_ids is None:
655
  device = input_ids.device if input_ids is not None else inputs_embeds.device
656
  position_ids = torch.arange(
657
- past_key_values_length,
658
- seq_length + past_key_values_length,
659
- dtype=torch.long,
660
- device=device,
661
  )
662
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
663
- else:
664
- position_ids = position_ids.view(-1, seq_length).long()
665
 
666
  if inputs_embeds is None:
667
  inputs_embeds = self.embed_tokens(input_ids)
668
- # Embed positions
669
- if self._use_flash_attention_2:
670
  # 2d mask is passed through the layers
671
  attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
672
  else:
673
- if attention_mask is None:
674
- attention_mask = torch.ones(
675
- (batch_size, seq_length_with_past),
676
- dtype=torch.bool,
677
- device=inputs_embeds.device,
678
- )
679
- attention_mask = self._prepare_decoder_attention_mask(
680
- attention_mask,
681
- (batch_size, seq_length),
682
- inputs_embeds,
683
- past_key_values_length,
684
  )
685
 
686
  hidden_states = inputs_embeds
687
 
688
- if self.gradient_checkpointing and self.training:
689
- if use_cache:
690
- logger.warning(
691
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
692
- )
693
- use_cache = False
694
-
695
- # Decoder layers
696
  all_hidden_states = () if output_hidden_states else None
697
  all_self_attns = () if output_attentions else None
698
- next_decoder_cache = () if use_cache else None
699
 
700
- for idx, decoder_layer in enumerate(self.layers):
701
  if output_hidden_states:
702
  all_hidden_states += (hidden_states,)
703
 
704
- past_key_value = (
705
- past_key_values[idx] if past_key_values is not None else None
706
- )
707
-
708
  if self.gradient_checkpointing and self.training:
709
-
710
- def create_custom_forward(module):
711
- def custom_forward(*inputs):
712
- # None for past_key_value
713
- return module(*inputs, past_key_value, output_attentions)
714
-
715
- return custom_forward
716
-
717
- layer_outputs = torch.utils.checkpoint.checkpoint(
718
- create_custom_forward(decoder_layer),
719
  hidden_states,
720
  attention_mask,
721
  position_ids,
 
 
722
  )
723
  else:
724
  layer_outputs = decoder_layer(
725
  hidden_states,
726
  attention_mask=attention_mask,
727
  position_ids=position_ids,
728
- past_key_value=past_key_value,
729
  output_attentions=output_attentions,
730
  use_cache=use_cache,
731
  )
@@ -733,24 +897,23 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
733
  hidden_states = layer_outputs[0]
734
 
735
  if use_cache:
736
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
737
 
738
  if output_attentions:
739
  all_self_attns += (layer_outputs[1],)
740
 
741
  hidden_states = self.norm(hidden_states)
742
 
743
- # Add hidden states from the last decoder layer
744
  if output_hidden_states:
745
  all_hidden_states += (hidden_states,)
746
 
747
- next_cache = next_decoder_cache if use_cache else None
 
 
 
748
  if not return_dict:
749
- return tuple(
750
- v
751
- for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
752
- if v is not None
753
- )
754
  return BaseModelOutputWithPast(
755
  last_hidden_state=hidden_states,
756
  past_key_values=next_cache,
@@ -759,42 +922,53 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
759
  )
760
 
761
 
762
- class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
 
763
  _tied_weights_keys = ["lm_head.weight"]
764
 
765
- def __init__(self, config: StableLMEpochConfig):
 
766
  super().__init__(config)
767
-
768
- self.model = StableLMEpochModel(config)
769
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
770
 
771
  # Initialize weights and apply final processing
772
  self.post_init()
773
 
 
774
  def get_input_embeddings(self):
775
  return self.model.embed_tokens
776
 
 
777
  def set_input_embeddings(self, value):
778
  self.model.embed_tokens = value
779
 
 
780
  def get_output_embeddings(self):
781
  return self.lm_head
782
 
783
- def set_output_embeddings(self, new_embeddings: nn.Module):
 
784
  self.lm_head = new_embeddings
785
 
786
- def get_decoder(self):
787
- return self.model
788
-
789
  def set_decoder(self, decoder):
790
  self.model = decoder
791
 
 
 
 
 
 
 
 
792
  def forward(
793
  self,
794
- input_ids: Optional[torch.LongTensor] = None,
795
- attention_mask: Optional[torch.FloatTensor] = None,
796
  position_ids: Optional[torch.LongTensor] = None,
797
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
798
  inputs_embeds: Optional[torch.FloatTensor] = None,
799
  labels: Optional[torch.LongTensor] = None,
800
  use_cache: Optional[bool] = None,
@@ -802,23 +976,40 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
802
  output_hidden_states: Optional[bool] = None,
803
  return_dict: Optional[bool] = None,
804
  ) -> Union[Tuple, CausalLMOutputWithPast]:
805
- output_attentions = (
806
- output_attentions
807
- if output_attentions is not None
808
- else self.config.output_attentions
809
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
810
  output_hidden_states = (
811
- output_hidden_states
812
- if output_hidden_states is not None
813
- else self.config.output_hidden_states
814
- )
815
- return_dict = (
816
- return_dict if return_dict is not None else self.config.use_return_dict
817
  )
 
818
 
819
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
820
  outputs = self.model(
821
- input_ids,
822
  attention_mask=attention_mask,
823
  position_ids=position_ids,
824
  past_key_values=past_key_values,
@@ -830,7 +1021,7 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
830
  )
831
 
832
  hidden_states = outputs[0]
833
- logits = self.lm_head(hidden_states).float()
834
 
835
  loss = None
836
  if labels is not None:
@@ -857,36 +1048,54 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
857
  attentions=outputs.attentions,
858
  )
859
 
 
860
  def prepare_inputs_for_generation(
861
- self,
862
- input_ids,
863
- past_key_values: Optional[torch.Tensor] = None,
864
- attention_mask: Optional[torch.Tensor] = None,
865
- inputs_embeds: Optional[torch.Tensor] = None,
866
- **kwargs,
867
  ):
868
- # Trim decoder_input_ids if past is used
869
  if past_key_values is not None:
870
- past_length = past_key_values[0][0].shape[2]
871
-
872
- # Some generation methods already pass only the last input ID
873
- if input_ids.shape[1] > past_length:
874
- remove_prefix_length = past_length
875
  else:
876
- # Default to old behavior: keep only final ID
877
- remove_prefix_length = input_ids.shape[1] - 1
878
-
879
- input_ids = input_ids[:, remove_prefix_length:]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
880
 
881
  position_ids = kwargs.get("position_ids", None)
882
  if attention_mask is not None and position_ids is None:
883
- # Create position_ids on the fly for batch generation
884
  position_ids = attention_mask.long().cumsum(-1) - 1
885
  position_ids.masked_fill_(attention_mask == 0, 1)
886
  if past_key_values:
887
- position_ids = position_ids[:, -1].unsqueeze(-1)
 
 
 
 
 
 
888
 
889
- # If `inputs_embeds` are passed, we only want to use them in the 1st generation step
890
  if inputs_embeds is not None and past_key_values is None:
891
  model_inputs = {"inputs_embeds": inputs_embeds}
892
  else:
@@ -894,10 +1103,10 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
894
 
895
  model_inputs.update(
896
  {
897
- "attention_mask": attention_mask,
898
  "past_key_values": past_key_values,
899
  "use_cache": kwargs.get("use_cache"),
900
- "position_ids": position_ids,
901
  }
902
  )
903
  return model_inputs
@@ -907,13 +1116,130 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
907
  reordered_past = ()
908
  for layer_past in past_key_values:
909
  reordered_past += (
910
- tuple(
911
- past_state.index_select(0, beam_idx.to(past_state.device))
912
- for past_state in layer_past
913
- ),
914
  )
915
  return reordered_past
916
 
917
 
918
- StableLMEpochConfig.register_for_auto_class()
919
- StableLMEpochForCausalLM.register_for_auto_class("AutoModelForCausalLM")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # coding=utf-8
2
+ # Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
  #
9
  # Licensed under the Apache License, Version 2.0 (the "License");
10
  # you may not use this file except in compliance with the License.
 
17
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
  # See the License for the specific language governing permissions and
19
  # limitations under the License.
20
+ """ PyTorch StableLM model."""
 
 
 
 
 
21
  import math
22
+ from typing import List, Optional, Tuple, Union
23
 
24
  import torch
25
  import torch.nn.functional as F
26
  import torch.utils.checkpoint
27
  from torch import nn
28
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
29
 
30
+ from transformers.activations import ACT2FN
31
+ from transformers.cache_utils import Cache, DynamicCache
32
+ from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
33
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 
34
  from transformers.modeling_utils import PreTrainedModel
35
+ from transformers.utils import (
36
+ add_start_docstrings,
37
+ add_start_docstrings_to_model_forward,
38
+ is_flash_attn_2_available,
39
+ is_flash_attn_greater_or_equal_2_10,
40
+ logging,
41
+ replace_return_docstrings,
42
+ )
43
+ from .configuration_stablelm import StableLmConfig
44
 
 
45
 
46
+ if is_flash_attn_2_available():
47
  from flash_attn import flash_attn_func, flash_attn_varlen_func
48
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
 
 
 
49
 
50
 
51
  logger = logging.get_logger(__name__)
52
 
53
+ _CONFIG_FOR_DOC = "StableLmConfig"
54
+
55
 
56
  # Copied from transformers.models.llama.modeling_llama._get_unpad_data
57
  def _get_unpad_data(attention_mask):
58
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
59
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
60
  max_seqlen_in_batch = seqlens_in_batch.max().item()
61
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
62
  return (
63
  indices,
64
  cu_seqlens,
 
66
  )
67
 
68
 
69
+ # Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->StableLm
70
+ class StableLmRotaryEmbedding(nn.Module):
71
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  super().__init__()
73
 
74
  self.dim = dim
75
  self.max_position_embeddings = max_position_embeddings
76
  self.base = base
77
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
78
  self.register_buffer("inv_freq", inv_freq, persistent=False)
79
 
80
  # Build here to make `torch.jit.trace` work.
81
  self._set_cos_sin_cache(
82
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
83
  )
84
 
85
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
86
  self.max_seq_len_cached = seq_len
87
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
88
 
 
 
89
  freqs = torch.outer(t, self.inv_freq)
90
  # Different from paper, but it uses a different permutation in order to obtain the same calculation
91
  emb = torch.cat((freqs, freqs), dim=-1)
92
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
93
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
94
 
95
+ def forward(self, x, seq_len=None):
96
+ # x: [bs, num_attention_heads, seq_len, head_size]
97
  if seq_len > self.max_seq_len_cached:
98
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
99
+
100
  return (
101
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
102
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
103
  )
104
 
105
 
106
+ # Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->StableLm
107
+ class StableLmLinearScalingRotaryEmbedding(StableLmRotaryEmbedding):
108
+ """StableLmRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
109
+
110
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
111
+ self.scaling_factor = scaling_factor
112
+ super().__init__(dim, max_position_embeddings, base, device)
113
+
114
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
115
+ self.max_seq_len_cached = seq_len
116
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
117
+ t = t / self.scaling_factor
118
+
119
+ freqs = torch.outer(t, self.inv_freq)
120
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
121
+ emb = torch.cat((freqs, freqs), dim=-1)
122
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
123
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
124
+
125
+
126
+ # Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->StableLm
127
+ class StableLmDynamicNTKScalingRotaryEmbedding(StableLmRotaryEmbedding):
128
+ """StableLmRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
129
+
130
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
131
+ self.scaling_factor = scaling_factor
132
+ super().__init__(dim, max_position_embeddings, base, device)
133
+
134
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
135
+ self.max_seq_len_cached = seq_len
136
+
137
+ if seq_len > self.max_position_embeddings:
138
+ base = self.base * (
139
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
140
+ ) ** (self.dim / (self.dim - 2))
141
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
142
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
143
+
144
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
145
+
146
+ freqs = torch.outer(t, self.inv_freq)
147
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
148
+ emb = torch.cat((freqs, freqs), dim=-1)
149
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
150
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
151
+
152
+
153
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
154
+ def rotate_half(x):
155
  """Rotates half the hidden dims of the input."""
156
+ x1 = x[..., : x.shape[-1] // 2]
157
+ x2 = x[..., x.shape[-1] // 2 :]
158
  return torch.cat((-x2, x1), dim=-1)
159
 
160
 
161
+ # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
162
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
163
+ """Applies Rotary Position Embedding to the query and key tensors.
164
+
165
+ Args:
166
+ q (`torch.Tensor`): The query tensor.
167
+ k (`torch.Tensor`): The key tensor.
168
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
169
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
170
+ position_ids (`torch.Tensor`):
171
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
172
+ used to pass offsetted position ids when working with a KV-cache.
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[position_ids].unsqueeze(unsqueeze_dim)
184
+ sin = sin[position_ids].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
+ # Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->StableLm
191
+ class StableLmMLP(nn.Module):
192
+ def __init__(self, config):
193
  super().__init__()
194
  self.config = config
195
  self.hidden_size = config.hidden_size
196
  self.intermediate_size = config.intermediate_size
197
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
198
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
199
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
200
+ self.act_fn = ACT2FN[config.hidden_act]
201
 
202
+ def forward(self, x):
203
  return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
204
 
205
 
206
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
207
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
208
  """
209
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
 
216
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
217
 
218
 
219
+ class StableLmAttention(nn.Module):
220
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
221
+
222
+ def __init__(self, config: StableLmConfig, layer_idx: Optional[int] = None):
223
  super().__init__()
224
  self.config = config
225
+ self.layer_idx = layer_idx
226
+ if layer_idx is None:
227
+ logger.warning_once(
228
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
229
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
230
+ "when creating this class."
231
+ )
232
+
233
  self.hidden_size = config.hidden_size
234
  self.num_heads = config.num_attention_heads
235
  self.head_dim = self.hidden_size // self.num_heads
236
  self.num_key_value_heads = config.num_key_value_heads
237
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
238
  self.max_position_embeddings = config.max_position_embeddings
239
+ self.rope_theta = config.rope_theta
240
+ self.partial_rotary_factor = config.partial_rotary_factor
241
  self.is_causal = True
 
242
 
243
  if (self.head_dim * self.num_heads) != self.hidden_size:
244
  raise ValueError(
245
  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
246
  f" and `num_heads`: {self.num_heads})."
247
  )
 
248
  self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.use_qkv_bias)
249
  self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_qkv_bias)
250
  self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.use_qkv_bias)
251
  self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
252
 
253
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
254
  self._init_rope()
255
 
256
+ # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonAttention._init_rope with Persimmon->StableLm
257
  def _init_rope(self):
258
+ if self.config.rope_scaling is None:
259
+ self.rotary_emb = StableLmRotaryEmbedding(
260
+ int(self.partial_rotary_factor * self.head_dim),
261
+ max_position_embeddings=self.max_position_embeddings,
262
+ base=self.rope_theta,
263
+ )
264
+ else:
265
+ scaling_type = self.config.rope_scaling["type"]
266
+ scaling_factor = self.config.rope_scaling["factor"]
267
+ if scaling_type == "linear":
268
+ self.rotary_emb = StableLmLinearScalingRotaryEmbedding(
269
+ int(self.partial_rotary_factor * self.head_dim),
270
+ max_position_embeddings=self.max_position_embeddings,
271
+ scaling_factor=scaling_factor,
272
+ base=self.rope_theta,
273
+ )
274
+ elif scaling_type == "dynamic":
275
+ self.rotary_emb = StableLmDynamicNTKScalingRotaryEmbedding(
276
+ int(self.partial_rotary_factor * self.head_dim),
277
+ max_position_embeddings=self.max_position_embeddings,
278
+ scaling_factor=scaling_factor,
279
+ base=self.rope_theta,
280
+ )
281
+ else:
282
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
283
 
284
  def forward(
285
  self,
286
+ hidden_states: torch.Tensor,
287
+ attention_mask: Optional[torch.Tensor] = None,
288
+ position_ids: Optional[torch.LongTensor] = None,
289
+ past_key_value: Optional[Cache] = None,
290
+ output_attentions: bool = False,
291
+ use_cache: bool = False,
292
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
293
  bsz, q_len, _ = hidden_states.size()
294
 
 
300
  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
301
  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
302
 
 
 
 
 
 
303
  kv_seq_len = key_states.shape[-2]
304
  if past_key_value is not None:
305
+ if self.layer_idx is None:
306
+ raise ValueError(
307
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
308
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
309
+ "with a layer index."
310
+ )
311
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
312
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
 
313
 
314
+ # Partial rotary embedding
315
+ query_rot, query_pass = (
316
+ query_states[..., : self.rotary_emb.dim],
317
+ query_states[..., self.rotary_emb.dim :],
318
+ )
319
+ key_rot, key_pass = (
320
+ key_states[..., : self.rotary_emb.dim],
321
+ key_states[..., self.rotary_emb.dim :],
322
+ )
323
+ # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
324
+ query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
325
 
326
+ # [batch_size, seq_length, num_heads, head_dim]
327
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
328
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
 
329
 
330
+ if past_key_value is not None:
331
+ # Specific to RoPE models with partial rotation
332
+ cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
333
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
334
 
335
  # Repeat k/v heads if n_kv_heads < n_heads
336
  key_states = repeat_kv(key_states, self.num_key_value_groups)
 
351
  )
352
  attn_weights = attn_weights + attention_mask
353
 
354
+ # upcast attention to fp32
355
+ attn_weights = nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query_states.dtype)
356
+ attn_weights = self.attention_dropout(attn_weights)
357
+
358
  attn_output = torch.matmul(attn_weights, value_states)
359
 
360
  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
 
363
  f" {attn_output.size()}"
364
  )
365
 
 
366
  attn_output = attn_output.transpose(1, 2).contiguous()
367
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
368
 
 
369
  attn_output = self.o_proj(attn_output)
370
 
371
  if not output_attentions:
 
374
  return attn_output, attn_weights, past_key_value
375
 
376
 
377
+ class StableLmFlashAttention2(StableLmAttention):
378
  """
379
+ StableLM flash attention module. This module inherits from `StableLmAttention` as the weights of the module stays
380
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
381
+ flash attention and deal with padding tokens in case the input contains any of them.
382
  """
383
 
384
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
385
  def __init__(self, *args, **kwargs):
386
  super().__init__(*args, **kwargs)
387
 
 
400
  use_cache: bool = False,
401
  **kwargs,
402
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
403
+ # StableLmFlashAttention2 attention does not support output_attentions
 
 
 
 
 
 
 
404
 
405
  output_attentions = False
406
 
 
417
  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
418
  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
419
 
 
 
 
 
 
420
  kv_seq_len = key_states.shape[-2]
421
  if past_key_value is not None:
422
+ if self.layer_idx is None:
423
+ raise ValueError(
424
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
425
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
426
+ "with a layer index."
427
+ )
428
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
429
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
 
430
 
431
+ # Partial rotary embedding
432
+ query_rot, query_pass = (
433
+ query_states[..., : self.rotary_emb.dim],
434
+ query_states[..., self.rotary_emb.dim :],
435
+ )
436
+ key_rot, key_pass = (
437
+ key_states[..., : self.rotary_emb.dim],
438
+ key_states[..., self.rotary_emb.dim :],
439
+ )
440
+ query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
441
+
442
+ # [batch_size, seq_length, num_heads, head_dim]
443
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
444
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
445
 
446
  if past_key_value is not None:
447
+ cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
448
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
 
449
 
450
  # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
451
  # to be able to avoid many of these transpose/reshape/view.
 
456
  dropout_rate = self.attention_dropout if self.training else 0.0
457
 
458
  attn_output = self._flash_attention_forward(
459
+ query_states,
460
+ key_states,
461
+ value_states,
462
+ attention_mask,
463
+ q_len,
464
+ dropout=dropout_rate,
465
  )
466
+
467
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
468
  attn_output = self.o_proj(attn_output)
469
 
 
472
 
473
  return attn_output, attn_weights, past_key_value
474
 
475
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
476
  def _flash_attention_forward(
477
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
478
  ):
 
498
  if not self._flash_attn_uses_top_left_mask:
499
  causal = self.is_causal
500
  else:
501
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
502
  causal = self.is_causal and query_length != 1
503
 
504
  # Contains at least one padding token in the sequence
 
532
 
533
  return attn_output
534
 
535
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
536
  def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
537
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
538
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
 
573
 
574
 
575
  ATTENTION_CLASSES = {
576
+ "eager": StableLmAttention,
577
+ "flash_attention_2": StableLmFlashAttention2,
578
  }
579
 
580
 
581
+ class StableLmDecoderLayer(nn.Module):
582
+ def __init__(self, config: StableLmConfig, layer_idx: int):
583
  super().__init__()
584
+ self.hidden_size = config.hidden_size
585
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
586
+ self.mlp = StableLmMLP(config)
587
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
588
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
589
+ self.dropout = nn.Dropout(config.hidden_dropout)
590
 
591
  def forward(
592
  self,
593
+ hidden_states: torch.Tensor,
594
+ attention_mask: Optional[torch.Tensor] = None,
595
  position_ids: Optional[torch.LongTensor] = None,
596
  past_key_value: Optional[Tuple[torch.Tensor]] = None,
597
  output_attentions: Optional[bool] = False,
598
  use_cache: Optional[bool] = False,
599
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
600
+ """
601
+ Args:
602
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
603
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
604
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
605
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
606
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
607
+ `[0, config.n_positions - 1]`.
608
+
609
+ [What are position IDs?](../glossary#position-ids)
610
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
611
+ cached past key and value projection states
612
+ output_attentions (`bool`, *optional*):
613
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
614
+ returned tensors for more detail.
615
+ use_cache (`bool`, *optional*):
616
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
617
+ (see `past_key_values`).
618
+ """
619
+
620
  residual = hidden_states
621
 
622
  hidden_states = self.input_layernorm(hidden_states)
 
636
  residual = hidden_states
637
  hidden_states = self.post_attention_layernorm(hidden_states)
638
  hidden_states = self.mlp(hidden_states)
639
+
640
+ hidden_states = self.dropout(hidden_states)
641
+ hidden_states = hidden_states + residual
642
 
643
  outputs = (hidden_states,)
644
 
 
651
  return outputs
652
 
653
 
654
+ STABLELM_START_DOCSTRING = r"""
655
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
656
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
657
+ etc.)
658
+
659
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
660
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
661
+ and behavior.
662
+
663
+ Parameters:
664
+ config ([`StableLmConfig`]):
665
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
666
+ load the weights associated with the model, only the configuration. Check out the
667
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
668
+ """
669
+
670
 
671
+ @add_start_docstrings(
672
+ "The bare StableLM Model outputting raw hidden-states without any specific head on top.",
673
+ STABLELM_START_DOCSTRING,
674
+ )
675
+ class StableLmPreTrainedModel(PreTrainedModel):
676
+ config_class = StableLmConfig
677
  base_model_prefix = "model"
678
  supports_gradient_checkpointing = True
679
+ _no_split_modules = ["StableLmDecoderLayer"]
680
  _skip_keys_device_placement = "past_key_values"
681
  _supports_flash_attn_2 = True
682
+ _supports_cache_class = True
683
 
684
+ def _init_weights(self, module):
685
+ std = self.config.initializer_range
686
  if isinstance(module, nn.Linear):
687
+ module.weight.data.normal_(mean=0.0, std=std)
688
  if module.bias is not None:
689
  module.bias.data.zero_()
690
  elif isinstance(module, nn.Embedding):
691
+ module.weight.data.normal_(mean=0.0, std=std)
692
  if module.padding_idx is not None:
693
  module.weight.data[module.padding_idx].zero_()
 
 
 
694
 
 
 
 
695
 
696
+ STABLELM_INPUTS_DOCSTRING = r"""
697
+ Args:
698
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
699
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
700
+ it.
701
+
702
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
703
+ [`PreTrainedTokenizer.__call__`] for details.
704
+
705
+ [What are input IDs?](../glossary#input-ids)
706
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
707
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
708
+
709
+ - 1 for tokens that are **not masked**,
710
+ - 0 for tokens that are **masked**.
711
+
712
+ [What are attention masks?](../glossary#attention-mask)
713
+
714
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
715
+ [`PreTrainedTokenizer.__call__`] for details.
716
+
717
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
718
+ `past_key_values`).
719
+
720
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
721
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
722
+ information on the default strategy.
723
+
724
+ - 1 indicates the head is **not masked**,
725
+ - 0 indicates the head is **masked**.
726
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
727
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
728
+ config.n_positions - 1]`.
729
+
730
+ [What are position IDs?](../glossary#position-ids)
731
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
732
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
733
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
734
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
735
+
736
+ Two formats are allowed:
737
+ - a [`~cache_utils.Cache`] instance;
738
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
739
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
740
+ cache format.
741
+
742
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
743
+ legacy cache format will be returned.
744
+
745
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
746
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
747
+ of shape `(batch_size, sequence_length)`.
748
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
749
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
750
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
751
+ model's internal embedding lookup matrix.
752
+ use_cache (`bool`, *optional*):
753
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
754
+ `past_key_values`).
755
+ output_attentions (`bool`, *optional*):
756
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
757
+ tensors for more detail.
758
+ output_hidden_states (`bool`, *optional*):
759
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
760
+ more detail.
761
+ return_dict (`bool`, *optional*):
762
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
763
+ """
764
+
765
+
766
+ @add_start_docstrings(
767
+ "The bare StableLM Model outputting raw hidden-states without any specific head on top.",
768
+ STABLELM_START_DOCSTRING,
769
+ )
770
+ class StableLmModel(StableLmPreTrainedModel):
771
+ """
772
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`StableLmDecoderLayer`]
773
 
774
+ Args:
775
+ config: StableLmConfig
776
+ """
777
+
778
+ def __init__(self, config: StableLmConfig):
779
  super().__init__(config)
780
+ self.padding_idx = config.pad_token_id
781
+ self.vocab_size = config.vocab_size
782
+
783
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
784
+ self.layers = nn.ModuleList(
785
+ [StableLmDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
786
+ )
787
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
788
 
789
+ self._attn_implementation = config._attn_implementation
790
  self.gradient_checkpointing = False
791
  # Initialize weights and apply final processing
792
  self.post_init()
 
794
  def get_input_embeddings(self):
795
  return self.embed_tokens
796
 
797
+ def set_input_embeddings(self, value):
798
  self.embed_tokens = value
799
 
800
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
801
  def forward(
802
  self,
803
+ input_ids: torch.LongTensor = None,
804
+ attention_mask: Optional[torch.Tensor] = None,
805
  position_ids: Optional[torch.LongTensor] = None,
806
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
807
  inputs_embeds: Optional[torch.FloatTensor] = None,
808
  use_cache: Optional[bool] = None,
809
  output_attentions: Optional[bool] = None,
 
811
  return_dict: Optional[bool] = None,
812
  ) -> Union[Tuple, BaseModelOutputWithPast]:
813
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
814
+ output_hidden_states = (
815
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
816
+ )
817
  use_cache = use_cache if use_cache is not None else self.config.use_cache
818
 
819
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
820
 
821
+ # retrieve input_ids and inputs_embeds
822
  if input_ids is not None and inputs_embeds is not None:
823
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
 
 
824
  elif input_ids is not None:
825
  batch_size, seq_length = input_ids.shape
826
  elif inputs_embeds is not None:
827
  batch_size, seq_length, _ = inputs_embeds.shape
828
  else:
829
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
 
830
 
831
  seq_length_with_past = seq_length
832
  past_key_values_length = 0
833
 
834
+ if self.gradient_checkpointing and self.training:
835
+ if use_cache:
836
+ logger.warning_once(
837
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
838
+ )
839
+ use_cache = False
840
+
841
+ if use_cache:
842
+ use_legacy_cache = not isinstance(past_key_values, Cache)
843
+ if use_legacy_cache:
844
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
845
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
846
+ seq_length_with_past = seq_length_with_past + past_key_values_length
847
+
848
  if position_ids is None:
849
  device = input_ids.device if input_ids is not None else inputs_embeds.device
850
  position_ids = torch.arange(
851
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
 
 
 
852
  )
853
+ position_ids = position_ids.unsqueeze(0)
 
 
854
 
855
  if inputs_embeds is None:
856
  inputs_embeds = self.embed_tokens(input_ids)
857
+ # embed positions
858
+ if self._attn_implementation == "flash_attention_2":
859
  # 2d mask is passed through the layers
860
  attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
861
  else:
862
+ # 4d mask is passed through the layers
863
+ attention_mask = _prepare_4d_causal_attention_mask(
864
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
 
 
 
 
 
 
 
 
865
  )
866
 
867
  hidden_states = inputs_embeds
868
 
869
+ # decoder layers
 
 
 
 
 
 
 
870
  all_hidden_states = () if output_hidden_states else None
871
  all_self_attns = () if output_attentions else None
872
+ next_decoder_cache = None
873
 
874
+ for decoder_layer in self.layers:
875
  if output_hidden_states:
876
  all_hidden_states += (hidden_states,)
877
 
 
 
 
 
878
  if self.gradient_checkpointing and self.training:
879
+ layer_outputs = self._gradient_checkpointing_func(
880
+ decoder_layer.__call__,
 
 
 
 
 
 
 
 
881
  hidden_states,
882
  attention_mask,
883
  position_ids,
884
+ past_key_values,
885
+ output_attentions,
886
  )
887
  else:
888
  layer_outputs = decoder_layer(
889
  hidden_states,
890
  attention_mask=attention_mask,
891
  position_ids=position_ids,
892
+ past_key_value=past_key_values,
893
  output_attentions=output_attentions,
894
  use_cache=use_cache,
895
  )
 
897
  hidden_states = layer_outputs[0]
898
 
899
  if use_cache:
900
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
901
 
902
  if output_attentions:
903
  all_self_attns += (layer_outputs[1],)
904
 
905
  hidden_states = self.norm(hidden_states)
906
 
907
+ # add hidden states from the last decoder layer
908
  if output_hidden_states:
909
  all_hidden_states += (hidden_states,)
910
 
911
+ next_cache = None
912
+ if use_cache:
913
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
914
+
915
  if not return_dict:
916
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
 
 
 
 
917
  return BaseModelOutputWithPast(
918
  last_hidden_state=hidden_states,
919
  past_key_values=next_cache,
 
922
  )
923
 
924
 
925
+ # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonForCausalLM with PERSIMMON->STABLELM,Persimmon->StableLm
926
+ class StableLmForCausalLM(StableLmPreTrainedModel):
927
  _tied_weights_keys = ["lm_head.weight"]
928
 
929
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__ with LLAMA->STABLELM,Llama->StableLm
930
+ def __init__(self, config):
931
  super().__init__(config)
932
+ self.model = StableLmModel(config)
933
+ self.vocab_size = config.vocab_size
934
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
935
 
936
  # Initialize weights and apply final processing
937
  self.post_init()
938
 
939
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
940
  def get_input_embeddings(self):
941
  return self.model.embed_tokens
942
 
943
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
944
  def set_input_embeddings(self, value):
945
  self.model.embed_tokens = value
946
 
947
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
948
  def get_output_embeddings(self):
949
  return self.lm_head
950
 
951
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
952
+ def set_output_embeddings(self, new_embeddings):
953
  self.lm_head = new_embeddings
954
 
955
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
 
 
956
  def set_decoder(self, decoder):
957
  self.model = decoder
958
 
959
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
960
+ def get_decoder(self):
961
+ return self.model
962
+
963
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
964
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
965
+ # Ignore copy
966
  def forward(
967
  self,
968
+ input_ids: torch.LongTensor = None,
969
+ attention_mask: Optional[torch.Tensor] = None,
970
  position_ids: Optional[torch.LongTensor] = None,
971
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
972
  inputs_embeds: Optional[torch.FloatTensor] = None,
973
  labels: Optional[torch.LongTensor] = None,
974
  use_cache: Optional[bool] = None,
 
976
  output_hidden_states: Optional[bool] = None,
977
  return_dict: Optional[bool] = None,
978
  ) -> Union[Tuple, CausalLMOutputWithPast]:
979
+ r"""
980
+ Args:
981
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
982
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
983
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
984
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
985
+
986
+ Returns:
987
+
988
+ Example:
989
+
990
+ ```python
991
+ >>> from transformers import AutoTokenizer, StableLmForCausalLM
992
+
993
+ >>> model = StableLmForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t")
994
+ >>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t")
995
+
996
+ >>> prompt = "The weather is always wonderful in"
997
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
998
+
999
+ >>> # Generate
1000
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1001
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1002
+ 'The weather is always wonderful in the San Juan Islands, and whether you're a vacationer or a resident, here are some ideas for fun!'
1003
+ ```"""
1004
+
1005
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1006
  output_hidden_states = (
1007
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 
 
 
 
 
1008
  )
1009
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1010
 
 
1011
  outputs = self.model(
1012
+ input_ids=input_ids,
1013
  attention_mask=attention_mask,
1014
  position_ids=position_ids,
1015
  past_key_values=past_key_values,
 
1021
  )
1022
 
1023
  hidden_states = outputs[0]
1024
+ logits = self.lm_head(hidden_states)
1025
 
1026
  loss = None
1027
  if labels is not None:
 
1048
  attentions=outputs.attentions,
1049
  )
1050
 
1051
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.prepare_inputs_for_generation
1052
  def prepare_inputs_for_generation(
1053
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
 
 
 
 
 
1054
  ):
 
1055
  if past_key_values is not None:
1056
+ if isinstance(past_key_values, Cache):
1057
+ cache_length = past_key_values.get_seq_length()
1058
+ past_length = past_key_values.seen_tokens
1059
+ max_cache_length = past_key_values.get_max_length()
 
1060
  else:
1061
+ cache_length = past_length = past_key_values[0][0].shape[2]
1062
+ max_cache_length = None
1063
+
1064
+ # Keep only the unprocessed tokens:
1065
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1066
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1067
+ # input)
1068
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1069
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1070
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1071
+ # input_ids based on the past_length.
1072
+ elif past_length < input_ids.shape[1]:
1073
+ input_ids = input_ids[:, past_length:]
1074
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1075
+
1076
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1077
+ if (
1078
+ max_cache_length is not None
1079
+ and attention_mask is not None
1080
+ and cache_length + input_ids.shape[1] > max_cache_length
1081
+ ):
1082
+ attention_mask = attention_mask[:, -max_cache_length:]
1083
 
1084
  position_ids = kwargs.get("position_ids", None)
1085
  if attention_mask is not None and position_ids is None:
1086
+ # create position_ids on the fly for batch generation
1087
  position_ids = attention_mask.long().cumsum(-1) - 1
1088
  position_ids.masked_fill_(attention_mask == 0, 1)
1089
  if past_key_values:
1090
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1091
+
1092
+ if past_key_value := getattr(self.model.layers[0].self_attn, "past_key_value", None):
1093
+ # generation with static cache
1094
+ seen_tokens = past_key_value.get_seq_length()
1095
+ input_ids = input_ids[:, seen_tokens:]
1096
+ position_ids = position_ids[:, seen_tokens:]
1097
 
1098
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1099
  if inputs_embeds is not None and past_key_values is None:
1100
  model_inputs = {"inputs_embeds": inputs_embeds}
1101
  else:
 
1103
 
1104
  model_inputs.update(
1105
  {
1106
+ "position_ids": position_ids,
1107
  "past_key_values": past_key_values,
1108
  "use_cache": kwargs.get("use_cache"),
1109
+ "attention_mask": attention_mask,
1110
  }
1111
  )
1112
  return model_inputs
 
1116
  reordered_past = ()
1117
  for layer_past in past_key_values:
1118
  reordered_past += (
1119
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
 
 
 
1120
  )
1121
  return reordered_past
1122
 
1123
 
1124
+ @add_start_docstrings(
1125
+ """
1126
+ The StableLM transformer with a sequence classification head on top (linear layer).
1127
+
1128
+ [`StableLmForSequenceClassification`] uses the last token in order to do the classification, as other causal
1129
+ models (e.g. GPT-2) do.
1130
+
1131
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1132
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1133
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1134
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1135
+ each row of the batch).
1136
+ """,
1137
+ STABLELM_START_DOCSTRING,
1138
+ )
1139
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->STABLELM,Llama->StableLm
1140
+ class StableLmForSequenceClassification(StableLmPreTrainedModel):
1141
+ def __init__(self, config):
1142
+ super().__init__(config)
1143
+ self.num_labels = config.num_labels
1144
+ self.model = StableLmModel(config)
1145
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1146
+
1147
+ # Initialize weights and apply final processing
1148
+ self.post_init()
1149
+
1150
+ def get_input_embeddings(self):
1151
+ return self.model.embed_tokens
1152
+
1153
+ def set_input_embeddings(self, value):
1154
+ self.model.embed_tokens = value
1155
+
1156
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
1157
+ def forward(
1158
+ self,
1159
+ input_ids: torch.LongTensor = None,
1160
+ attention_mask: Optional[torch.Tensor] = None,
1161
+ position_ids: Optional[torch.LongTensor] = None,
1162
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1163
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1164
+ labels: Optional[torch.LongTensor] = None,
1165
+ use_cache: Optional[bool] = None,
1166
+ output_attentions: Optional[bool] = None,
1167
+ output_hidden_states: Optional[bool] = None,
1168
+ return_dict: Optional[bool] = None,
1169
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1170
+ r"""
1171
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1172
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1173
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1174
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1175
+ """
1176
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1177
+
1178
+ transformer_outputs = self.model(
1179
+ input_ids,
1180
+ attention_mask=attention_mask,
1181
+ position_ids=position_ids,
1182
+ past_key_values=past_key_values,
1183
+ inputs_embeds=inputs_embeds,
1184
+ use_cache=use_cache,
1185
+ output_attentions=output_attentions,
1186
+ output_hidden_states=output_hidden_states,
1187
+ return_dict=return_dict,
1188
+ )
1189
+ hidden_states = transformer_outputs[0]
1190
+ logits = self.score(hidden_states)
1191
+
1192
+ if input_ids is not None:
1193
+ batch_size = input_ids.shape[0]
1194
+ else:
1195
+ batch_size = inputs_embeds.shape[0]
1196
+
1197
+ if self.config.pad_token_id is None and batch_size != 1:
1198
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1199
+ if self.config.pad_token_id is None:
1200
+ sequence_lengths = -1
1201
+ else:
1202
+ if input_ids is not None:
1203
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1204
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1205
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1206
+ sequence_lengths = sequence_lengths.to(logits.device)
1207
+ else:
1208
+ sequence_lengths = -1
1209
+
1210
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1211
+
1212
+ loss = None
1213
+ if labels is not None:
1214
+ labels = labels.to(logits.device)
1215
+ if self.config.problem_type is None:
1216
+ if self.num_labels == 1:
1217
+ self.config.problem_type = "regression"
1218
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1219
+ self.config.problem_type = "single_label_classification"
1220
+ else:
1221
+ self.config.problem_type = "multi_label_classification"
1222
+
1223
+ if self.config.problem_type == "regression":
1224
+ loss_fct = MSELoss()
1225
+ if self.num_labels == 1:
1226
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1227
+ else:
1228
+ loss = loss_fct(pooled_logits, labels)
1229
+ elif self.config.problem_type == "single_label_classification":
1230
+ loss_fct = CrossEntropyLoss()
1231
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1232
+ elif self.config.problem_type == "multi_label_classification":
1233
+ loss_fct = BCEWithLogitsLoss()
1234
+ loss = loss_fct(pooled_logits, labels)
1235
+ if not return_dict:
1236
+ output = (pooled_logits,) + transformer_outputs[1:]
1237
+ return ((loss,) + output) if loss is not None else output
1238
+
1239
+ return SequenceClassifierOutputWithPast(
1240
+ loss=loss,
1241
+ logits=pooled_logits,
1242
+ past_key_values=transformer_outputs.past_key_values,
1243
+ hidden_states=transformer_outputs.hidden_states,
1244
+ attentions=transformer_outputs.attentions,
1245
+ )