jon-tow commited on
Commit
ebf5f38
1 Parent(s): 39d1453

merge: upload transformers implementation (#6)

Browse files

- merge: upload `transformers` implementation (efd6eabee32eabd8675eae49359287ad061fe197)
- feat: use latest modeling code (163f149452b3ecef0f4c3df5ce63c69652e45b78)

README.md CHANGED
@@ -33,7 +33,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
33
  tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-1_6b", trust_remote_code=True)
34
  model = AutoModelForCausalLM.from_pretrained(
35
  "stabilityai/stablelm-2-1_6b",
36
- trust_remote_code=True,
37
  torch_dtype="auto",
38
  )
39
  model.cuda()
@@ -58,7 +57,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
58
  tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-1_6b", trust_remote_code=True)
59
  model = AutoModelForCausalLM.from_pretrained(
60
  "stabilityai/stablelm-2-1_6b",
61
- trust_remote_code=True,
62
  torch_dtype="auto",
63
  attn_implementation="flash_attention_2",
64
  )
 
33
  tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-1_6b", trust_remote_code=True)
34
  model = AutoModelForCausalLM.from_pretrained(
35
  "stabilityai/stablelm-2-1_6b",
 
36
  torch_dtype="auto",
37
  )
38
  model.cuda()
 
57
  tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-1_6b", trust_remote_code=True)
58
  model = AutoModelForCausalLM.from_pretrained(
59
  "stabilityai/stablelm-2-1_6b",
 
60
  torch_dtype="auto",
61
  attn_implementation="flash_attention_2",
62
  )
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,32 +12,45 @@
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):
34
- Dimension of the decoder layers and the pooler layer.
35
  num_hidden_layers (`int`, *optional*, defaults to 32):
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,124 @@ 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):
48
+ Number of hidden layers in the Transformer decoder.
49
  num_hidden_layers (`int`, *optional*, defaults to 32):
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
+
138
  self.hidden_size = hidden_size
139
+ self.intermediate_size = intermediate_size
140
  self.num_hidden_layers = num_hidden_layers
141
  self.num_attention_heads = num_attention_heads
142
  self.num_key_value_heads = num_key_value_heads
143
  self.hidden_act = hidden_act
144
+
 
145
  self.initializer_range = initializer_range
146
+ self.layer_norm_eps = layer_norm_eps
147
  self.use_cache = use_cache
148
+ self.rope_theta = rope_theta
149
+ self.rope_scaling = rope_scaling
150
  self.use_qkv_bias = use_qkv_bias
151
+ self.hidden_dropout = hidden_dropout
152
  self.attention_dropout = attention_dropout
153
+ self.partial_rotary_factor = partial_rotary_factor
154
+ self._rope_scaling_validation()
155
+
156
  super().__init__(
157
  bos_token_id=bos_token_id,
158
  eos_token_id=eos_token_id,
159
  tie_word_embeddings=tie_word_embeddings,
160
  **kwargs,
161
  )
162
+
163
+ # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
164
+ def _rope_scaling_validation(self):
165
+ """
166
+ Validate the `rope_scaling` configuration.
167
+ """
168
+ if self.rope_scaling is None:
169
+ return
170
+
171
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
172
+ raise ValueError(
173
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
174
+ f"got {self.rope_scaling}"
175
+ )
176
+ rope_scaling_type = self.rope_scaling.get("type", None)
177
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
178
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
179
+ raise ValueError(
180
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
181
+ )
182
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
183
+ 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,110 @@ 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 +496,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 +513,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 +552,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 +568,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 +594,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 +628,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 +669,51 @@ 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 +733,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 +748,143 @@ 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 +892,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 +909,90 @@ 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 +1000,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 +1025,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 +1079,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 +1124,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:
@@ -858,35 +1152,46 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
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 +1199,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 +1212,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.falcon.modeling_falcon.FalconLinearScalingRotaryEmbedding with Falcon->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.falcon.modeling_falcon.FalconDynamicNTKScalingRotaryEmbedding with Falcon->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 StableLmSdpaAttention(StableLmAttention):
378
+ def forward(
379
+ self,
380
+ hidden_states: torch.Tensor,
381
+ attention_mask: Optional[torch.Tensor] = None,
382
+ position_ids: Optional[torch.LongTensor] = None,
383
+ past_key_value: Optional[Cache] = None,
384
+ output_attentions: bool = False,
385
+ use_cache: bool = False,
386
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
387
+ if output_attentions:
388
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
389
+ logger.warning_once(
390
+ "StableLmModel is using StableLmSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
391
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
392
+ )
393
+ return super().forward(
394
+ hidden_states=hidden_states,
395
+ attention_mask=attention_mask,
396
+ position_ids=position_ids,
397
+ past_key_value=past_key_value,
398
+ output_attentions=output_attentions,
399
+ use_cache=use_cache,
400
+ )
401
+
402
+ bsz, q_len, _ = hidden_states.size()
403
+
404
+ query_states = self.q_proj(hidden_states)
405
+ key_states = self.k_proj(hidden_states)
406
+ value_states = self.v_proj(hidden_states)
407
+
408
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
409
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
410
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
411
+
412
+ kv_seq_len = key_states.shape[-2]
413
+ if past_key_value is not None:
414
+ if self.layer_idx is None:
415
+ raise ValueError(
416
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
417
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
418
+ "with a layer index."
419
+ )
420
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
421
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
422
+
423
+ # Partial rotary embedding
424
+ query_rot, query_pass = (
425
+ query_states[..., : self.rotary_emb.dim],
426
+ query_states[..., self.rotary_emb.dim :],
427
+ )
428
+ key_rot, key_pass = (
429
+ key_states[..., : self.rotary_emb.dim],
430
+ key_states[..., self.rotary_emb.dim :],
431
+ )
432
+ # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
433
+ query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
434
+
435
+ # [batch_size, seq_length, num_heads, head_dim]
436
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
437
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
438
+
439
+ if past_key_value is not None:
440
+ # Specific to RoPE models with partial rotation
441
+ cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
442
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
443
+
444
+ # Repeat k/v heads if n_kv_heads < n_heads
445
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
446
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
447
+
448
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
449
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
450
+ if query_states.device.type == "cuda" and attention_mask is not None:
451
+ query_states = query_states.contiguous()
452
+ key_states = key_states.contiguous()
453
+ value_states = value_states.contiguous()
454
+
455
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
456
+ query_states,
457
+ key_states,
458
+ value_states,
459
+ attn_mask=attention_mask,
460
+ dropout_p=self.attention_dropout.p if self.training else 0.0,
461
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
462
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
463
+ )
464
+
465
+ attn_output = attn_output.transpose(1, 2).contiguous()
466
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
467
+
468
+ attn_output = self.o_proj(attn_output)
469
+
470
+ return attn_output, None, past_key_value
471
+
472
+
473
+ class StableLmFlashAttention2(StableLmAttention):
474
  """
475
+ StableLM flash attention module. This module inherits from `StableLmAttention` as the weights of the module stays
476
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
477
+ flash attention and deal with padding tokens in case the input contains any of them.
478
  """
479
 
480
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
481
  def __init__(self, *args, **kwargs):
482
  super().__init__(*args, **kwargs)
483
 
 
496
  use_cache: bool = False,
497
  **kwargs,
498
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
499
+ # StableLmFlashAttention2 attention does not support output_attentions
 
 
 
 
 
 
 
500
 
501
  output_attentions = False
502
 
 
513
  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
514
  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
515
 
 
 
 
 
 
516
  kv_seq_len = key_states.shape[-2]
517
  if past_key_value is not None:
518
+ if self.layer_idx is None:
519
+ raise ValueError(
520
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
521
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
522
+ "with a layer index."
523
+ )
524
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
525
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
 
526
 
527
+ # Partial rotary embedding
528
+ query_rot, query_pass = (
529
+ query_states[..., : self.rotary_emb.dim],
530
+ query_states[..., self.rotary_emb.dim :],
531
+ )
532
+ key_rot, key_pass = (
533
+ key_states[..., : self.rotary_emb.dim],
534
+ key_states[..., self.rotary_emb.dim :],
535
+ )
536
+ query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
537
+
538
+ # [batch_size, seq_length, num_heads, head_dim]
539
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
540
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
541
 
542
  if past_key_value is not None:
543
+ cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
544
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
 
 
545
 
546
  # 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
547
  # to be able to avoid many of these transpose/reshape/view.
 
552
  dropout_rate = self.attention_dropout if self.training else 0.0
553
 
554
  attn_output = self._flash_attention_forward(
555
+ query_states,
556
+ key_states,
557
+ value_states,
558
+ attention_mask,
559
+ q_len,
560
+ dropout=dropout_rate,
561
  )
562
+
563
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
564
  attn_output = self.o_proj(attn_output)
565
 
 
568
 
569
  return attn_output, attn_weights, past_key_value
570
 
571
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
572
  def _flash_attention_forward(
573
  self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
574
  ):
 
594
  if not self._flash_attn_uses_top_left_mask:
595
  causal = self.is_causal
596
  else:
597
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
598
  causal = self.is_causal and query_length != 1
599
 
600
  # Contains at least one padding token in the sequence
 
628
 
629
  return attn_output
630
 
631
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
632
  def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
633
  indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
634
  batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
 
669
 
670
 
671
  ATTENTION_CLASSES = {
672
+ "eager": StableLmAttention,
673
+ "sdpa": StableLmSdpaAttention,
674
+ "flash_attention_2": StableLmFlashAttention2,
675
  }
676
 
677
 
678
+ class StableLmDecoderLayer(nn.Module):
679
+ def __init__(self, config: StableLmConfig, layer_idx: int):
680
  super().__init__()
681
+ self.hidden_size = config.hidden_size
682
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
683
+ self.mlp = StableLmMLP(config)
684
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
685
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
686
+ self.dropout = nn.Dropout(config.hidden_dropout)
687
 
688
  def forward(
689
  self,
690
+ hidden_states: torch.Tensor,
691
+ attention_mask: Optional[torch.Tensor] = None,
692
  position_ids: Optional[torch.LongTensor] = None,
693
  past_key_value: Optional[Tuple[torch.Tensor]] = None,
694
  output_attentions: Optional[bool] = False,
695
  use_cache: Optional[bool] = False,
696
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
697
+ """
698
+ Args:
699
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
700
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
701
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
702
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
703
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
704
+ `[0, config.n_positions - 1]`.
705
+
706
+ [What are position IDs?](../glossary#position-ids)
707
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
708
+ cached past key and value projection states
709
+ output_attentions (`bool`, *optional*):
710
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
711
+ returned tensors for more detail.
712
+ use_cache (`bool`, *optional*):
713
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
714
+ (see `past_key_values`).
715
+ """
716
+
717
  residual = hidden_states
718
 
719
  hidden_states = self.input_layernorm(hidden_states)
 
733
  residual = hidden_states
734
  hidden_states = self.post_attention_layernorm(hidden_states)
735
  hidden_states = self.mlp(hidden_states)
736
+
737
+ hidden_states = self.dropout(hidden_states)
738
+ hidden_states = hidden_states + residual
739
 
740
  outputs = (hidden_states,)
741
 
 
748
  return outputs
749
 
750
 
751
+ STABLELM_START_DOCSTRING = r"""
752
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
753
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
754
+ etc.)
755
+
756
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
757
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
758
+ and behavior.
759
+
760
+ Parameters:
761
+ config ([`StableLmConfig`]):
762
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
763
+ load the weights associated with the model, only the configuration. Check out the
764
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
765
+ """
766
 
767
+
768
+ @add_start_docstrings(
769
+ "The bare StableLm Model outputting raw hidden-states without any specific head on top.",
770
+ STABLELM_START_DOCSTRING,
771
+ )
772
+ class StableLmPreTrainedModel(PreTrainedModel):
773
+ config_class = StableLmConfig
774
  base_model_prefix = "model"
775
  supports_gradient_checkpointing = True
776
+ _no_split_modules = ["StableLmDecoderLayer"]
777
  _skip_keys_device_placement = "past_key_values"
778
  _supports_flash_attn_2 = True
779
+ _supports_cache_class = True
780
+ _supports_sdpa = True
781
 
782
+ def _init_weights(self, module):
783
+ std = self.config.initializer_range
784
  if isinstance(module, nn.Linear):
785
+ module.weight.data.normal_(mean=0.0, std=std)
786
  if module.bias is not None:
787
  module.bias.data.zero_()
788
  elif isinstance(module, nn.Embedding):
789
+ module.weight.data.normal_(mean=0.0, std=std)
790
  if module.padding_idx is not None:
791
  module.weight.data[module.padding_idx].zero_()
 
 
 
792
 
 
 
 
793
 
794
+ STABLELM_INPUTS_DOCSTRING = r"""
795
+ Args:
796
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
797
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
798
+ it.
799
+
800
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
801
+ [`PreTrainedTokenizer.__call__`] for details.
802
+
803
+ [What are input IDs?](../glossary#input-ids)
804
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
805
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
806
+
807
+ - 1 for tokens that are **not masked**,
808
+ - 0 for tokens that are **masked**.
809
+
810
+ [What are attention masks?](../glossary#attention-mask)
811
+
812
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
813
+ [`PreTrainedTokenizer.__call__`] for details.
814
+
815
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
816
+ `past_key_values`).
817
+
818
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
819
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
820
+ information on the default strategy.
821
+
822
+ - 1 indicates the head is **not masked**,
823
+ - 0 indicates the head is **masked**.
824
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
825
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
826
+ config.n_positions - 1]`.
827
+
828
+ [What are position IDs?](../glossary#position-ids)
829
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
830
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
831
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
832
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
833
+
834
+ Two formats are allowed:
835
+ - a [`~cache_utils.Cache`] instance;
836
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
837
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
838
+ cache format.
839
+
840
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
841
+ legacy cache format will be returned.
842
+
843
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
844
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
845
+ of shape `(batch_size, sequence_length)`.
846
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
847
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
848
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
849
+ model's internal embedding lookup matrix.
850
+ use_cache (`bool`, *optional*):
851
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
852
+ `past_key_values`).
853
+ output_attentions (`bool`, *optional*):
854
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
855
+ tensors for more detail.
856
+ output_hidden_states (`bool`, *optional*):
857
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
858
+ more detail.
859
+ return_dict (`bool`, *optional*):
860
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
861
+ """
862
+
863
+
864
+ @add_start_docstrings(
865
+ "The bare StableLm Model outputting raw hidden-states without any specific head on top.",
866
+ STABLELM_START_DOCSTRING,
867
+ )
868
+ class StableLmModel(StableLmPreTrainedModel):
869
+ """
870
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`StableLmDecoderLayer`]
871
+
872
+ Args:
873
+ config: StableLmConfig
874
+ """
875
 
876
+ def __init__(self, config: StableLmConfig):
 
877
  super().__init__(config)
878
+ self.padding_idx = config.pad_token_id
879
+ self.vocab_size = config.vocab_size
880
+
881
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
882
+ self.layers = nn.ModuleList(
883
+ [StableLmDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
884
+ )
885
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
886
 
887
+ self._attn_implementation = config._attn_implementation
888
  self.gradient_checkpointing = False
889
  # Initialize weights and apply final processing
890
  self.post_init()
 
892
  def get_input_embeddings(self):
893
  return self.embed_tokens
894
 
895
+ def set_input_embeddings(self, value):
896
  self.embed_tokens = value
897
 
898
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
899
  def forward(
900
  self,
901
+ input_ids: torch.LongTensor = None,
902
+ attention_mask: Optional[torch.Tensor] = None,
903
  position_ids: Optional[torch.LongTensor] = None,
904
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
905
  inputs_embeds: Optional[torch.FloatTensor] = None,
906
  use_cache: Optional[bool] = None,
907
  output_attentions: Optional[bool] = None,
 
909
  return_dict: Optional[bool] = None,
910
  ) -> Union[Tuple, BaseModelOutputWithPast]:
911
  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
912
+ output_hidden_states = (
913
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
914
+ )
915
  use_cache = use_cache if use_cache is not None else self.config.use_cache
916
 
917
  return_dict = return_dict if return_dict is not None else self.config.use_return_dict
918
 
919
+ # retrieve input_ids and inputs_embeds
920
  if input_ids is not None and inputs_embeds is not None:
921
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
 
 
922
  elif input_ids is not None:
923
  batch_size, seq_length = input_ids.shape
924
  elif inputs_embeds is not None:
925
  batch_size, seq_length, _ = inputs_embeds.shape
926
  else:
927
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
 
928
 
929
  seq_length_with_past = seq_length
930
  past_key_values_length = 0
931
 
932
+ if self.gradient_checkpointing and self.training:
933
+ if use_cache:
934
+ logger.warning_once(
935
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
936
+ )
937
+ use_cache = False
938
+
939
+ if use_cache:
940
+ use_legacy_cache = not isinstance(past_key_values, Cache)
941
+ if use_legacy_cache:
942
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
943
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
944
+ seq_length_with_past = seq_length_with_past + past_key_values_length
945
+
946
  if position_ids is None:
947
  device = input_ids.device if input_ids is not None else inputs_embeds.device
948
  position_ids = torch.arange(
949
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
 
 
 
950
  )
951
+ position_ids = position_ids.unsqueeze(0)
 
 
952
 
953
  if inputs_embeds is None:
954
  inputs_embeds = self.embed_tokens(input_ids)
955
+ # embed positions
956
+ if self._attn_implementation == "flash_attention_2":
957
  # 2d mask is passed through the layers
958
  attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
959
+ # for output_attentions case used fallback to eager attention realization
960
+ elif self._attn_implementation == "sdpa" and not output_attentions:
961
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
962
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
963
+ )
964
  else:
965
+ # 4d mask is passed through the layers
966
+ attention_mask = _prepare_4d_causal_attention_mask(
967
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
 
 
 
 
 
 
 
 
968
  )
969
 
970
  hidden_states = inputs_embeds
971
 
972
+ # decoder layers
 
 
 
 
 
 
 
973
  all_hidden_states = () if output_hidden_states else None
974
  all_self_attns = () if output_attentions else None
975
+ next_decoder_cache = None
976
 
977
+ for decoder_layer in self.layers:
978
  if output_hidden_states:
979
  all_hidden_states += (hidden_states,)
980
 
 
 
 
 
981
  if self.gradient_checkpointing and self.training:
982
+ layer_outputs = self._gradient_checkpointing_func(
983
+ decoder_layer.__call__,
 
 
 
 
 
 
 
 
984
  hidden_states,
985
  attention_mask,
986
  position_ids,
987
+ past_key_values,
988
+ output_attentions,
989
  )
990
  else:
991
  layer_outputs = decoder_layer(
992
  hidden_states,
993
  attention_mask=attention_mask,
994
  position_ids=position_ids,
995
+ past_key_value=past_key_values,
996
  output_attentions=output_attentions,
997
  use_cache=use_cache,
998
  )
 
1000
  hidden_states = layer_outputs[0]
1001
 
1002
  if use_cache:
1003
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1004
 
1005
  if output_attentions:
1006
  all_self_attns += (layer_outputs[1],)
1007
 
1008
  hidden_states = self.norm(hidden_states)
1009
 
1010
+ # add hidden states from the last decoder layer
1011
  if output_hidden_states:
1012
  all_hidden_states += (hidden_states,)
1013
 
1014
+ next_cache = None
1015
+ if use_cache:
1016
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1017
+
1018
  if not return_dict:
1019
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
 
 
 
 
1020
  return BaseModelOutputWithPast(
1021
  last_hidden_state=hidden_states,
1022
  past_key_values=next_cache,
 
1025
  )
1026
 
1027
 
1028
+ # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonForCausalLM with PERSIMMON->STABLELM,Persimmon->StableLm
1029
+ class StableLmForCausalLM(StableLmPreTrainedModel):
1030
  _tied_weights_keys = ["lm_head.weight"]
1031
 
1032
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__ with LLAMA->STABLELM,Llama->StableLm
1033
+ def __init__(self, config):
1034
  super().__init__(config)
1035
+ self.model = StableLmModel(config)
1036
+ self.vocab_size = config.vocab_size
1037
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1038
 
1039
  # Initialize weights and apply final processing
1040
  self.post_init()
1041
 
1042
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
1043
  def get_input_embeddings(self):
1044
  return self.model.embed_tokens
1045
 
1046
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
1047
  def set_input_embeddings(self, value):
1048
  self.model.embed_tokens = value
1049
 
1050
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
1051
  def get_output_embeddings(self):
1052
  return self.lm_head
1053
 
1054
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
1055
+ def set_output_embeddings(self, new_embeddings):
1056
  self.lm_head = new_embeddings
1057
 
1058
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
 
 
1059
  def set_decoder(self, decoder):
1060
  self.model = decoder
1061
 
1062
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
1063
+ def get_decoder(self):
1064
+ return self.model
1065
+
1066
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
1067
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1068
+ # Ignore copy
1069
  def forward(
1070
  self,
1071
+ input_ids: torch.LongTensor = None,
1072
+ attention_mask: Optional[torch.Tensor] = None,
1073
  position_ids: Optional[torch.LongTensor] = None,
1074
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1075
  inputs_embeds: Optional[torch.FloatTensor] = None,
1076
  labels: Optional[torch.LongTensor] = None,
1077
  use_cache: Optional[bool] = None,
 
1079
  output_hidden_states: Optional[bool] = None,
1080
  return_dict: Optional[bool] = None,
1081
  ) -> Union[Tuple, CausalLMOutputWithPast]:
1082
+ r"""
1083
+ Args:
1084
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1085
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1086
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1087
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1088
+
1089
+ Returns:
1090
+
1091
+ Example:
1092
+
1093
+ ```python
1094
+ >>> from transformers import AutoTokenizer, StableLmForCausalLM
1095
+
1096
+ >>> model = StableLmForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t")
1097
+ >>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t")
1098
+
1099
+ >>> prompt = "The weather is always wonderful in"
1100
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1101
+
1102
+ >>> # Generate
1103
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1104
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1105
+ 'The weather is always wonderful in the summer in the city of San Diego. The city is located on the coast of the Pacific Ocean and is surrounded by'
1106
+ ```"""
1107
+
1108
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1109
  output_hidden_states = (
1110
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 
 
 
 
 
1111
  )
1112
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1113
 
 
1114
  outputs = self.model(
1115
+ input_ids=input_ids,
1116
  attention_mask=attention_mask,
1117
  position_ids=position_ids,
1118
  past_key_values=past_key_values,
 
1124
  )
1125
 
1126
  hidden_states = outputs[0]
1127
+ logits = self.lm_head(hidden_states)
1128
 
1129
  loss = None
1130
  if labels is not None:
 
1152
  )
1153
 
1154
  def prepare_inputs_for_generation(
1155
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
 
 
 
 
 
1156
  ):
 
1157
  if past_key_values is not None:
1158
+ if isinstance(past_key_values, Cache):
1159
+ cache_length = past_key_values.get_seq_length()
1160
+ past_length = past_key_values.seen_tokens
1161
+ max_cache_length = past_key_values.get_max_length()
 
1162
  else:
1163
+ cache_length = past_length = past_key_values[0][0].shape[2]
1164
+ max_cache_length = None
1165
+
1166
+ # Keep only the unprocessed tokens:
1167
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1168
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1169
+ # input)
1170
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1171
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1172
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1173
+ # input_ids based on the past_length.
1174
+ elif past_length < input_ids.shape[1]:
1175
+ input_ids = input_ids[:, past_length:]
1176
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1177
+
1178
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1179
+ if (
1180
+ max_cache_length is not None
1181
+ and attention_mask is not None
1182
+ and cache_length + input_ids.shape[1] > max_cache_length
1183
+ ):
1184
+ attention_mask = attention_mask[:, -max_cache_length:]
1185
 
1186
  position_ids = kwargs.get("position_ids", None)
1187
  if attention_mask is not None and position_ids is None:
1188
+ # create position_ids on the fly for batch generation
1189
  position_ids = attention_mask.long().cumsum(-1) - 1
1190
  position_ids.masked_fill_(attention_mask == 0, 1)
1191
  if past_key_values:
1192
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1193
 
1194
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1195
  if inputs_embeds is not None and past_key_values is None:
1196
  model_inputs = {"inputs_embeds": inputs_embeds}
1197
  else:
 
1199
 
1200
  model_inputs.update(
1201
  {
1202
+ "position_ids": position_ids,
1203
  "past_key_values": past_key_values,
1204
  "use_cache": kwargs.get("use_cache"),
1205
+ "attention_mask": attention_mask,
1206
  }
1207
  )
1208
  return model_inputs
 
1212
  reordered_past = ()
1213
  for layer_past in past_key_values:
1214
  reordered_past += (
1215
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
 
 
 
1216
  )
1217
  return reordered_past
1218
 
1219
 
1220
+ @add_start_docstrings(
1221
+ """
1222
+ The StableLm transformer with a sequence classification head on top (linear layer).
1223
+
1224
+ [`StableLmForSequenceClassification`] uses the last token in order to do the classification, as other causal
1225
+ models (e.g. GPT-2) do.
1226
+
1227
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1228
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1229
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1230
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1231
+ each row of the batch).
1232
+ """,
1233
+ STABLELM_START_DOCSTRING,
1234
+ )
1235
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->STABLELM,Llama->StableLm
1236
+ class StableLmForSequenceClassification(StableLmPreTrainedModel):
1237
+ def __init__(self, config):
1238
+ super().__init__(config)
1239
+ self.num_labels = config.num_labels
1240
+ self.model = StableLmModel(config)
1241
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1242
+
1243
+ # Initialize weights and apply final processing
1244
+ self.post_init()
1245
+
1246
+ def get_input_embeddings(self):
1247
+ return self.model.embed_tokens
1248
+
1249
+ def set_input_embeddings(self, value):
1250
+ self.model.embed_tokens = value
1251
+
1252
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
1253
+ def forward(
1254
+ self,
1255
+ input_ids: torch.LongTensor = None,
1256
+ attention_mask: Optional[torch.Tensor] = None,
1257
+ position_ids: Optional[torch.LongTensor] = None,
1258
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1259
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1260
+ labels: Optional[torch.LongTensor] = None,
1261
+ use_cache: Optional[bool] = None,
1262
+ output_attentions: Optional[bool] = None,
1263
+ output_hidden_states: Optional[bool] = None,
1264
+ return_dict: Optional[bool] = None,
1265
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1266
+ r"""
1267
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1268
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1269
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1270
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1271
+ """
1272
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1273
+
1274
+ transformer_outputs = self.model(
1275
+ input_ids,
1276
+ attention_mask=attention_mask,
1277
+ position_ids=position_ids,
1278
+ past_key_values=past_key_values,
1279
+ inputs_embeds=inputs_embeds,
1280
+ use_cache=use_cache,
1281
+ output_attentions=output_attentions,
1282
+ output_hidden_states=output_hidden_states,
1283
+ return_dict=return_dict,
1284
+ )
1285
+ hidden_states = transformer_outputs[0]
1286
+ logits = self.score(hidden_states)
1287
+
1288
+ if input_ids is not None:
1289
+ batch_size = input_ids.shape[0]
1290
+ else:
1291
+ batch_size = inputs_embeds.shape[0]
1292
+
1293
+ if self.config.pad_token_id is None and batch_size != 1:
1294
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1295
+ if self.config.pad_token_id is None:
1296
+ sequence_lengths = -1
1297
+ else:
1298
+ if input_ids is not None:
1299
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1300
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1301
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1302
+ sequence_lengths = sequence_lengths.to(logits.device)
1303
+ else:
1304
+ sequence_lengths = -1
1305
+
1306
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1307
+
1308
+ loss = None
1309
+ if labels is not None:
1310
+ labels = labels.to(logits.device)
1311
+ if self.config.problem_type is None:
1312
+ if self.num_labels == 1:
1313
+ self.config.problem_type = "regression"
1314
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1315
+ self.config.problem_type = "single_label_classification"
1316
+ else:
1317
+ self.config.problem_type = "multi_label_classification"
1318
+
1319
+ if self.config.problem_type == "regression":
1320
+ loss_fct = MSELoss()
1321
+ if self.num_labels == 1:
1322
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1323
+ else:
1324
+ loss = loss_fct(pooled_logits, labels)
1325
+ elif self.config.problem_type == "single_label_classification":
1326
+ loss_fct = CrossEntropyLoss()
1327
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1328
+ elif self.config.problem_type == "multi_label_classification":
1329
+ loss_fct = BCEWithLogitsLoss()
1330
+ loss = loss_fct(pooled_logits, labels)
1331
+ if not return_dict:
1332
+ output = (pooled_logits,) + transformer_outputs[1:]
1333
+ return ((loss,) + output) if loss is not None else output
1334
+
1335
+ return SequenceClassifierOutputWithPast(
1336
+ loss=loss,
1337
+ logits=pooled_logits,
1338
+ past_key_values=transformer_outputs.past_key_values,
1339
+ hidden_states=transformer_outputs.hidden_states,
1340
+ attentions=transformer_outputs.attentions,
1341
+ )