portilla911 commited on
Commit
84b6988
1 Parent(s): 089323b

Upload StableLMEpochForCausalLM

Browse files
config.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "stabilityai/stable-code-3b",
3
+ "architectures": [
4
+ "StableLMEpochForCausalLM"
5
+ ],
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_stablelm_epoch.StableLMEpochConfig",
8
+ "AutoModelForCausalLM": "modeling_stablelm_epoch.StableLMEpochForCausalLM"
9
+ },
10
+ "bos_token_id": 0,
11
+ "eos_token_id": 0,
12
+ "hidden_act": "silu",
13
+ "hidden_size": 2560,
14
+ "initializer_range": 0.02,
15
+ "intermediate_size": 6912,
16
+ "max_position_embeddings": 16384,
17
+ "model_type": "stablelm_epoch",
18
+ "norm_eps": 1e-05,
19
+ "num_attention_heads": 32,
20
+ "num_heads": 32,
21
+ "num_hidden_layers": 32,
22
+ "num_key_value_heads": 32,
23
+ "rope_pct": 0.25,
24
+ "rope_theta": 1000000,
25
+ "rotary_scaling_factor": 1.0,
26
+ "tie_word_embeddings": false,
27
+ "torch_dtype": "float32",
28
+ "transformers_version": "4.36.2",
29
+ "use_cache": true,
30
+ "vocab_size": 50304
31
+ }
configuration_stablelm_epoch.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Stability 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
+ """ StableLM Epoch model configuration"""
16
+ from transformers import PretrainedConfig
17
+ from transformers.utils import logging
18
+
19
+
20
+ logger = logging.get_logger(__name__)
21
+
22
+
23
+ class StableLMEpochConfig(PretrainedConfig):
24
+ r"""
25
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
26
+ documentation from [`PretrainedConfig`] for more information.
27
+
28
+ Args:
29
+ vocab_size (`int`, *optional*, defaults to 50_304):
30
+ Vocabulary size of the StableLM model. Defines the number of different tokens that
31
+ can be represented by the `inputs_ids` passed when calling [`StableLMEpochModel`].
32
+ intermediate_size (`int`, *optional*, defaults to 6912):
33
+ Dimension of the MLP representations.
34
+ hidden_size (`int`, *optional*, defaults to 2560):
35
+ Dimension of the decoder layers and the pooler layer.
36
+ num_hidden_layers (`int`, *optional*, defaults to 32):
37
+ Number of hidden layers in the Transformer decoder.
38
+ num_attention_heads (`int`, *optional*, defaults to 32):
39
+ Number of attention heads for each attention layer in the Transformer encoder.
40
+ num_key_value_heads (`int`, *optional*):
41
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
42
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
43
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
44
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
45
+ by meanpooling all the original heads within that group. For more details checkout [this
46
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
47
+ `num_attention_heads`.
48
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
49
+ The non-linear activation function (function or string).
50
+ rope_pct (`float`, *optional*, defaults to 1.0):
51
+ Percentage of hidden dimensions to allocate to rotary embeddings.
52
+ rope_theta (`float`, *optional*, defaults to 10000.0):
53
+ The base period of the RoPE embeddings.
54
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
55
+ The maximum sequence length that this model might ever be used with.
56
+ Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
57
+ initializer_range (`float`, *optional*, defaults to 1e-5):
58
+ The standard deviation of the truncated_normal_initializer for initializing
59
+ all weight matrices.
60
+ norm_eps (`float`, *optional*, defaults to 1e-8):
61
+ The epsilon used by the normalization layers.
62
+ use_cache (`bool`, *optional*, defaults to `True`):
63
+ Whether or not the model should return the last key/values attentions
64
+ (not used by all models). Only relevant if `config.is_decoder=True`.
65
+ tie_word_embeddings(`bool`, *optional*, defaults to `False`):
66
+ Whether to tie weight embeddings
67
+ """
68
+ model_type = "stablelm_epoch"
69
+ keys_to_ignore_at_inference = ["past_key_values"]
70
+
71
+ def __init__(
72
+ self,
73
+ vocab_size=50_304,
74
+ intermediate_size=6912,
75
+ hidden_size=2560,
76
+ num_hidden_layers=32,
77
+ num_attention_heads=32,
78
+ num_key_value_heads=32,
79
+ hidden_act="silu",
80
+ rope_pct=0.25,
81
+ rope_theta=10_000,
82
+ max_position_embeddings=4096,
83
+ initializer_range=0.02,
84
+ norm_eps=1.0e-5,
85
+ use_cache=True,
86
+ bos_token_id=0,
87
+ eos_token_id=2,
88
+ tie_word_embeddings=False,
89
+ **kwargs,
90
+ ):
91
+ self.vocab_size = vocab_size
92
+ self.max_position_embeddings = max_position_embeddings
93
+ self.intermediate_size = intermediate_size
94
+ self.hidden_size = hidden_size
95
+ self.num_hidden_layers = num_hidden_layers
96
+ self.num_attention_heads = num_attention_heads
97
+ self.num_key_value_heads = num_key_value_heads
98
+ self.hidden_act = hidden_act
99
+ self.rope_pct = rope_pct
100
+ self.rope_theta = rope_theta
101
+ self.initializer_range = initializer_range
102
+ self.norm_eps = norm_eps
103
+ self.use_cache = use_cache
104
+ self.tie_word_embeddings = tie_word_embeddings
105
+ super().__init__(
106
+ bos_token_id=bos_token_id,
107
+ eos_token_id=eos_token_id,
108
+ tie_word_embeddings=tie_word_embeddings,
109
+ **kwargs,
110
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 0,
4
+ "eos_token_id": 0,
5
+ "transformers_version": "4.36.2"
6
+ }
model-00001-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:49c5f9b485b4cdb9bcafbbe0815bdf6da6b270e7b69fb0af13a2ef0cd3b5df49
3
+ size 1959695944
model-00002-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:36f37af830a590aeb5e124f6e0669b5316885a7033d8fdc64528dd34037f0dc7
3
+ size 1974197696
model-00003-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b36eb50719f77d996b909cef63c51633ebdddce63037011da3d59104dd5209cc
3
+ size 1974239144
model-00004-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:87d098a48eb07a51c400ec14bb8f285b63357a84e39f750c2e9d307b651c2355
3
+ size 1982062296
model-00005-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1be080e358c162b896ba974e498168baa9a995a9efb7c9f3f8a95488893063c2
3
+ size 1929633264
model-00006-of-00006.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6598d417dfbe99fd45f45ca25aa60384dcf628c06e182bbb15e1bce0d6d83bf1
3
+ size 1361985040
model.safetensors.index.json ADDED
@@ -0,0 +1,363 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 11181772800
4
+ },
5
+ "weight_map": {
6
+ "lm_head.weight": "model-00006-of-00006.safetensors",
7
+ "model.embed_tokens.weight": "model-00001-of-00006.safetensors",
8
+ "model.layers.0.input_layernorm.bias": "model-00001-of-00006.safetensors",
9
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00006.safetensors",
10
+ "model.layers.0.mlp.down_proj.weight": "model-00001-of-00006.safetensors",
11
+ "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00006.safetensors",
12
+ "model.layers.0.mlp.up_proj.weight": "model-00001-of-00006.safetensors",
13
+ "model.layers.0.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
14
+ "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
15
+ "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00006.safetensors",
16
+ "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00006.safetensors",
17
+ "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00006.safetensors",
18
+ "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00006.safetensors",
19
+ "model.layers.1.input_layernorm.bias": "model-00001-of-00006.safetensors",
20
+ "model.layers.1.input_layernorm.weight": "model-00001-of-00006.safetensors",
21
+ "model.layers.1.mlp.down_proj.weight": "model-00001-of-00006.safetensors",
22
+ "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00006.safetensors",
23
+ "model.layers.1.mlp.up_proj.weight": "model-00001-of-00006.safetensors",
24
+ "model.layers.1.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
25
+ "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
26
+ "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00006.safetensors",
27
+ "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00006.safetensors",
28
+ "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00006.safetensors",
29
+ "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00006.safetensors",
30
+ "model.layers.10.input_layernorm.bias": "model-00003-of-00006.safetensors",
31
+ "model.layers.10.input_layernorm.weight": "model-00003-of-00006.safetensors",
32
+ "model.layers.10.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
33
+ "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
34
+ "model.layers.10.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
35
+ "model.layers.10.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
36
+ "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
37
+ "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
38
+ "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
39
+ "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
40
+ "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
41
+ "model.layers.11.input_layernorm.bias": "model-00003-of-00006.safetensors",
42
+ "model.layers.11.input_layernorm.weight": "model-00003-of-00006.safetensors",
43
+ "model.layers.11.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
44
+ "model.layers.11.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
45
+ "model.layers.11.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
46
+ "model.layers.11.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
47
+ "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
48
+ "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
49
+ "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
50
+ "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
51
+ "model.layers.11.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
52
+ "model.layers.12.input_layernorm.bias": "model-00003-of-00006.safetensors",
53
+ "model.layers.12.input_layernorm.weight": "model-00003-of-00006.safetensors",
54
+ "model.layers.12.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
55
+ "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
56
+ "model.layers.12.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
57
+ "model.layers.12.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
58
+ "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
59
+ "model.layers.12.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
60
+ "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
61
+ "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
62
+ "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
63
+ "model.layers.13.input_layernorm.bias": "model-00003-of-00006.safetensors",
64
+ "model.layers.13.input_layernorm.weight": "model-00003-of-00006.safetensors",
65
+ "model.layers.13.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
66
+ "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
67
+ "model.layers.13.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
68
+ "model.layers.13.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
69
+ "model.layers.13.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
70
+ "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
71
+ "model.layers.13.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
72
+ "model.layers.13.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
73
+ "model.layers.13.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
74
+ "model.layers.14.input_layernorm.bias": "model-00003-of-00006.safetensors",
75
+ "model.layers.14.input_layernorm.weight": "model-00003-of-00006.safetensors",
76
+ "model.layers.14.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
77
+ "model.layers.14.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
78
+ "model.layers.14.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
79
+ "model.layers.14.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
80
+ "model.layers.14.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
81
+ "model.layers.14.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
82
+ "model.layers.14.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
83
+ "model.layers.14.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
84
+ "model.layers.14.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
85
+ "model.layers.15.input_layernorm.bias": "model-00003-of-00006.safetensors",
86
+ "model.layers.15.input_layernorm.weight": "model-00003-of-00006.safetensors",
87
+ "model.layers.15.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
88
+ "model.layers.15.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
89
+ "model.layers.15.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
90
+ "model.layers.15.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
91
+ "model.layers.15.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
92
+ "model.layers.15.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
93
+ "model.layers.15.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
94
+ "model.layers.15.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
95
+ "model.layers.15.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
96
+ "model.layers.16.input_layernorm.bias": "model-00003-of-00006.safetensors",
97
+ "model.layers.16.input_layernorm.weight": "model-00003-of-00006.safetensors",
98
+ "model.layers.16.mlp.down_proj.weight": "model-00003-of-00006.safetensors",
99
+ "model.layers.16.mlp.gate_proj.weight": "model-00003-of-00006.safetensors",
100
+ "model.layers.16.mlp.up_proj.weight": "model-00003-of-00006.safetensors",
101
+ "model.layers.16.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
102
+ "model.layers.16.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
103
+ "model.layers.16.self_attn.k_proj.weight": "model-00003-of-00006.safetensors",
104
+ "model.layers.16.self_attn.o_proj.weight": "model-00003-of-00006.safetensors",
105
+ "model.layers.16.self_attn.q_proj.weight": "model-00003-of-00006.safetensors",
106
+ "model.layers.16.self_attn.v_proj.weight": "model-00003-of-00006.safetensors",
107
+ "model.layers.17.input_layernorm.bias": "model-00004-of-00006.safetensors",
108
+ "model.layers.17.input_layernorm.weight": "model-00004-of-00006.safetensors",
109
+ "model.layers.17.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
110
+ "model.layers.17.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
111
+ "model.layers.17.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
112
+ "model.layers.17.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
113
+ "model.layers.17.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
114
+ "model.layers.17.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
115
+ "model.layers.17.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
116
+ "model.layers.17.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
117
+ "model.layers.17.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
118
+ "model.layers.18.input_layernorm.bias": "model-00004-of-00006.safetensors",
119
+ "model.layers.18.input_layernorm.weight": "model-00004-of-00006.safetensors",
120
+ "model.layers.18.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
121
+ "model.layers.18.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
122
+ "model.layers.18.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
123
+ "model.layers.18.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
124
+ "model.layers.18.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
125
+ "model.layers.18.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
126
+ "model.layers.18.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
127
+ "model.layers.18.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
128
+ "model.layers.18.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
129
+ "model.layers.19.input_layernorm.bias": "model-00004-of-00006.safetensors",
130
+ "model.layers.19.input_layernorm.weight": "model-00004-of-00006.safetensors",
131
+ "model.layers.19.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
132
+ "model.layers.19.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
133
+ "model.layers.19.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
134
+ "model.layers.19.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
135
+ "model.layers.19.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
136
+ "model.layers.19.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
137
+ "model.layers.19.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
138
+ "model.layers.19.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
139
+ "model.layers.19.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
140
+ "model.layers.2.input_layernorm.bias": "model-00001-of-00006.safetensors",
141
+ "model.layers.2.input_layernorm.weight": "model-00001-of-00006.safetensors",
142
+ "model.layers.2.mlp.down_proj.weight": "model-00001-of-00006.safetensors",
143
+ "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00006.safetensors",
144
+ "model.layers.2.mlp.up_proj.weight": "model-00001-of-00006.safetensors",
145
+ "model.layers.2.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
146
+ "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
147
+ "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00006.safetensors",
148
+ "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00006.safetensors",
149
+ "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00006.safetensors",
150
+ "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00006.safetensors",
151
+ "model.layers.20.input_layernorm.bias": "model-00004-of-00006.safetensors",
152
+ "model.layers.20.input_layernorm.weight": "model-00004-of-00006.safetensors",
153
+ "model.layers.20.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
154
+ "model.layers.20.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
155
+ "model.layers.20.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
156
+ "model.layers.20.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
157
+ "model.layers.20.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
158
+ "model.layers.20.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
159
+ "model.layers.20.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
160
+ "model.layers.20.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
161
+ "model.layers.20.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
162
+ "model.layers.21.input_layernorm.bias": "model-00004-of-00006.safetensors",
163
+ "model.layers.21.input_layernorm.weight": "model-00004-of-00006.safetensors",
164
+ "model.layers.21.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
165
+ "model.layers.21.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
166
+ "model.layers.21.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
167
+ "model.layers.21.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
168
+ "model.layers.21.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
169
+ "model.layers.21.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
170
+ "model.layers.21.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
171
+ "model.layers.21.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
172
+ "model.layers.21.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
173
+ "model.layers.22.input_layernorm.bias": "model-00004-of-00006.safetensors",
174
+ "model.layers.22.input_layernorm.weight": "model-00004-of-00006.safetensors",
175
+ "model.layers.22.mlp.down_proj.weight": "model-00004-of-00006.safetensors",
176
+ "model.layers.22.mlp.gate_proj.weight": "model-00004-of-00006.safetensors",
177
+ "model.layers.22.mlp.up_proj.weight": "model-00004-of-00006.safetensors",
178
+ "model.layers.22.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
179
+ "model.layers.22.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
180
+ "model.layers.22.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
181
+ "model.layers.22.self_attn.o_proj.weight": "model-00004-of-00006.safetensors",
182
+ "model.layers.22.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
183
+ "model.layers.22.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
184
+ "model.layers.23.input_layernorm.bias": "model-00005-of-00006.safetensors",
185
+ "model.layers.23.input_layernorm.weight": "model-00005-of-00006.safetensors",
186
+ "model.layers.23.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
187
+ "model.layers.23.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
188
+ "model.layers.23.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
189
+ "model.layers.23.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
190
+ "model.layers.23.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
191
+ "model.layers.23.self_attn.k_proj.weight": "model-00004-of-00006.safetensors",
192
+ "model.layers.23.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
193
+ "model.layers.23.self_attn.q_proj.weight": "model-00004-of-00006.safetensors",
194
+ "model.layers.23.self_attn.v_proj.weight": "model-00004-of-00006.safetensors",
195
+ "model.layers.24.input_layernorm.bias": "model-00005-of-00006.safetensors",
196
+ "model.layers.24.input_layernorm.weight": "model-00005-of-00006.safetensors",
197
+ "model.layers.24.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
198
+ "model.layers.24.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
199
+ "model.layers.24.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
200
+ "model.layers.24.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
201
+ "model.layers.24.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
202
+ "model.layers.24.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
203
+ "model.layers.24.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
204
+ "model.layers.24.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
205
+ "model.layers.24.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
206
+ "model.layers.25.input_layernorm.bias": "model-00005-of-00006.safetensors",
207
+ "model.layers.25.input_layernorm.weight": "model-00005-of-00006.safetensors",
208
+ "model.layers.25.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
209
+ "model.layers.25.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
210
+ "model.layers.25.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
211
+ "model.layers.25.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
212
+ "model.layers.25.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
213
+ "model.layers.25.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
214
+ "model.layers.25.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
215
+ "model.layers.25.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
216
+ "model.layers.25.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
217
+ "model.layers.26.input_layernorm.bias": "model-00005-of-00006.safetensors",
218
+ "model.layers.26.input_layernorm.weight": "model-00005-of-00006.safetensors",
219
+ "model.layers.26.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
220
+ "model.layers.26.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
221
+ "model.layers.26.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
222
+ "model.layers.26.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
223
+ "model.layers.26.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
224
+ "model.layers.26.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
225
+ "model.layers.26.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
226
+ "model.layers.26.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
227
+ "model.layers.26.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
228
+ "model.layers.27.input_layernorm.bias": "model-00005-of-00006.safetensors",
229
+ "model.layers.27.input_layernorm.weight": "model-00005-of-00006.safetensors",
230
+ "model.layers.27.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
231
+ "model.layers.27.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
232
+ "model.layers.27.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
233
+ "model.layers.27.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
234
+ "model.layers.27.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
235
+ "model.layers.27.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
236
+ "model.layers.27.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
237
+ "model.layers.27.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
238
+ "model.layers.27.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
239
+ "model.layers.28.input_layernorm.bias": "model-00005-of-00006.safetensors",
240
+ "model.layers.28.input_layernorm.weight": "model-00005-of-00006.safetensors",
241
+ "model.layers.28.mlp.down_proj.weight": "model-00005-of-00006.safetensors",
242
+ "model.layers.28.mlp.gate_proj.weight": "model-00005-of-00006.safetensors",
243
+ "model.layers.28.mlp.up_proj.weight": "model-00005-of-00006.safetensors",
244
+ "model.layers.28.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
245
+ "model.layers.28.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
246
+ "model.layers.28.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
247
+ "model.layers.28.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
248
+ "model.layers.28.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
249
+ "model.layers.28.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
250
+ "model.layers.29.input_layernorm.bias": "model-00006-of-00006.safetensors",
251
+ "model.layers.29.input_layernorm.weight": "model-00006-of-00006.safetensors",
252
+ "model.layers.29.mlp.down_proj.weight": "model-00006-of-00006.safetensors",
253
+ "model.layers.29.mlp.gate_proj.weight": "model-00006-of-00006.safetensors",
254
+ "model.layers.29.mlp.up_proj.weight": "model-00006-of-00006.safetensors",
255
+ "model.layers.29.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
256
+ "model.layers.29.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
257
+ "model.layers.29.self_attn.k_proj.weight": "model-00005-of-00006.safetensors",
258
+ "model.layers.29.self_attn.o_proj.weight": "model-00005-of-00006.safetensors",
259
+ "model.layers.29.self_attn.q_proj.weight": "model-00005-of-00006.safetensors",
260
+ "model.layers.29.self_attn.v_proj.weight": "model-00005-of-00006.safetensors",
261
+ "model.layers.3.input_layernorm.bias": "model-00001-of-00006.safetensors",
262
+ "model.layers.3.input_layernorm.weight": "model-00001-of-00006.safetensors",
263
+ "model.layers.3.mlp.down_proj.weight": "model-00001-of-00006.safetensors",
264
+ "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00006.safetensors",
265
+ "model.layers.3.mlp.up_proj.weight": "model-00001-of-00006.safetensors",
266
+ "model.layers.3.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
267
+ "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
268
+ "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00006.safetensors",
269
+ "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00006.safetensors",
270
+ "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00006.safetensors",
271
+ "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00006.safetensors",
272
+ "model.layers.30.input_layernorm.bias": "model-00006-of-00006.safetensors",
273
+ "model.layers.30.input_layernorm.weight": "model-00006-of-00006.safetensors",
274
+ "model.layers.30.mlp.down_proj.weight": "model-00006-of-00006.safetensors",
275
+ "model.layers.30.mlp.gate_proj.weight": "model-00006-of-00006.safetensors",
276
+ "model.layers.30.mlp.up_proj.weight": "model-00006-of-00006.safetensors",
277
+ "model.layers.30.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
278
+ "model.layers.30.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
279
+ "model.layers.30.self_attn.k_proj.weight": "model-00006-of-00006.safetensors",
280
+ "model.layers.30.self_attn.o_proj.weight": "model-00006-of-00006.safetensors",
281
+ "model.layers.30.self_attn.q_proj.weight": "model-00006-of-00006.safetensors",
282
+ "model.layers.30.self_attn.v_proj.weight": "model-00006-of-00006.safetensors",
283
+ "model.layers.31.input_layernorm.bias": "model-00006-of-00006.safetensors",
284
+ "model.layers.31.input_layernorm.weight": "model-00006-of-00006.safetensors",
285
+ "model.layers.31.mlp.down_proj.weight": "model-00006-of-00006.safetensors",
286
+ "model.layers.31.mlp.gate_proj.weight": "model-00006-of-00006.safetensors",
287
+ "model.layers.31.mlp.up_proj.weight": "model-00006-of-00006.safetensors",
288
+ "model.layers.31.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
289
+ "model.layers.31.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
290
+ "model.layers.31.self_attn.k_proj.weight": "model-00006-of-00006.safetensors",
291
+ "model.layers.31.self_attn.o_proj.weight": "model-00006-of-00006.safetensors",
292
+ "model.layers.31.self_attn.q_proj.weight": "model-00006-of-00006.safetensors",
293
+ "model.layers.31.self_attn.v_proj.weight": "model-00006-of-00006.safetensors",
294
+ "model.layers.4.input_layernorm.bias": "model-00002-of-00006.safetensors",
295
+ "model.layers.4.input_layernorm.weight": "model-00002-of-00006.safetensors",
296
+ "model.layers.4.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
297
+ "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00006.safetensors",
298
+ "model.layers.4.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
299
+ "model.layers.4.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
300
+ "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
301
+ "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00006.safetensors",
302
+ "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00006.safetensors",
303
+ "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00006.safetensors",
304
+ "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00006.safetensors",
305
+ "model.layers.5.input_layernorm.bias": "model-00002-of-00006.safetensors",
306
+ "model.layers.5.input_layernorm.weight": "model-00002-of-00006.safetensors",
307
+ "model.layers.5.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
308
+ "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
309
+ "model.layers.5.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
310
+ "model.layers.5.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
311
+ "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
312
+ "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
313
+ "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
314
+ "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
315
+ "model.layers.5.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
316
+ "model.layers.6.input_layernorm.bias": "model-00002-of-00006.safetensors",
317
+ "model.layers.6.input_layernorm.weight": "model-00002-of-00006.safetensors",
318
+ "model.layers.6.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
319
+ "model.layers.6.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
320
+ "model.layers.6.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
321
+ "model.layers.6.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
322
+ "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
323
+ "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
324
+ "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
325
+ "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
326
+ "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
327
+ "model.layers.7.input_layernorm.bias": "model-00002-of-00006.safetensors",
328
+ "model.layers.7.input_layernorm.weight": "model-00002-of-00006.safetensors",
329
+ "model.layers.7.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
330
+ "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
331
+ "model.layers.7.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
332
+ "model.layers.7.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
333
+ "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
334
+ "model.layers.7.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
335
+ "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
336
+ "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
337
+ "model.layers.7.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
338
+ "model.layers.8.input_layernorm.bias": "model-00002-of-00006.safetensors",
339
+ "model.layers.8.input_layernorm.weight": "model-00002-of-00006.safetensors",
340
+ "model.layers.8.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
341
+ "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
342
+ "model.layers.8.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
343
+ "model.layers.8.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
344
+ "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
345
+ "model.layers.8.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
346
+ "model.layers.8.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
347
+ "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
348
+ "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
349
+ "model.layers.9.input_layernorm.bias": "model-00002-of-00006.safetensors",
350
+ "model.layers.9.input_layernorm.weight": "model-00002-of-00006.safetensors",
351
+ "model.layers.9.mlp.down_proj.weight": "model-00002-of-00006.safetensors",
352
+ "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00006.safetensors",
353
+ "model.layers.9.mlp.up_proj.weight": "model-00002-of-00006.safetensors",
354
+ "model.layers.9.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
355
+ "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
356
+ "model.layers.9.self_attn.k_proj.weight": "model-00002-of-00006.safetensors",
357
+ "model.layers.9.self_attn.o_proj.weight": "model-00002-of-00006.safetensors",
358
+ "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00006.safetensors",
359
+ "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00006.safetensors",
360
+ "model.norm.bias": "model-00006-of-00006.safetensors",
361
+ "model.norm.weight": "model-00006-of-00006.safetensors"
362
+ }
363
+ }
modeling_stablelm_epoch.py ADDED
@@ -0,0 +1,916 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
195
+ if (self.head_dim * self.num_heads) != self.hidden_size:
196
+ raise ValueError(
197
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
198
+ f" and `num_heads`: {self.num_heads})."
199
+ )
200
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
201
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
202
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
203
+ self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
204
+
205
+ self._init_rope()
206
+
207
+ def _init_rope(self):
208
+ self.rotary_ndims = int(self.head_dim * self.config.rope_pct)
209
+ self.rotary_emb = RotaryEmbedding(
210
+ self.rotary_ndims,
211
+ max_position_embeddings=self.config.max_position_embeddings,
212
+ base=self.config.rope_theta,
213
+ )
214
+
215
+ def forward(
216
+ self,
217
+ hidden_states: torch.FloatTensor,
218
+ attention_mask: torch.FloatTensor,
219
+ position_ids: torch.LongTensor,
220
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
221
+ output_attentions: Optional[bool] = False,
222
+ use_cache: Optional[bool] = False,
223
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
224
+ bsz, q_len, _ = hidden_states.size()
225
+
226
+ query_states = self.q_proj(hidden_states)
227
+ key_states = self.k_proj(hidden_states)
228
+ value_states = self.v_proj(hidden_states)
229
+
230
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
231
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
232
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
233
+
234
+ query_rot = query_states[..., : self.rotary_ndims]
235
+ query_pass = query_states[..., self.rotary_ndims :]
236
+ key_rot = key_states[..., : self.rotary_ndims]
237
+ key_pass = key_states[..., self.rotary_ndims :]
238
+
239
+ kv_seq_len = key_states.shape[-2]
240
+ if past_key_value is not None:
241
+ kv_seq_len += past_key_value[0].shape[-2]
242
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
243
+ query_states, key_states = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
244
+
245
+ # [batch_size, num_heads, seq_len, head_dim]
246
+ query_states = torch.cat((query_states, query_pass), dim=-1)
247
+ key_states = torch.cat((key_states, key_pass), dim=-1)
248
+
249
+ if past_key_value is not None:
250
+ # Reuse k, v, self_attention
251
+ key_states = torch.cat((past_key_value[0], key_states), dim=2)
252
+ value_states = torch.cat((past_key_value[1], value_states), dim=2)
253
+
254
+ past_key_value = (key_states, value_states) if use_cache else None
255
+
256
+ # Repeat k/v heads if n_kv_heads < n_heads
257
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
258
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
259
+
260
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
261
+
262
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
263
+ raise ValueError(
264
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
265
+ f" {attn_weights.size()}"
266
+ )
267
+
268
+ if attention_mask is not None:
269
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
270
+ raise ValueError(
271
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
272
+ )
273
+ attn_weights = attn_weights + attention_mask
274
+
275
+ # Upcast attention to fp32
276
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
277
+ attn_output = torch.matmul(attn_weights, value_states)
278
+
279
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
280
+ raise ValueError(
281
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
282
+ f" {attn_output.size()}"
283
+ )
284
+
285
+ # Merge heads
286
+ attn_output = attn_output.transpose(1, 2).contiguous()
287
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
288
+
289
+ # Final linear projection
290
+ attn_output = self.o_proj(attn_output)
291
+
292
+ if not output_attentions:
293
+ attn_weights = None
294
+
295
+ return attn_output, attn_weights, past_key_value
296
+
297
+
298
+ class FlashAttention2(Attention):
299
+ """
300
+ Reference: https://github.com/huggingface/transformers/blob/5d36025ca13d05151b7a0c761e90d429c4644a30/src/transformers/models/llama/modeling_llama.py#L456
301
+ """
302
+
303
+ def __init__(self, *args, **kwargs):
304
+ super().__init__(*args, **kwargs)
305
+
306
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
307
+ # 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.
308
+ # 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).
309
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
310
+
311
+ def forward(
312
+ self,
313
+ hidden_states: torch.Tensor,
314
+ attention_mask: Optional[torch.LongTensor] = None,
315
+ position_ids: Optional[torch.LongTensor] = None,
316
+ past_key_value: Optional[Cache] = None,
317
+ output_attentions: bool = False,
318
+ use_cache: bool = False,
319
+ **kwargs,
320
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
321
+ # FlashAttention2 attention does not support output_attentions
322
+ if "padding_mask" in kwargs:
323
+ warnings.warn(
324
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
325
+ )
326
+
327
+ # overwrite attention_mask with padding_mask
328
+ attention_mask = kwargs.pop("padding_mask")
329
+
330
+ output_attentions = False
331
+
332
+ bsz, q_len, _ = hidden_states.size()
333
+
334
+ query_states = self.q_proj(hidden_states)
335
+ key_states = self.k_proj(hidden_states)
336
+ value_states = self.v_proj(hidden_states)
337
+
338
+ # Flash attention requires the input to have the shape
339
+ # batch_size x seq_length x head_dim x hidden_dim
340
+ # therefore we just need to keep the original shape
341
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
342
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
343
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
344
+
345
+ query_rot = query_states[..., : self.rotary_ndims]
346
+ query_pass = query_states[..., self.rotary_ndims :]
347
+ key_rot = key_states[..., : self.rotary_ndims]
348
+ key_pass = key_states[..., self.rotary_ndims :]
349
+
350
+ kv_seq_len = key_states.shape[-2]
351
+ if past_key_value is not None:
352
+ kv_seq_len += past_key_value[0].shape[-2]
353
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
354
+ query_states, key_states = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
355
+
356
+ # [batch_size, num_heads, seq_len, head_dim]
357
+ query_states = torch.cat((query_states, query_pass), dim=-1)
358
+ key_states = torch.cat((key_states, key_pass), dim=-1)
359
+
360
+ if past_key_value is not None:
361
+ # Reuse k, v, self_attention
362
+ key_states = torch.cat((past_key_value[0], key_states), dim=2)
363
+ value_states = torch.cat((past_key_value[1], value_states), dim=2)
364
+
365
+ past_key_value = (key_states, value_states) if use_cache else None
366
+
367
+ # 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
368
+ # to be able to avoid many of these transpose/reshape/view.
369
+ query_states = query_states.transpose(1, 2)
370
+ key_states = key_states.transpose(1, 2)
371
+ value_states = value_states.transpose(1, 2)
372
+
373
+ dropout_rate = self.attention_dropout if self.training else 0.0
374
+
375
+ attn_output = self._flash_attention_forward(
376
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
377
+ )
378
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
379
+ attn_output = self.o_proj(attn_output)
380
+
381
+ if not output_attentions:
382
+ attn_weights = None
383
+
384
+ return attn_output, attn_weights, past_key_value
385
+
386
+ def _flash_attention_forward(
387
+ self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
388
+ ):
389
+ """
390
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
391
+ first unpad the input, then computes the attention scores and pad the final attention scores.
392
+
393
+ Args:
394
+ query_states (`torch.Tensor`):
395
+ Input query states to be passed to Flash Attention API
396
+ key_states (`torch.Tensor`):
397
+ Input key states to be passed to Flash Attention API
398
+ value_states (`torch.Tensor`):
399
+ Input value states to be passed to Flash Attention API
400
+ attention_mask (`torch.Tensor`):
401
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
402
+ position of padding tokens and 1 for the position of non-padding tokens.
403
+ dropout (`int`, *optional*):
404
+ Attention dropout
405
+ softmax_scale (`float`, *optional*):
406
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
407
+ """
408
+ if not self._flash_attn_uses_top_left_mask:
409
+ causal = self.is_causal
410
+ else:
411
+ # 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__.
412
+ causal = self.is_causal and query_length != 1
413
+
414
+ # Contains at least one padding token in the sequence
415
+ if attention_mask is not None:
416
+ batch_size = query_states.shape[0]
417
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
418
+ query_states, key_states, value_states, attention_mask, query_length
419
+ )
420
+
421
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
422
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
423
+
424
+ attn_output_unpad = flash_attn_varlen_func(
425
+ query_states,
426
+ key_states,
427
+ value_states,
428
+ cu_seqlens_q=cu_seqlens_q,
429
+ cu_seqlens_k=cu_seqlens_k,
430
+ max_seqlen_q=max_seqlen_in_batch_q,
431
+ max_seqlen_k=max_seqlen_in_batch_k,
432
+ dropout_p=dropout,
433
+ softmax_scale=softmax_scale,
434
+ causal=causal,
435
+ )
436
+
437
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
438
+ else:
439
+ attn_output = flash_attn_func(
440
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
441
+ )
442
+
443
+ return attn_output
444
+
445
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
446
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
447
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
448
+
449
+ key_layer = index_first_axis(
450
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
451
+ )
452
+ value_layer = index_first_axis(
453
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
454
+ )
455
+ if query_length == kv_seq_len:
456
+ query_layer = index_first_axis(
457
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
458
+ )
459
+ cu_seqlens_q = cu_seqlens_k
460
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
461
+ indices_q = indices_k
462
+ elif query_length == 1:
463
+ max_seqlen_in_batch_q = 1
464
+ cu_seqlens_q = torch.arange(
465
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
466
+ ) # There is a memcpy here, that is very bad.
467
+ indices_q = cu_seqlens_q[:-1]
468
+ query_layer = query_layer.squeeze(1)
469
+ else:
470
+ # The -q_len: slice assumes left padding.
471
+ attention_mask = attention_mask[:, -query_length:]
472
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
473
+
474
+ return (
475
+ query_layer,
476
+ key_layer,
477
+ value_layer,
478
+ indices_q,
479
+ (cu_seqlens_q, cu_seqlens_k),
480
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
481
+ )
482
+
483
+
484
+ ATTENTION_CLASSES = {
485
+ "eager": Attention,
486
+ "flash_attention_2": FlashAttention2,
487
+ }
488
+
489
+
490
+ class DecoderLayer(nn.Module):
491
+ def __init__(self, config: StableLMEpochConfig):
492
+ super().__init__()
493
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](config=config)
494
+ self.mlp = MLP(config)
495
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
496
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
497
+
498
+ def forward(
499
+ self,
500
+ hidden_states: Optional[torch.FloatTensor],
501
+ attention_mask: Optional[torch.FloatTensor] = None,
502
+ position_ids: Optional[torch.LongTensor] = None,
503
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
504
+ output_attentions: Optional[bool] = False,
505
+ use_cache: Optional[bool] = False,
506
+ ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
507
+ residual = hidden_states
508
+
509
+ hidden_states = self.input_layernorm(hidden_states)
510
+
511
+ # Self Attention
512
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
513
+ hidden_states=hidden_states,
514
+ attention_mask=attention_mask,
515
+ position_ids=position_ids,
516
+ past_key_value=past_key_value,
517
+ output_attentions=output_attentions,
518
+ use_cache=use_cache,
519
+ )
520
+ hidden_states = residual + hidden_states
521
+
522
+ # Fully Connected
523
+ residual = hidden_states
524
+ hidden_states = self.post_attention_layernorm(hidden_states)
525
+ hidden_states = self.mlp(hidden_states)
526
+ hidden_states = residual + hidden_states
527
+
528
+ outputs = (hidden_states,)
529
+
530
+ if output_attentions:
531
+ outputs += (self_attn_weights,)
532
+
533
+ if use_cache:
534
+ outputs += (present_key_value,)
535
+
536
+ return outputs
537
+
538
+
539
+ class StableLMEpochPreTrainedModel(PreTrainedModel):
540
+ """An abstract class to handle weights initialization and a simple interface
541
+ for downloading and loading pretrained models.
542
+ """
543
+
544
+ config_class = StableLMEpochConfig
545
+ base_model_prefix = "transformer"
546
+ supports_gradient_checkpointing = True
547
+ _no_split_modules = ["DecoderLayer"]
548
+ _skip_keys_device_placement = "past_key_values"
549
+ _supports_flash_attn_2 = True
550
+
551
+ def _init_weights(self, module: nn.Module):
552
+ """Initialize the weights"""
553
+ if isinstance(module, nn.Linear):
554
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
555
+ if module.bias is not None:
556
+ module.bias.data.zero_()
557
+ elif isinstance(module, nn.Embedding):
558
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
559
+ if module.padding_idx is not None:
560
+ module.weight.data[module.padding_idx].zero_()
561
+ elif isinstance(module, nn.LayerNorm):
562
+ module.bias.data.zero_()
563
+ module.weight.data.fill_(1.0)
564
+
565
+ def _set_gradient_checkpointing(self, module: nn.Module, value=False):
566
+ if isinstance(module, StableLMEpochModel):
567
+ module.gradient_checkpointing = value
568
+
569
+
570
+ class StableLMEpochModel(StableLMEpochPreTrainedModel):
571
+ def __init__(self, config: StableLMEpochConfig):
572
+ super().__init__(config)
573
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
574
+ self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
575
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
576
+
577
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
578
+ self.gradient_checkpointing = False
579
+ # Initialize weights and apply final processing
580
+ self.post_init()
581
+
582
+ def get_input_embeddings(self):
583
+ return self.embed_tokens
584
+
585
+ def set_input_embeddings(self, value: nn.Module):
586
+ self.embed_tokens = value
587
+
588
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
589
+ def _prepare_decoder_attention_mask(
590
+ self,
591
+ attention_mask: torch.Tensor,
592
+ input_shape: torch.Size,
593
+ inputs_embeds: torch.Tensor,
594
+ past_key_values_length: int,
595
+ ):
596
+ # Create causal mask
597
+ # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
598
+ combined_attention_mask = None
599
+ if input_shape[-1] > 1:
600
+ combined_attention_mask = _make_causal_mask(
601
+ input_shape,
602
+ inputs_embeds.dtype,
603
+ device=inputs_embeds.device,
604
+ past_key_values_length=past_key_values_length,
605
+ )
606
+
607
+ if attention_mask is not None:
608
+ # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
609
+ expanded_attn_mask = _expand_mask(
610
+ attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
611
+ ).to(inputs_embeds.device)
612
+ combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
613
+
614
+ return combined_attention_mask
615
+
616
+ def forward(
617
+ self,
618
+ input_ids: Optional[torch.LongTensor] = None,
619
+ attention_mask: Optional[torch.FloatTensor] = None,
620
+ position_ids: Optional[torch.LongTensor] = None,
621
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
622
+ inputs_embeds: Optional[torch.FloatTensor] = None,
623
+ use_cache: Optional[bool] = None,
624
+ output_attentions: Optional[bool] = None,
625
+ output_hidden_states: Optional[bool] = None,
626
+ return_dict: Optional[bool] = None,
627
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
628
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
629
+ output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
630
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
631
+
632
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
633
+
634
+ # Retrieve input_ids and inputs_embeds
635
+ if input_ids is not None and inputs_embeds is not None:
636
+ raise ValueError(
637
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
638
+ )
639
+ elif input_ids is not None:
640
+ batch_size, seq_length = input_ids.shape
641
+ elif inputs_embeds is not None:
642
+ batch_size, seq_length, _ = inputs_embeds.shape
643
+ else:
644
+ raise ValueError(
645
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
646
+ )
647
+
648
+ seq_length_with_past = seq_length
649
+ past_key_values_length = 0
650
+
651
+ if position_ids is None:
652
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
653
+ position_ids = torch.arange(
654
+ past_key_values_length,
655
+ seq_length + past_key_values_length,
656
+ dtype=torch.long,
657
+ device=device,
658
+ )
659
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
660
+ else:
661
+ position_ids = position_ids.view(-1, seq_length).long()
662
+
663
+ if inputs_embeds is None:
664
+ inputs_embeds = self.embed_tokens(input_ids)
665
+ # Embed positions
666
+ if self._use_flash_attention_2:
667
+ # 2d mask is passed through the layers
668
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
669
+ else:
670
+ if attention_mask is None:
671
+ attention_mask = torch.ones(
672
+ (batch_size, seq_length_with_past),
673
+ dtype=torch.bool,
674
+ device=inputs_embeds.device,
675
+ )
676
+ attention_mask = self._prepare_decoder_attention_mask(
677
+ attention_mask,
678
+ (batch_size, seq_length),
679
+ inputs_embeds,
680
+ past_key_values_length,
681
+ )
682
+
683
+ hidden_states = inputs_embeds
684
+
685
+ if self.gradient_checkpointing and self.training:
686
+ if use_cache:
687
+ logger.warning(
688
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
689
+ )
690
+ use_cache = False
691
+
692
+ # Decoder layers
693
+ all_hidden_states = () if output_hidden_states else None
694
+ all_self_attns = () if output_attentions else None
695
+ next_decoder_cache = () if use_cache else None
696
+
697
+ for idx, decoder_layer in enumerate(self.layers):
698
+ if output_hidden_states:
699
+ all_hidden_states += (hidden_states,)
700
+
701
+ past_key_value = (
702
+ past_key_values[idx] if past_key_values is not None else None
703
+ )
704
+
705
+ if self.gradient_checkpointing and self.training:
706
+
707
+ def create_custom_forward(module):
708
+ def custom_forward(*inputs):
709
+ # None for past_key_value
710
+ return module(*inputs, past_key_value, output_attentions)
711
+
712
+ return custom_forward
713
+
714
+ layer_outputs = torch.utils.checkpoint.checkpoint(
715
+ create_custom_forward(decoder_layer),
716
+ hidden_states,
717
+ attention_mask,
718
+ position_ids,
719
+ )
720
+ else:
721
+ layer_outputs = decoder_layer(
722
+ hidden_states,
723
+ attention_mask=attention_mask,
724
+ position_ids=position_ids,
725
+ past_key_value=past_key_value,
726
+ output_attentions=output_attentions,
727
+ use_cache=use_cache,
728
+ )
729
+
730
+ hidden_states = layer_outputs[0]
731
+
732
+ if use_cache:
733
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
734
+
735
+ if output_attentions:
736
+ all_self_attns += (layer_outputs[1],)
737
+
738
+ hidden_states = self.norm(hidden_states)
739
+
740
+ # Add hidden states from the last decoder layer
741
+ if output_hidden_states:
742
+ all_hidden_states += (hidden_states,)
743
+
744
+ next_cache = next_decoder_cache if use_cache else None
745
+ if not return_dict:
746
+ return tuple(
747
+ v
748
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
749
+ if v is not None
750
+ )
751
+ return BaseModelOutputWithPast(
752
+ last_hidden_state=hidden_states,
753
+ past_key_values=next_cache,
754
+ hidden_states=all_hidden_states,
755
+ attentions=all_self_attns,
756
+ )
757
+
758
+
759
+ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
760
+ _tied_weights_keys = ["lm_head.weight"]
761
+
762
+ def __init__(self, config: StableLMEpochConfig):
763
+ super().__init__(config)
764
+
765
+ self.model = StableLMEpochModel(config)
766
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
767
+
768
+ # Initialize weights and apply final processing
769
+ self.post_init()
770
+
771
+ def get_input_embeddings(self):
772
+ return self.model.embed_tokens
773
+
774
+ def set_input_embeddings(self, value):
775
+ self.model.embed_tokens = value
776
+
777
+ def get_output_embeddings(self):
778
+ return self.lm_head
779
+
780
+ def set_output_embeddings(self, new_embeddings: nn.Module):
781
+ self.lm_head = new_embeddings
782
+
783
+ def get_decoder(self):
784
+ return self.model
785
+
786
+ def set_decoder(self, decoder):
787
+ self.model = decoder
788
+
789
+ def forward(
790
+ self,
791
+ input_ids: Optional[torch.LongTensor] = None,
792
+ attention_mask: Optional[torch.FloatTensor] = None,
793
+ position_ids: Optional[torch.LongTensor] = None,
794
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
795
+ inputs_embeds: Optional[torch.FloatTensor] = None,
796
+ labels: Optional[torch.LongTensor] = None,
797
+ use_cache: Optional[bool] = None,
798
+ output_attentions: Optional[bool] = None,
799
+ output_hidden_states: Optional[bool] = None,
800
+ return_dict: Optional[bool] = None,
801
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
802
+ output_attentions = (
803
+ output_attentions
804
+ if output_attentions is not None
805
+ else self.config.output_attentions
806
+ )
807
+ output_hidden_states = (
808
+ output_hidden_states
809
+ if output_hidden_states is not None
810
+ else self.config.output_hidden_states
811
+ )
812
+ return_dict = (
813
+ return_dict if return_dict is not None else self.config.use_return_dict
814
+ )
815
+
816
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
817
+ outputs = self.model(
818
+ input_ids,
819
+ attention_mask=attention_mask,
820
+ position_ids=position_ids,
821
+ past_key_values=past_key_values,
822
+ inputs_embeds=inputs_embeds,
823
+ use_cache=use_cache,
824
+ output_attentions=output_attentions,
825
+ output_hidden_states=output_hidden_states,
826
+ return_dict=return_dict,
827
+ )
828
+
829
+ hidden_states = outputs[0]
830
+ logits = self.lm_head(hidden_states).float()
831
+
832
+ loss = None
833
+ if labels is not None:
834
+ # Shift so that tokens < n predict n
835
+ shift_logits = logits[..., :-1, :].contiguous()
836
+ shift_labels = labels[..., 1:].contiguous()
837
+ # Flatten the tokens
838
+ loss_fct = CrossEntropyLoss()
839
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
840
+ shift_labels = shift_labels.view(-1)
841
+ # Enable model parallelism
842
+ shift_labels = shift_labels.to(shift_logits.device)
843
+ loss = loss_fct(shift_logits, shift_labels)
844
+
845
+ if not return_dict:
846
+ output = (logits,) + outputs[1:]
847
+ return (loss,) + output if loss is not None else output
848
+
849
+ return CausalLMOutputWithPast(
850
+ loss=loss,
851
+ logits=logits,
852
+ past_key_values=outputs.past_key_values,
853
+ hidden_states=outputs.hidden_states,
854
+ attentions=outputs.attentions,
855
+ )
856
+
857
+ def prepare_inputs_for_generation(
858
+ self,
859
+ input_ids,
860
+ past_key_values: Optional[torch.Tensor] = None,
861
+ attention_mask: Optional[torch.Tensor] = None,
862
+ inputs_embeds: Optional[torch.Tensor] = None,
863
+ **kwargs,
864
+ ):
865
+ # Trim decoder_input_ids if past is used
866
+ if past_key_values is not None:
867
+ past_length = past_key_values[0][0].shape[2]
868
+
869
+ # Some generation methods already pass only the last input ID
870
+ if input_ids.shape[1] > past_length:
871
+ remove_prefix_length = past_length
872
+ else:
873
+ # Default to old behavior: keep only final ID
874
+ remove_prefix_length = input_ids.shape[1] - 1
875
+
876
+ input_ids = input_ids[:, remove_prefix_length:]
877
+
878
+ position_ids = kwargs.get("position_ids", None)
879
+ if attention_mask is not None and position_ids is None:
880
+ # Create position_ids on the fly for batch generation
881
+ position_ids = attention_mask.long().cumsum(-1) - 1
882
+ position_ids.masked_fill_(attention_mask == 0, 1)
883
+ if past_key_values:
884
+ position_ids = position_ids[:, -1].unsqueeze(-1)
885
+
886
+ # If `inputs_embeds` are passed, we only want to use them in the 1st generation step
887
+ if inputs_embeds is not None and past_key_values is None:
888
+ model_inputs = {"inputs_embeds": inputs_embeds}
889
+ else:
890
+ model_inputs = {"input_ids": input_ids}
891
+
892
+ model_inputs.update(
893
+ {
894
+ "attention_mask": attention_mask,
895
+ "past_key_values": past_key_values,
896
+ "use_cache": kwargs.get("use_cache"),
897
+ "position_ids": position_ids,
898
+ }
899
+ )
900
+ return model_inputs
901
+
902
+ @staticmethod
903
+ def _reorder_cache(past_key_values, beam_idx):
904
+ reordered_past = ()
905
+ for layer_past in past_key_values:
906
+ reordered_past += (
907
+ tuple(
908
+ past_state.index_select(0, beam_idx.to(past_state.device))
909
+ for past_state in layer_past
910
+ ),
911
+ )
912
+ return reordered_past
913
+
914
+
915
+ StableLMEpochConfig.register_for_auto_class()
916
+ StableLMEpochForCausalLM.register_for_auto_class("AutoModelForCausalLM")