aloobun commited on
Commit
c36ca75
1 Parent(s): 96039b2

Upload StableLMEpochForCausalLM

Browse files
README.md ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ tags: []
4
+ ---
5
+
6
+ # Model Card for Model ID
7
+
8
+ <!-- Provide a quick summary of what the model is/does. -->
9
+
10
+
11
+
12
+ ## Model Details
13
+
14
+ ### Model Description
15
+
16
+ <!-- Provide a longer summary of what this model is. -->
17
+
18
+ This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
19
+
20
+ - **Developed by:** [More Information Needed]
21
+ - **Funded by [optional]:** [More Information Needed]
22
+ - **Shared by [optional]:** [More Information Needed]
23
+ - **Model type:** [More Information Needed]
24
+ - **Language(s) (NLP):** [More Information Needed]
25
+ - **License:** [More Information Needed]
26
+ - **Finetuned from model [optional]:** [More Information Needed]
27
+
28
+ ### Model Sources [optional]
29
+
30
+ <!-- Provide the basic links for the model. -->
31
+
32
+ - **Repository:** [More Information Needed]
33
+ - **Paper [optional]:** [More Information Needed]
34
+ - **Demo [optional]:** [More Information Needed]
35
+
36
+ ## Uses
37
+
38
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
39
+
40
+ ### Direct Use
41
+
42
+ <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
43
+
44
+ [More Information Needed]
45
+
46
+ ### Downstream Use [optional]
47
+
48
+ <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
49
+
50
+ [More Information Needed]
51
+
52
+ ### Out-of-Scope Use
53
+
54
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
55
+
56
+ [More Information Needed]
57
+
58
+ ## Bias, Risks, and Limitations
59
+
60
+ <!-- This section is meant to convey both technical and sociotechnical limitations. -->
61
+
62
+ [More Information Needed]
63
+
64
+ ### Recommendations
65
+
66
+ <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
67
+
68
+ Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
69
+
70
+ ## How to Get Started with the Model
71
+
72
+ Use the code below to get started with the model.
73
+
74
+ [More Information Needed]
75
+
76
+ ## Training Details
77
+
78
+ ### Training Data
79
+
80
+ <!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
81
+
82
+ [More Information Needed]
83
+
84
+ ### Training Procedure
85
+
86
+ <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
87
+
88
+ #### Preprocessing [optional]
89
+
90
+ [More Information Needed]
91
+
92
+
93
+ #### Training Hyperparameters
94
+
95
+ - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
96
+
97
+ #### Speeds, Sizes, Times [optional]
98
+
99
+ <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
100
+
101
+ [More Information Needed]
102
+
103
+ ## Evaluation
104
+
105
+ <!-- This section describes the evaluation protocols and provides the results. -->
106
+
107
+ ### Testing Data, Factors & Metrics
108
+
109
+ #### Testing Data
110
+
111
+ <!-- This should link to a Dataset Card if possible. -->
112
+
113
+ [More Information Needed]
114
+
115
+ #### Factors
116
+
117
+ <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
118
+
119
+ [More Information Needed]
120
+
121
+ #### Metrics
122
+
123
+ <!-- These are the evaluation metrics being used, ideally with a description of why. -->
124
+
125
+ [More Information Needed]
126
+
127
+ ### Results
128
+
129
+ [More Information Needed]
130
+
131
+ #### Summary
132
+
133
+
134
+
135
+ ## Model Examination [optional]
136
+
137
+ <!-- Relevant interpretability work for the model goes here -->
138
+
139
+ [More Information Needed]
140
+
141
+ ## Environmental Impact
142
+
143
+ <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
144
+
145
+ Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
146
+
147
+ - **Hardware Type:** [More Information Needed]
148
+ - **Hours used:** [More Information Needed]
149
+ - **Cloud Provider:** [More Information Needed]
150
+ - **Compute Region:** [More Information Needed]
151
+ - **Carbon Emitted:** [More Information Needed]
152
+
153
+ ## Technical Specifications [optional]
154
+
155
+ ### Model Architecture and Objective
156
+
157
+ [More Information Needed]
158
+
159
+ ### Compute Infrastructure
160
+
161
+ [More Information Needed]
162
+
163
+ #### Hardware
164
+
165
+ [More Information Needed]
166
+
167
+ #### Software
168
+
169
+ [More Information Needed]
170
+
171
+ ## Citation [optional]
172
+
173
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
174
+
175
+ **BibTeX:**
176
+
177
+ [More Information Needed]
178
+
179
+ **APA:**
180
+
181
+ [More Information Needed]
182
+
183
+ ## Glossary [optional]
184
+
185
+ <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
186
+
187
+ [More Information Needed]
188
+
189
+ ## More Information [optional]
190
+
191
+ [More Information Needed]
192
+
193
+ ## Model Card Authors [optional]
194
+
195
+ [More Information Needed]
196
+
197
+ ## Model Card Contact
198
+
199
+ [More Information Needed]
200
+
201
+
config.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/content/out-9deb715f-949f-46fd-b9fa-9080bed03e92/checkpoint-8190",
3
+ "architectures": [
4
+ "StableLMEpochForCausalLM"
5
+ ],
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_stablelm_epoch.StableLMEpochConfig",
9
+ "AutoModelForCausalLM": "modeling_stablelm_epoch.StableLMEpochForCausalLM"
10
+ },
11
+ "bos_token_id": 100257,
12
+ "eos_token_id": 100278,
13
+ "hidden_act": "silu",
14
+ "hidden_size": 2048,
15
+ "initializer_range": 0.02,
16
+ "intermediate_size": 5632,
17
+ "max_position_embeddings": 4096,
18
+ "model_type": "stablelm_epoch",
19
+ "norm_eps": 1e-05,
20
+ "num_attention_heads": 32,
21
+ "num_heads": 32,
22
+ "num_hidden_layers": 24,
23
+ "num_key_value_heads": 32,
24
+ "rope_pct": 0.25,
25
+ "rope_theta": 10000,
26
+ "rotary_scaling_factor": 1.0,
27
+ "tie_word_embeddings": false,
28
+ "torch_dtype": "bfloat16",
29
+ "transformers_version": "4.37.1",
30
+ "use_cache": true,
31
+ "use_qkv_bias": true,
32
+ "vocab_size": 100352
33
+ }
configuration_stablelm_epoch.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
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
43
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
44
+ by meanpooling all the original heads within that group. For more details checkout [this
45
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
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
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 100257,
4
+ "eos_token_id": 100257,
5
+ "transformers_version": "4.37.1"
6
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1ab21104e173ff2c5c03ca4914b78ea6167dfff40182e8412334d1ff228f8670
3
+ size 3289069520
modeling_stablelm_epoch.py ADDED
@@ -0,0 +1,919 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
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,
60
+ max_seqlen_in_batch,
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,
174
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
175
+ """
176
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
177
+ if n_rep == 1:
178
+ return hidden_states
179
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
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
+
228
+ query_states = self.q_proj(hidden_states)
229
+ key_states = self.k_proj(hidden_states)
230
+ value_states = self.v_proj(hidden_states)
231
+
232
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
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)
260
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
261
+
262
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
263
+
264
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
265
+ raise ValueError(
266
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
267
+ f" {attn_weights.size()}"
268
+ )
269
+
270
+ if attention_mask is not None:
271
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
272
+ raise ValueError(
273
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
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):
283
+ raise ValueError(
284
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
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:
296
+ attn_weights = None
297
+
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
+
309
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
310
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
311
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
312
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
313
+
314
+ def forward(
315
+ self,
316
+ hidden_states: torch.Tensor,
317
+ attention_mask: Optional[torch.LongTensor] = None,
318
+ position_ids: Optional[torch.LongTensor] = None,
319
+ past_key_value: Optional[Cache] = None,
320
+ output_attentions: bool = False,
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
+
335
+ bsz, q_len, _ = hidden_states.size()
336
+
337
+ query_states = self.q_proj(hidden_states)
338
+ key_states = self.k_proj(hidden_states)
339
+ value_states = self.v_proj(hidden_states)
340
+
341
+ # Flash attention requires the input to have the shape
342
+ # batch_size x seq_length x head_dim x hidden_dim
343
+ # therefore we just need to keep the original shape
344
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
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.
372
+ query_states = query_states.transpose(1, 2)
373
+ key_states = key_states.transpose(1, 2)
374
+ value_states = value_states.transpose(1, 2)
375
+
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
+
384
+ if not output_attentions:
385
+ attn_weights = None
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
+ ):
392
+ """
393
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
394
+ first unpad the input, then computes the attention scores and pad the final attention scores.
395
+
396
+ Args:
397
+ query_states (`torch.Tensor`):
398
+ Input query states to be passed to Flash Attention API
399
+ key_states (`torch.Tensor`):
400
+ Input key states to be passed to Flash Attention API
401
+ value_states (`torch.Tensor`):
402
+ Input value states to be passed to Flash Attention API
403
+ attention_mask (`torch.Tensor`):
404
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
405
+ position of padding tokens and 1 for the position of non-padding tokens.
406
+ dropout (`int`, *optional*):
407
+ Attention dropout
408
+ softmax_scale (`float`, *optional*):
409
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
410
+ """
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
418
+ if attention_mask is not None:
419
+ batch_size = query_states.shape[0]
420
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
421
+ query_states, key_states, value_states, attention_mask, query_length
422
+ )
423
+
424
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
425
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
426
+
427
+ attn_output_unpad = flash_attn_varlen_func(
428
+ query_states,
429
+ key_states,
430
+ value_states,
431
+ cu_seqlens_q=cu_seqlens_q,
432
+ cu_seqlens_k=cu_seqlens_k,
433
+ max_seqlen_q=max_seqlen_in_batch_q,
434
+ max_seqlen_k=max_seqlen_in_batch_k,
435
+ dropout_p=dropout,
436
+ softmax_scale=softmax_scale,
437
+ causal=causal,
438
+ )
439
+
440
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
441
+ else:
442
+ attn_output = flash_attn_func(
443
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
444
+ )
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
451
+
452
+ key_layer = index_first_axis(
453
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
454
+ )
455
+ value_layer = index_first_axis(
456
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
457
+ )
458
+ if query_length == kv_seq_len:
459
+ query_layer = index_first_axis(
460
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
461
+ )
462
+ cu_seqlens_q = cu_seqlens_k
463
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
464
+ indices_q = indices_k
465
+ elif query_length == 1:
466
+ max_seqlen_in_batch_q = 1
467
+ cu_seqlens_q = torch.arange(
468
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
469
+ ) # There is a memcpy here, that is very bad.
470
+ indices_q = cu_seqlens_q[:-1]
471
+ query_layer = query_layer.squeeze(1)
472
+ else:
473
+ # The -q_len: slice assumes left padding.
474
+ attention_mask = attention_mask[:, -query_length:]
475
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
476
+
477
+ return (
478
+ query_layer,
479
+ key_layer,
480
+ value_layer,
481
+ indices_q,
482
+ (cu_seqlens_q, cu_seqlens_k),
483
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
484
+ )
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)
513
+
514
+ # Self Attention
515
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
516
+ hidden_states=hidden_states,
517
+ attention_mask=attention_mask,
518
+ position_ids=position_ids,
519
+ past_key_value=past_key_value,
520
+ output_attentions=output_attentions,
521
+ use_cache=use_cache,
522
+ )
523
+ hidden_states = residual + hidden_states
524
+
525
+ # Fully Connected
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
+
533
+ if output_attentions:
534
+ outputs += (self_attn_weights,)
535
+
536
+ if use_cache:
537
+ outputs += (present_key_value,)
538
+
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 = "transformer"
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()
584
+
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,
628
+ output_hidden_states: Optional[bool] = None,
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
+ )
732
+
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,
757
+ hidden_states=all_hidden_states,
758
+ attentions=all_self_attns,
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,
801
+ output_attentions: Optional[bool] = None,
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,
825
+ inputs_embeds=inputs_embeds,
826
+ use_cache=use_cache,
827
+ output_attentions=output_attentions,
828
+ output_hidden_states=output_hidden_states,
829
+ return_dict=return_dict,
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:
837
+ # Shift so that tokens < n predict n
838
+ shift_logits = logits[..., :-1, :].contiguous()
839
+ shift_labels = labels[..., 1:].contiguous()
840
+ # Flatten the tokens
841
+ loss_fct = CrossEntropyLoss()
842
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
843
+ shift_labels = shift_labels.view(-1)
844
+ # Enable model parallelism
845
+ shift_labels = shift_labels.to(shift_logits.device)
846
+ loss = loss_fct(shift_logits, shift_labels)
847
+
848
+ if not return_dict:
849
+ output = (logits,) + outputs[1:]
850
+ return (loss,) + output if loss is not None else output
851
+
852
+ return CausalLMOutputWithPast(
853
+ loss=loss,
854
+ logits=logits,
855
+ past_key_values=outputs.past_key_values,
856
+ hidden_states=outputs.hidden_states,
857
+ attentions=outputs.attentions,
858
+ )
859
+
860
+ def prepare_inputs_for_generation(
861
+ self,
862
+ input_ids,
863
+ past_key_values: Optional[torch.Tensor] = None,
864
+ attention_mask: Optional[torch.Tensor] = None,
865
+ inputs_embeds: Optional[torch.Tensor] = None,
866
+ **kwargs,
867
+ ):
868
+ # Trim decoder_input_ids if past is used
869
+ if past_key_values is not None:
870
+ past_length = past_key_values[0][0].shape[2]
871
+
872
+ # Some generation methods already pass only the last input ID
873
+ if input_ids.shape[1] > past_length:
874
+ remove_prefix_length = past_length
875
+ else:
876
+ # Default to old behavior: keep only final ID
877
+ remove_prefix_length = input_ids.shape[1] - 1
878
+
879
+ input_ids = input_ids[:, remove_prefix_length:]
880
+
881
+ position_ids = kwargs.get("position_ids", None)
882
+ if attention_mask is not None and position_ids is None:
883
+ # Create position_ids on the fly for batch generation
884
+ position_ids = attention_mask.long().cumsum(-1) - 1
885
+ position_ids.masked_fill_(attention_mask == 0, 1)
886
+ if past_key_values:
887
+ position_ids = position_ids[:, -1].unsqueeze(-1)
888
+
889
+ # If `inputs_embeds` are passed, we only want to use them in the 1st generation step
890
+ if inputs_embeds is not None and past_key_values is None:
891
+ model_inputs = {"inputs_embeds": inputs_embeds}
892
+ else:
893
+ model_inputs = {"input_ids": input_ids}
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
904
+
905
+ @staticmethod
906
+ def _reorder_cache(past_key_values, beam_idx):
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")