rhysjones commited on
Commit
e6379cd
1 Parent(s): 7e1743f

Upload folder using huggingface_hub

Browse files
added_tokens.json ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "\t\t": 50294,
3
+ "\t\t\t": 50293,
4
+ "\t\t\t\t": 50292,
5
+ "\t\t\t\t\t": 50291,
6
+ "\t\t\t\t\t\t": 50290,
7
+ "\t\t\t\t\t\t\t": 50289,
8
+ "\t\t\t\t\t\t\t\t": 50288,
9
+ "\t\t\t\t\t\t\t\t\t": 50287,
10
+ " ": 50286,
11
+ " ": 50285,
12
+ " ": 50284,
13
+ " ": 50283,
14
+ " ": 50282,
15
+ " ": 50281,
16
+ " ": 50280,
17
+ " ": 50279,
18
+ " ": 50278,
19
+ " ": 50277,
20
+ " ": 50276,
21
+ " ": 50275,
22
+ " ": 50274,
23
+ " ": 50273,
24
+ " ": 50272,
25
+ " ": 50271,
26
+ " ": 50270,
27
+ " ": 50269,
28
+ " ": 50268,
29
+ " ": 50267,
30
+ " ": 50266,
31
+ " ": 50265,
32
+ " ": 50264,
33
+ " ": 50263,
34
+ " ": 50262,
35
+ " ": 50261,
36
+ " ": 50260,
37
+ " ": 50259,
38
+ " ": 50258,
39
+ " ": 50257
40
+ }
config.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "microsoft/phi-2",
3
+ "activation_function": "gelu_new",
4
+ "architectures": [
5
+ "PhiForCausalLM"
6
+ ],
7
+ "attn_pdrop": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_phi.PhiConfig",
10
+ "AutoModelForCausalLM": "modeling_phi.PhiForCausalLM"
11
+ },
12
+ "embd_pdrop": 0.0,
13
+ "flash_attn": false,
14
+ "flash_rotary": false,
15
+ "fused_dense": false,
16
+ "img_processor": null,
17
+ "initializer_range": 0.02,
18
+ "layer_norm_epsilon": 1e-05,
19
+ "model_type": "phi-msft",
20
+ "n_embd": 2560,
21
+ "n_head": 32,
22
+ "n_head_kv": null,
23
+ "n_inner": null,
24
+ "n_layer": 32,
25
+ "n_positions": 2048,
26
+ "resid_pdrop": 0.1,
27
+ "rotary_dim": 32,
28
+ "tie_word_embeddings": false,
29
+ "torch_dtype": "bfloat16",
30
+ "transformers_version": "4.37.0.dev0",
31
+ "use_cache": false,
32
+ "vocab_size": 51200
33
+ }
configuration_phi.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Microsoft Corporation.
2
+ # Licensed under the MIT license.
3
+
4
+ import math
5
+ from typing import Optional
6
+
7
+ from transformers import PretrainedConfig
8
+
9
+
10
+ class PhiConfig(PretrainedConfig):
11
+ """Phi configuration."""
12
+
13
+ model_type = "phi-msft"
14
+ attribute_map = {
15
+ "max_position_embeddings": "n_positions",
16
+ "hidden_size": "n_embd",
17
+ "num_attention_heads": "n_head",
18
+ "num_hidden_layers": "n_layer",
19
+ }
20
+
21
+ def __init__(
22
+ self,
23
+ vocab_size: int = 50304,
24
+ n_positions: int = 2048,
25
+ n_embd: int = 1024,
26
+ n_layer: int = 20,
27
+ n_inner: Optional[int] = None,
28
+ n_head: int = 16,
29
+ n_head_kv: Optional[int] = None,
30
+ rotary_dim: Optional[int] = 32,
31
+ activation_function: Optional[str] = "gelu_new",
32
+ flash_attn: bool = False,
33
+ flash_rotary: bool = False,
34
+ fused_dense: bool = False,
35
+ attn_pdrop: float = 0.0,
36
+ embd_pdrop: float = 0.0,
37
+ resid_pdrop: float = 0.0,
38
+ layer_norm_epsilon: float = 1e-5,
39
+ initializer_range: float = 0.02,
40
+ tie_word_embeddings: bool = False,
41
+ pad_vocab_size_multiple: int = 64,
42
+ **kwargs
43
+ ) -> None:
44
+ self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
45
+ self.n_positions = n_positions
46
+ self.n_embd = n_embd
47
+ self.n_layer = n_layer
48
+ self.n_inner = n_inner
49
+ self.n_head = n_head
50
+ self.n_head_kv = n_head_kv
51
+ self.rotary_dim = min(rotary_dim, n_embd // n_head)
52
+ self.activation_function = activation_function
53
+ self.flash_attn = flash_attn
54
+ self.flash_rotary = flash_rotary
55
+ self.fused_dense = fused_dense
56
+ self.attn_pdrop = attn_pdrop
57
+ self.embd_pdrop = embd_pdrop
58
+ self.resid_pdrop = resid_pdrop
59
+ self.layer_norm_epsilon = layer_norm_epsilon
60
+ self.initializer_range = initializer_range
61
+
62
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
generation_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.37.0.dev0"
4
+ }
model-00001-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ff0f1d31399f121bf2f28d842c3d6f70033ca76a848ce486d4ef143eb9730a3b
3
+ size 4982468168
model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1eedaa8ceb5341ef141e5dc5e4fbf1ec2a3e25057bad636ffd2866887d498739
3
+ size 576934152
model.safetensors.index.json ADDED
@@ -0,0 +1,332 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 5559367680
4
+ },
5
+ "weight_map": {
6
+ "lm_head.linear.bias": "model-00002-of-00002.safetensors",
7
+ "lm_head.linear.weight": "model-00002-of-00002.safetensors",
8
+ "lm_head.ln.bias": "model-00002-of-00002.safetensors",
9
+ "lm_head.ln.weight": "model-00002-of-00002.safetensors",
10
+ "transformer.embd.wte.weight": "model-00001-of-00002.safetensors",
11
+ "transformer.h.0.ln.bias": "model-00001-of-00002.safetensors",
12
+ "transformer.h.0.ln.weight": "model-00001-of-00002.safetensors",
13
+ "transformer.h.0.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
14
+ "transformer.h.0.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
15
+ "transformer.h.0.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
16
+ "transformer.h.0.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
17
+ "transformer.h.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
18
+ "transformer.h.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
19
+ "transformer.h.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
20
+ "transformer.h.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
21
+ "transformer.h.1.ln.bias": "model-00001-of-00002.safetensors",
22
+ "transformer.h.1.ln.weight": "model-00001-of-00002.safetensors",
23
+ "transformer.h.1.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
24
+ "transformer.h.1.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
25
+ "transformer.h.1.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
26
+ "transformer.h.1.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
27
+ "transformer.h.1.mlp.fc1.bias": "model-00001-of-00002.safetensors",
28
+ "transformer.h.1.mlp.fc1.weight": "model-00001-of-00002.safetensors",
29
+ "transformer.h.1.mlp.fc2.bias": "model-00001-of-00002.safetensors",
30
+ "transformer.h.1.mlp.fc2.weight": "model-00001-of-00002.safetensors",
31
+ "transformer.h.10.ln.bias": "model-00001-of-00002.safetensors",
32
+ "transformer.h.10.ln.weight": "model-00001-of-00002.safetensors",
33
+ "transformer.h.10.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
34
+ "transformer.h.10.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
35
+ "transformer.h.10.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
36
+ "transformer.h.10.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
37
+ "transformer.h.10.mlp.fc1.bias": "model-00001-of-00002.safetensors",
38
+ "transformer.h.10.mlp.fc1.weight": "model-00001-of-00002.safetensors",
39
+ "transformer.h.10.mlp.fc2.bias": "model-00001-of-00002.safetensors",
40
+ "transformer.h.10.mlp.fc2.weight": "model-00001-of-00002.safetensors",
41
+ "transformer.h.11.ln.bias": "model-00001-of-00002.safetensors",
42
+ "transformer.h.11.ln.weight": "model-00001-of-00002.safetensors",
43
+ "transformer.h.11.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
44
+ "transformer.h.11.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
45
+ "transformer.h.11.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
46
+ "transformer.h.11.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
47
+ "transformer.h.11.mlp.fc1.bias": "model-00001-of-00002.safetensors",
48
+ "transformer.h.11.mlp.fc1.weight": "model-00001-of-00002.safetensors",
49
+ "transformer.h.11.mlp.fc2.bias": "model-00001-of-00002.safetensors",
50
+ "transformer.h.11.mlp.fc2.weight": "model-00001-of-00002.safetensors",
51
+ "transformer.h.12.ln.bias": "model-00001-of-00002.safetensors",
52
+ "transformer.h.12.ln.weight": "model-00001-of-00002.safetensors",
53
+ "transformer.h.12.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
54
+ "transformer.h.12.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
55
+ "transformer.h.12.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
56
+ "transformer.h.12.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
57
+ "transformer.h.12.mlp.fc1.bias": "model-00001-of-00002.safetensors",
58
+ "transformer.h.12.mlp.fc1.weight": "model-00001-of-00002.safetensors",
59
+ "transformer.h.12.mlp.fc2.bias": "model-00001-of-00002.safetensors",
60
+ "transformer.h.12.mlp.fc2.weight": "model-00001-of-00002.safetensors",
61
+ "transformer.h.13.ln.bias": "model-00001-of-00002.safetensors",
62
+ "transformer.h.13.ln.weight": "model-00001-of-00002.safetensors",
63
+ "transformer.h.13.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
64
+ "transformer.h.13.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
65
+ "transformer.h.13.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
66
+ "transformer.h.13.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
67
+ "transformer.h.13.mlp.fc1.bias": "model-00001-of-00002.safetensors",
68
+ "transformer.h.13.mlp.fc1.weight": "model-00001-of-00002.safetensors",
69
+ "transformer.h.13.mlp.fc2.bias": "model-00001-of-00002.safetensors",
70
+ "transformer.h.13.mlp.fc2.weight": "model-00001-of-00002.safetensors",
71
+ "transformer.h.14.ln.bias": "model-00001-of-00002.safetensors",
72
+ "transformer.h.14.ln.weight": "model-00001-of-00002.safetensors",
73
+ "transformer.h.14.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
74
+ "transformer.h.14.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
75
+ "transformer.h.14.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
76
+ "transformer.h.14.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
77
+ "transformer.h.14.mlp.fc1.bias": "model-00001-of-00002.safetensors",
78
+ "transformer.h.14.mlp.fc1.weight": "model-00001-of-00002.safetensors",
79
+ "transformer.h.14.mlp.fc2.bias": "model-00001-of-00002.safetensors",
80
+ "transformer.h.14.mlp.fc2.weight": "model-00001-of-00002.safetensors",
81
+ "transformer.h.15.ln.bias": "model-00001-of-00002.safetensors",
82
+ "transformer.h.15.ln.weight": "model-00001-of-00002.safetensors",
83
+ "transformer.h.15.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
84
+ "transformer.h.15.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
85
+ "transformer.h.15.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
86
+ "transformer.h.15.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
87
+ "transformer.h.15.mlp.fc1.bias": "model-00001-of-00002.safetensors",
88
+ "transformer.h.15.mlp.fc1.weight": "model-00001-of-00002.safetensors",
89
+ "transformer.h.15.mlp.fc2.bias": "model-00001-of-00002.safetensors",
90
+ "transformer.h.15.mlp.fc2.weight": "model-00001-of-00002.safetensors",
91
+ "transformer.h.16.ln.bias": "model-00001-of-00002.safetensors",
92
+ "transformer.h.16.ln.weight": "model-00001-of-00002.safetensors",
93
+ "transformer.h.16.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
94
+ "transformer.h.16.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
95
+ "transformer.h.16.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
96
+ "transformer.h.16.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
97
+ "transformer.h.16.mlp.fc1.bias": "model-00001-of-00002.safetensors",
98
+ "transformer.h.16.mlp.fc1.weight": "model-00001-of-00002.safetensors",
99
+ "transformer.h.16.mlp.fc2.bias": "model-00001-of-00002.safetensors",
100
+ "transformer.h.16.mlp.fc2.weight": "model-00001-of-00002.safetensors",
101
+ "transformer.h.17.ln.bias": "model-00001-of-00002.safetensors",
102
+ "transformer.h.17.ln.weight": "model-00001-of-00002.safetensors",
103
+ "transformer.h.17.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
104
+ "transformer.h.17.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
105
+ "transformer.h.17.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
106
+ "transformer.h.17.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
107
+ "transformer.h.17.mlp.fc1.bias": "model-00001-of-00002.safetensors",
108
+ "transformer.h.17.mlp.fc1.weight": "model-00001-of-00002.safetensors",
109
+ "transformer.h.17.mlp.fc2.bias": "model-00001-of-00002.safetensors",
110
+ "transformer.h.17.mlp.fc2.weight": "model-00001-of-00002.safetensors",
111
+ "transformer.h.18.ln.bias": "model-00001-of-00002.safetensors",
112
+ "transformer.h.18.ln.weight": "model-00001-of-00002.safetensors",
113
+ "transformer.h.18.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
114
+ "transformer.h.18.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
115
+ "transformer.h.18.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
116
+ "transformer.h.18.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
117
+ "transformer.h.18.mlp.fc1.bias": "model-00001-of-00002.safetensors",
118
+ "transformer.h.18.mlp.fc1.weight": "model-00001-of-00002.safetensors",
119
+ "transformer.h.18.mlp.fc2.bias": "model-00001-of-00002.safetensors",
120
+ "transformer.h.18.mlp.fc2.weight": "model-00001-of-00002.safetensors",
121
+ "transformer.h.19.ln.bias": "model-00001-of-00002.safetensors",
122
+ "transformer.h.19.ln.weight": "model-00001-of-00002.safetensors",
123
+ "transformer.h.19.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
124
+ "transformer.h.19.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
125
+ "transformer.h.19.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
126
+ "transformer.h.19.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
127
+ "transformer.h.19.mlp.fc1.bias": "model-00001-of-00002.safetensors",
128
+ "transformer.h.19.mlp.fc1.weight": "model-00001-of-00002.safetensors",
129
+ "transformer.h.19.mlp.fc2.bias": "model-00001-of-00002.safetensors",
130
+ "transformer.h.19.mlp.fc2.weight": "model-00001-of-00002.safetensors",
131
+ "transformer.h.2.ln.bias": "model-00001-of-00002.safetensors",
132
+ "transformer.h.2.ln.weight": "model-00001-of-00002.safetensors",
133
+ "transformer.h.2.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
134
+ "transformer.h.2.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
135
+ "transformer.h.2.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
136
+ "transformer.h.2.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
137
+ "transformer.h.2.mlp.fc1.bias": "model-00001-of-00002.safetensors",
138
+ "transformer.h.2.mlp.fc1.weight": "model-00001-of-00002.safetensors",
139
+ "transformer.h.2.mlp.fc2.bias": "model-00001-of-00002.safetensors",
140
+ "transformer.h.2.mlp.fc2.weight": "model-00001-of-00002.safetensors",
141
+ "transformer.h.20.ln.bias": "model-00001-of-00002.safetensors",
142
+ "transformer.h.20.ln.weight": "model-00001-of-00002.safetensors",
143
+ "transformer.h.20.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
144
+ "transformer.h.20.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
145
+ "transformer.h.20.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
146
+ "transformer.h.20.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
147
+ "transformer.h.20.mlp.fc1.bias": "model-00001-of-00002.safetensors",
148
+ "transformer.h.20.mlp.fc1.weight": "model-00001-of-00002.safetensors",
149
+ "transformer.h.20.mlp.fc2.bias": "model-00001-of-00002.safetensors",
150
+ "transformer.h.20.mlp.fc2.weight": "model-00001-of-00002.safetensors",
151
+ "transformer.h.21.ln.bias": "model-00001-of-00002.safetensors",
152
+ "transformer.h.21.ln.weight": "model-00001-of-00002.safetensors",
153
+ "transformer.h.21.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
154
+ "transformer.h.21.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
155
+ "transformer.h.21.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
156
+ "transformer.h.21.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
157
+ "transformer.h.21.mlp.fc1.bias": "model-00001-of-00002.safetensors",
158
+ "transformer.h.21.mlp.fc1.weight": "model-00001-of-00002.safetensors",
159
+ "transformer.h.21.mlp.fc2.bias": "model-00001-of-00002.safetensors",
160
+ "transformer.h.21.mlp.fc2.weight": "model-00001-of-00002.safetensors",
161
+ "transformer.h.22.ln.bias": "model-00001-of-00002.safetensors",
162
+ "transformer.h.22.ln.weight": "model-00001-of-00002.safetensors",
163
+ "transformer.h.22.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
164
+ "transformer.h.22.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
165
+ "transformer.h.22.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
166
+ "transformer.h.22.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
167
+ "transformer.h.22.mlp.fc1.bias": "model-00001-of-00002.safetensors",
168
+ "transformer.h.22.mlp.fc1.weight": "model-00001-of-00002.safetensors",
169
+ "transformer.h.22.mlp.fc2.bias": "model-00001-of-00002.safetensors",
170
+ "transformer.h.22.mlp.fc2.weight": "model-00001-of-00002.safetensors",
171
+ "transformer.h.23.ln.bias": "model-00001-of-00002.safetensors",
172
+ "transformer.h.23.ln.weight": "model-00001-of-00002.safetensors",
173
+ "transformer.h.23.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
174
+ "transformer.h.23.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
175
+ "transformer.h.23.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
176
+ "transformer.h.23.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
177
+ "transformer.h.23.mlp.fc1.bias": "model-00001-of-00002.safetensors",
178
+ "transformer.h.23.mlp.fc1.weight": "model-00001-of-00002.safetensors",
179
+ "transformer.h.23.mlp.fc2.bias": "model-00001-of-00002.safetensors",
180
+ "transformer.h.23.mlp.fc2.weight": "model-00001-of-00002.safetensors",
181
+ "transformer.h.24.ln.bias": "model-00001-of-00002.safetensors",
182
+ "transformer.h.24.ln.weight": "model-00001-of-00002.safetensors",
183
+ "transformer.h.24.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
184
+ "transformer.h.24.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
185
+ "transformer.h.24.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
186
+ "transformer.h.24.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
187
+ "transformer.h.24.mlp.fc1.bias": "model-00001-of-00002.safetensors",
188
+ "transformer.h.24.mlp.fc1.weight": "model-00001-of-00002.safetensors",
189
+ "transformer.h.24.mlp.fc2.bias": "model-00001-of-00002.safetensors",
190
+ "transformer.h.24.mlp.fc2.weight": "model-00001-of-00002.safetensors",
191
+ "transformer.h.25.ln.bias": "model-00001-of-00002.safetensors",
192
+ "transformer.h.25.ln.weight": "model-00001-of-00002.safetensors",
193
+ "transformer.h.25.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
194
+ "transformer.h.25.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
195
+ "transformer.h.25.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
196
+ "transformer.h.25.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
197
+ "transformer.h.25.mlp.fc1.bias": "model-00001-of-00002.safetensors",
198
+ "transformer.h.25.mlp.fc1.weight": "model-00001-of-00002.safetensors",
199
+ "transformer.h.25.mlp.fc2.bias": "model-00001-of-00002.safetensors",
200
+ "transformer.h.25.mlp.fc2.weight": "model-00001-of-00002.safetensors",
201
+ "transformer.h.26.ln.bias": "model-00001-of-00002.safetensors",
202
+ "transformer.h.26.ln.weight": "model-00001-of-00002.safetensors",
203
+ "transformer.h.26.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
204
+ "transformer.h.26.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
205
+ "transformer.h.26.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
206
+ "transformer.h.26.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
207
+ "transformer.h.26.mlp.fc1.bias": "model-00001-of-00002.safetensors",
208
+ "transformer.h.26.mlp.fc1.weight": "model-00001-of-00002.safetensors",
209
+ "transformer.h.26.mlp.fc2.bias": "model-00001-of-00002.safetensors",
210
+ "transformer.h.26.mlp.fc2.weight": "model-00001-of-00002.safetensors",
211
+ "transformer.h.27.ln.bias": "model-00001-of-00002.safetensors",
212
+ "transformer.h.27.ln.weight": "model-00001-of-00002.safetensors",
213
+ "transformer.h.27.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
214
+ "transformer.h.27.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
215
+ "transformer.h.27.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
216
+ "transformer.h.27.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
217
+ "transformer.h.27.mlp.fc1.bias": "model-00001-of-00002.safetensors",
218
+ "transformer.h.27.mlp.fc1.weight": "model-00001-of-00002.safetensors",
219
+ "transformer.h.27.mlp.fc2.bias": "model-00001-of-00002.safetensors",
220
+ "transformer.h.27.mlp.fc2.weight": "model-00001-of-00002.safetensors",
221
+ "transformer.h.28.ln.bias": "model-00001-of-00002.safetensors",
222
+ "transformer.h.28.ln.weight": "model-00001-of-00002.safetensors",
223
+ "transformer.h.28.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
224
+ "transformer.h.28.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
225
+ "transformer.h.28.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
226
+ "transformer.h.28.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
227
+ "transformer.h.28.mlp.fc1.bias": "model-00001-of-00002.safetensors",
228
+ "transformer.h.28.mlp.fc1.weight": "model-00001-of-00002.safetensors",
229
+ "transformer.h.28.mlp.fc2.bias": "model-00001-of-00002.safetensors",
230
+ "transformer.h.28.mlp.fc2.weight": "model-00001-of-00002.safetensors",
231
+ "transformer.h.29.ln.bias": "model-00001-of-00002.safetensors",
232
+ "transformer.h.29.ln.weight": "model-00001-of-00002.safetensors",
233
+ "transformer.h.29.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
234
+ "transformer.h.29.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
235
+ "transformer.h.29.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
236
+ "transformer.h.29.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
237
+ "transformer.h.29.mlp.fc1.bias": "model-00001-of-00002.safetensors",
238
+ "transformer.h.29.mlp.fc1.weight": "model-00001-of-00002.safetensors",
239
+ "transformer.h.29.mlp.fc2.bias": "model-00001-of-00002.safetensors",
240
+ "transformer.h.29.mlp.fc2.weight": "model-00001-of-00002.safetensors",
241
+ "transformer.h.3.ln.bias": "model-00001-of-00002.safetensors",
242
+ "transformer.h.3.ln.weight": "model-00001-of-00002.safetensors",
243
+ "transformer.h.3.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
244
+ "transformer.h.3.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
245
+ "transformer.h.3.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
246
+ "transformer.h.3.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
247
+ "transformer.h.3.mlp.fc1.bias": "model-00001-of-00002.safetensors",
248
+ "transformer.h.3.mlp.fc1.weight": "model-00001-of-00002.safetensors",
249
+ "transformer.h.3.mlp.fc2.bias": "model-00001-of-00002.safetensors",
250
+ "transformer.h.3.mlp.fc2.weight": "model-00001-of-00002.safetensors",
251
+ "transformer.h.30.ln.bias": "model-00001-of-00002.safetensors",
252
+ "transformer.h.30.ln.weight": "model-00001-of-00002.safetensors",
253
+ "transformer.h.30.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
254
+ "transformer.h.30.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
255
+ "transformer.h.30.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
256
+ "transformer.h.30.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
257
+ "transformer.h.30.mlp.fc1.bias": "model-00002-of-00002.safetensors",
258
+ "transformer.h.30.mlp.fc1.weight": "model-00002-of-00002.safetensors",
259
+ "transformer.h.30.mlp.fc2.bias": "model-00002-of-00002.safetensors",
260
+ "transformer.h.30.mlp.fc2.weight": "model-00002-of-00002.safetensors",
261
+ "transformer.h.31.ln.bias": "model-00002-of-00002.safetensors",
262
+ "transformer.h.31.ln.weight": "model-00002-of-00002.safetensors",
263
+ "transformer.h.31.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
264
+ "transformer.h.31.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
265
+ "transformer.h.31.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
266
+ "transformer.h.31.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
267
+ "transformer.h.31.mlp.fc1.bias": "model-00002-of-00002.safetensors",
268
+ "transformer.h.31.mlp.fc1.weight": "model-00002-of-00002.safetensors",
269
+ "transformer.h.31.mlp.fc2.bias": "model-00002-of-00002.safetensors",
270
+ "transformer.h.31.mlp.fc2.weight": "model-00002-of-00002.safetensors",
271
+ "transformer.h.4.ln.bias": "model-00001-of-00002.safetensors",
272
+ "transformer.h.4.ln.weight": "model-00001-of-00002.safetensors",
273
+ "transformer.h.4.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
274
+ "transformer.h.4.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
275
+ "transformer.h.4.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
276
+ "transformer.h.4.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
277
+ "transformer.h.4.mlp.fc1.bias": "model-00001-of-00002.safetensors",
278
+ "transformer.h.4.mlp.fc1.weight": "model-00001-of-00002.safetensors",
279
+ "transformer.h.4.mlp.fc2.bias": "model-00001-of-00002.safetensors",
280
+ "transformer.h.4.mlp.fc2.weight": "model-00001-of-00002.safetensors",
281
+ "transformer.h.5.ln.bias": "model-00001-of-00002.safetensors",
282
+ "transformer.h.5.ln.weight": "model-00001-of-00002.safetensors",
283
+ "transformer.h.5.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
284
+ "transformer.h.5.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
285
+ "transformer.h.5.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
286
+ "transformer.h.5.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
287
+ "transformer.h.5.mlp.fc1.bias": "model-00001-of-00002.safetensors",
288
+ "transformer.h.5.mlp.fc1.weight": "model-00001-of-00002.safetensors",
289
+ "transformer.h.5.mlp.fc2.bias": "model-00001-of-00002.safetensors",
290
+ "transformer.h.5.mlp.fc2.weight": "model-00001-of-00002.safetensors",
291
+ "transformer.h.6.ln.bias": "model-00001-of-00002.safetensors",
292
+ "transformer.h.6.ln.weight": "model-00001-of-00002.safetensors",
293
+ "transformer.h.6.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
294
+ "transformer.h.6.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
295
+ "transformer.h.6.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
296
+ "transformer.h.6.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
297
+ "transformer.h.6.mlp.fc1.bias": "model-00001-of-00002.safetensors",
298
+ "transformer.h.6.mlp.fc1.weight": "model-00001-of-00002.safetensors",
299
+ "transformer.h.6.mlp.fc2.bias": "model-00001-of-00002.safetensors",
300
+ "transformer.h.6.mlp.fc2.weight": "model-00001-of-00002.safetensors",
301
+ "transformer.h.7.ln.bias": "model-00001-of-00002.safetensors",
302
+ "transformer.h.7.ln.weight": "model-00001-of-00002.safetensors",
303
+ "transformer.h.7.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
304
+ "transformer.h.7.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
305
+ "transformer.h.7.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
306
+ "transformer.h.7.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
307
+ "transformer.h.7.mlp.fc1.bias": "model-00001-of-00002.safetensors",
308
+ "transformer.h.7.mlp.fc1.weight": "model-00001-of-00002.safetensors",
309
+ "transformer.h.7.mlp.fc2.bias": "model-00001-of-00002.safetensors",
310
+ "transformer.h.7.mlp.fc2.weight": "model-00001-of-00002.safetensors",
311
+ "transformer.h.8.ln.bias": "model-00001-of-00002.safetensors",
312
+ "transformer.h.8.ln.weight": "model-00001-of-00002.safetensors",
313
+ "transformer.h.8.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
314
+ "transformer.h.8.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
315
+ "transformer.h.8.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
316
+ "transformer.h.8.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
317
+ "transformer.h.8.mlp.fc1.bias": "model-00001-of-00002.safetensors",
318
+ "transformer.h.8.mlp.fc1.weight": "model-00001-of-00002.safetensors",
319
+ "transformer.h.8.mlp.fc2.bias": "model-00001-of-00002.safetensors",
320
+ "transformer.h.8.mlp.fc2.weight": "model-00001-of-00002.safetensors",
321
+ "transformer.h.9.ln.bias": "model-00001-of-00002.safetensors",
322
+ "transformer.h.9.ln.weight": "model-00001-of-00002.safetensors",
323
+ "transformer.h.9.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
324
+ "transformer.h.9.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
325
+ "transformer.h.9.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
326
+ "transformer.h.9.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
327
+ "transformer.h.9.mlp.fc1.bias": "model-00001-of-00002.safetensors",
328
+ "transformer.h.9.mlp.fc1.weight": "model-00001-of-00002.safetensors",
329
+ "transformer.h.9.mlp.fc2.bias": "model-00001-of-00002.safetensors",
330
+ "transformer.h.9.mlp.fc2.weight": "model-00001-of-00002.safetensors"
331
+ }
332
+ }
modeling_phi.py ADDED
@@ -0,0 +1,960 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Microsoft Corporation.
2
+ # Licensed under the MIT license.
3
+ #
4
+ # Copyright (c) 2022, Tri Dao, [email protected].
5
+ # Licensed under the BSD 3-Clause License.
6
+
7
+ from __future__ import annotations
8
+
9
+ import math
10
+ from dataclasses import dataclass, field
11
+ from typing import Any, Dict, Optional, Tuple, Union
12
+
13
+ import torch
14
+ import torch.nn as nn
15
+ from einops import rearrange, repeat
16
+ from transformers import PretrainedConfig, PreTrainedModel
17
+ from transformers.activations import ACT2FN
18
+ from transformers.modeling_outputs import CausalLMOutputWithPast
19
+
20
+ from .configuration_phi import PhiConfig
21
+
22
+ try:
23
+ from flash_attn.bert_padding import pad_input, unpad_input
24
+ from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
25
+ from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
26
+ from flash_attn.ops.fused_dense import FusedDense
27
+ except:
28
+ pad_input, unpad_input = None, None
29
+ FlashRotaryEmbedding = None
30
+ FlashSelfAttention, FlashCrossAttention = None, None
31
+ FusedDense = None
32
+
33
+
34
+ @dataclass
35
+ class InferenceParams:
36
+ """Inference parameters passed to model to efficiently calculate
37
+ and store context during inference.
38
+
39
+ Reference:
40
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
41
+
42
+ Args:
43
+ max_seqlen: Maximum sequence length.
44
+ max_batch_size: Maximum batch size.
45
+ seqlen_offset: Sequence length offset.
46
+ batch_size_offset: Batch size offset.
47
+ key_value_memory_dict: Key value memory dictionary.
48
+ lengths_per_sample: Lengths per sample.
49
+
50
+ """
51
+
52
+ max_seqlen: int = field(metadata={"help": "Maximum sequence length."})
53
+
54
+ max_batch_size: int = field(metadata={"help": "Maximum batch size."})
55
+
56
+ seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
57
+
58
+ batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
59
+
60
+ key_value_memory_dict: Dict[str, Any] = field(
61
+ default_factory=dict, metadata={"help": "Key value memory dictionary."}
62
+ )
63
+
64
+ lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
65
+
66
+
67
+ class Embedding(nn.Module):
68
+ """Token embedding with dropout."""
69
+
70
+ def __init__(self, config: PretrainedConfig) -> None:
71
+ super().__init__()
72
+
73
+ self.wte = nn.Embedding(config.vocab_size, config.n_embd)
74
+ self.drop = nn.Dropout(config.embd_pdrop)
75
+
76
+ def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
77
+ input_shape = input_ids.size()
78
+ input_ids = input_ids.view(-1, input_shape[-1])
79
+
80
+ hidden_states = self.wte(input_ids)
81
+ hidden_states = self.drop(hidden_states)
82
+
83
+ return hidden_states
84
+
85
+
86
+ def _apply_rotary_emb(
87
+ x: torch.FloatTensor,
88
+ cos: torch.FloatTensor,
89
+ sin: torch.FloatTensor,
90
+ ) -> torch.FloatTensor:
91
+ _, seqlen, _, _ = x.shape
92
+ _, rotary_dim = cos.shape
93
+ rotary_dim *= 2
94
+
95
+ x_rot = x[:, :, :, :rotary_dim]
96
+ x_pass = x[:, :, :, rotary_dim:]
97
+
98
+ x1, x2 = x_rot.chunk(2, dim=-1)
99
+ c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
100
+ x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]
101
+
102
+ x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)
103
+
104
+ return torch.cat([x_rot, x_pass], axis=-1)
105
+
106
+
107
+ def _apply_rotary_emb_kv(
108
+ kv: torch.FloatTensor,
109
+ cos: torch.FloatTensor,
110
+ sin: torch.FloatTensor,
111
+ cos_k: Optional[torch.FloatTensor] = None,
112
+ sin_k: Optional[torch.FloatTensor] = None,
113
+ ) -> torch.FloatTensor:
114
+ _, seqlen, _, _, _ = kv.shape
115
+ _, rotary_dim = cos.shape
116
+ rotary_dim *= 2
117
+
118
+ k_rot = kv[:, :, 0, :, :rotary_dim]
119
+ k_pass = kv[:, :, 0, :, rotary_dim:]
120
+
121
+ k1, k2 = k_rot.chunk(2, dim=-1)
122
+ c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
123
+ k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]
124
+
125
+ k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)
126
+
127
+ return torch.cat(
128
+ [
129
+ torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
130
+ kv[:, :, 1:2, :, :],
131
+ ],
132
+ axis=2,
133
+ )
134
+
135
+
136
+ def _apply_rotary_emb_qkv(
137
+ qkv: torch.FloatTensor,
138
+ cos: torch.FloatTensor,
139
+ sin: torch.FloatTensor,
140
+ cos_k: Optional[torch.FloatTensor] = None,
141
+ sin_k: Optional[torch.FloatTensor] = None,
142
+ ) -> torch.FloatTensor:
143
+ _, seqlen, _, _, _ = qkv.shape
144
+ _, rotary_dim = cos.shape
145
+ rotary_dim *= 2
146
+
147
+ q_rot = qkv[:, :, 0, :, :rotary_dim]
148
+ q_pass = qkv[:, :, 0, :, rotary_dim:]
149
+
150
+ k_rot = qkv[:, :, 1, :, :rotary_dim]
151
+ k_pass = qkv[:, :, 1, :, rotary_dim:]
152
+
153
+ q1, q2 = q_rot.chunk(2, dim=-1)
154
+ k1, k2 = k_rot.chunk(2, dim=-1)
155
+ c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
156
+ q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
157
+
158
+ q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
159
+ k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
160
+
161
+ return torch.cat(
162
+ [
163
+ torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
164
+ torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
165
+ qkv[:, :, 2:3, :, :],
166
+ ],
167
+ axis=2,
168
+ )
169
+
170
+
171
+ class RotaryEmbedding(nn.Module):
172
+ """Rotary positional embedding (RoPE).
173
+
174
+ Reference:
175
+ RoFormer: Enhanced Transformer with Rotary Position Embedding.
176
+ https://arxiv.org/pdf/2104.09864.pdf.
177
+
178
+ """
179
+
180
+ def __init__(
181
+ self,
182
+ dim: int,
183
+ base: int = 10000,
184
+ scale_base: Optional[float] = None,
185
+ pos_idx_in_fp32: bool = True,
186
+ max_position_embeddings: int = 2048,
187
+ device: Optional[str] = None,
188
+ **kwargs,
189
+ ) -> None:
190
+ super().__init__()
191
+
192
+ if scale_base is not None:
193
+ raise NotImplementedError
194
+
195
+ self.dim = dim
196
+ self.base = float(base)
197
+ self.scale_base = scale_base
198
+ self.pos_idx_in_fp32 = pos_idx_in_fp32
199
+ self.max_position_embeddings = max_position_embeddings
200
+ self.device = device
201
+
202
+ # Generate and save the inverse frequency buffer (non-trainable)
203
+ inv_freq = self._compute_inv_freq(device)
204
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
205
+
206
+ # Generate and save the scale buffer (non-trainable)
207
+ scale = (
208
+ (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
209
+ if scale_base is not None
210
+ else None
211
+ )
212
+ self.register_buffer("scale", scale, persistent=False)
213
+
214
+ # Initialize cached attributes since ONNX can't rely on dynamic initialization
215
+ self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)
216
+
217
+ def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
218
+ return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
219
+
220
+ def _update_cos_sin_cache(
221
+ self,
222
+ seqlen: int,
223
+ device: Optional[str] = None,
224
+ dtype: Optional[torch.dtype] = None,
225
+ ) -> None:
226
+ self._seq_len_cached = seqlen
227
+
228
+ # fp32 is preferred since the output of `torch.arange` can be quite large
229
+ # and bf16 would lose a lot of precision
230
+ if self.pos_idx_in_fp32:
231
+ t = torch.arange(seqlen, device=device, dtype=torch.float32)
232
+ if self.inv_freq.dtype != torch.float32:
233
+ inv_freq = self._compute_inv_freq(device=device)
234
+ else:
235
+ inv_freq = self.inv_freq
236
+ else:
237
+ t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
238
+ inv_freq = self.inv_freq
239
+
240
+ # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP
241
+ freqs = torch.outer(t, inv_freq)
242
+ if self.scale is None:
243
+ self._cos_cached = torch.cos(freqs).to(dtype)
244
+ self._sin_cached = torch.sin(freqs).to(dtype)
245
+ else:
246
+ power = (
247
+ torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
248
+ ) / self.scale_base
249
+ scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
250
+
251
+ # Force the scale multiplication to happen in fp32
252
+ self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
253
+ self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
254
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
255
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
256
+
257
+ def forward(
258
+ self,
259
+ qkv: torch.Tensor,
260
+ kv: Optional[torch.Tensor] = None,
261
+ seqlen_offset: int = 0,
262
+ **kwargs,
263
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
264
+ if (
265
+ self._seq_len_cached < qkv.shape[1] + seqlen_offset
266
+ or self._cos_cached.device != qkv.device
267
+ or self._cos_cached.dtype != qkv.dtype
268
+ or (self.training and self._cos_cached.is_inference())
269
+ ):
270
+ self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
271
+
272
+ if kv is None:
273
+ return _apply_rotary_emb_qkv(
274
+ qkv,
275
+ self._cos_cached[seqlen_offset:],
276
+ self._sin_cached[seqlen_offset:],
277
+ )
278
+ else:
279
+ q = _apply_rotary_emb(
280
+ qkv,
281
+ self._cos_cached[seqlen_offset:],
282
+ self._sin_cached[seqlen_offset:],
283
+ )
284
+ kv = _apply_rotary_emb_kv(
285
+ kv,
286
+ self._cos_cached[seqlen_offset:],
287
+ self._sin_cached[seqlen_offset:],
288
+ )
289
+
290
+ return q, kv
291
+
292
+
293
+ class MLP(nn.Module):
294
+ """Multi-Layer Perceptron.
295
+
296
+ Reference:
297
+ Attention Is All You Need.
298
+ https://arxiv.org/pdf/1706.03762.pdf.
299
+
300
+ """
301
+
302
+ def __init__(
303
+ self,
304
+ config: PretrainedConfig,
305
+ n_inner: Optional[int] = None,
306
+ act_fn: Optional[str] = None,
307
+ ) -> None:
308
+ super().__init__()
309
+
310
+ act_fn = config.activation_function if act_fn is None else act_fn
311
+
312
+ n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
313
+ n_inner = n_inner if n_inner is not None else 4 * config.n_embd
314
+
315
+ self.fc1 = nn.Linear(config.n_embd, n_inner)
316
+ self.fc2 = nn.Linear(n_inner, config.n_embd)
317
+ self.act = ACT2FN[act_fn]
318
+
319
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
320
+ hidden_states = self.fc1(hidden_states)
321
+ hidden_states = self.act(hidden_states)
322
+ hidden_states = self.fc2(hidden_states)
323
+
324
+ return hidden_states
325
+
326
+
327
+ class SelfAttention(nn.Module):
328
+ """Self-attention layer (compatible with PyTorch).
329
+
330
+ Reference:
331
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
332
+
333
+ """
334
+
335
+ def __init__(
336
+ self,
337
+ causal: bool = True,
338
+ softmax_scale: Optional[float] = None,
339
+ attention_dropout: float = 0.0,
340
+ ) -> None:
341
+ super().__init__()
342
+
343
+ self.causal = causal
344
+ self.softmax_scale = softmax_scale
345
+ self.drop = nn.Dropout(attention_dropout)
346
+
347
+ @torch.autocast("cpu", enabled=False)
348
+ @torch.autocast("cuda", enabled=False)
349
+ def forward(
350
+ self,
351
+ qkv: torch.FloatTensor,
352
+ causal: bool = None,
353
+ key_padding_mask: Optional[torch.BoolTensor] = None,
354
+ **kwargs,
355
+ ) -> torch.FloatTensor:
356
+ batch_size, seqlen = qkv.shape[0], qkv.shape[1]
357
+ q, k, v = qkv.unbind(dim=2)
358
+
359
+ q = q.to(torch.float32)
360
+ k = k.to(torch.float32)
361
+
362
+ causal = self.causal if causal is None else causal
363
+ softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
364
+
365
+ # Autocast is manually disabled to avoid `torch.einsum` performing the operation
366
+ # using float16, which might lead to overflow
367
+ scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
368
+
369
+ if key_padding_mask is not None:
370
+ padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
371
+ padding_mask.masked_fill_(key_padding_mask, 0.0)
372
+
373
+ scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
374
+
375
+ if causal:
376
+ causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
377
+ scores = scores + causal_mask.to(dtype=scores.dtype)
378
+
379
+ attention = torch.softmax(scores, dim=-1).to(v.dtype)
380
+ attention = self.drop(attention)
381
+
382
+ output = torch.einsum("bhts,bshd->bthd", attention, v)
383
+
384
+ return output
385
+
386
+
387
+ class CrossAttention(nn.Module):
388
+ """Cross-attention layer (compatible with PyTorch).
389
+
390
+ Reference:
391
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
392
+
393
+ """
394
+
395
+ def __init__(
396
+ self,
397
+ causal: bool = True,
398
+ softmax_scale: Optional[float] = None,
399
+ attention_dropout: float = 0.0,
400
+ ) -> None:
401
+ super().__init__()
402
+
403
+ self.causal = causal
404
+ self.softmax_scale = softmax_scale
405
+ self.drop = nn.Dropout(attention_dropout)
406
+
407
+ @torch.autocast("cpu", enabled=False)
408
+ @torch.autocast("cuda", enabled=False)
409
+ def forward(
410
+ self,
411
+ q: torch.FloatTensor,
412
+ kv: torch.FloatTensor,
413
+ causal: bool = None,
414
+ key_padding_mask: Optional[torch.BoolTensor] = None,
415
+ **kwargs,
416
+ ) -> torch.FloatTensor:
417
+ batch_size, seqlen_q = q.shape[0], q.shape[1]
418
+ seqlen_k = kv.shape[1]
419
+
420
+ if kv.shape[3] != q.shape[2]:
421
+ kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
422
+ k, v = kv.unbind(dim=2)
423
+
424
+ q = q.to(torch.float32)
425
+ k = k.to(torch.float32)
426
+
427
+ causal = self.causal if causal is None else causal
428
+ softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
429
+
430
+ # Autocast is manually disabled to avoid `torch.einsum` performing the operation
431
+ # using float16, which might lead to overflow
432
+ scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
433
+
434
+ if key_padding_mask is not None:
435
+ padding_mask = torch.full(
436
+ (batch_size, seqlen_k),
437
+ -10000.0,
438
+ dtype=scores.dtype,
439
+ device=scores.device,
440
+ )
441
+ padding_mask.masked_fill_(key_padding_mask, 0.0)
442
+
443
+ scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
444
+
445
+ if causal:
446
+ rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")
447
+ cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
448
+ causal_mask = cols > rows + seqlen_k - seqlen_q
449
+
450
+ scores = scores.masked_fill(causal_mask, -10000.0)
451
+
452
+ attention = torch.softmax(scores, dim=-1).to(v.dtype)
453
+ attention = self.drop(attention)
454
+
455
+ output = torch.einsum("bhts,bshd->bthd", attention, v)
456
+
457
+ return output
458
+
459
+
460
+ def _find_mha_dims(
461
+ config: PretrainedConfig,
462
+ n_head: Optional[int] = None,
463
+ n_head_kv: Optional[int] = None,
464
+ head_dim: Optional[int] = None,
465
+ ) -> Tuple[int, int]:
466
+ if n_head is None and head_dim is None:
467
+ head_dim = config.n_embd // config.n_head
468
+ n_head = config.n_head
469
+ elif n_head is None or head_dim is None:
470
+ raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
471
+
472
+ if n_head_kv is None:
473
+ n_head_kv = getattr(config, "n_head_kv", None) or n_head
474
+
475
+ return n_head, n_head_kv, head_dim
476
+
477
+
478
+ def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
479
+ num_heads, head_dim = kv.shape[-2:]
480
+
481
+ if layer_idx not in inference_params.key_value_memory_dict:
482
+ inference_params.key_value_memory_dict[layer_idx] = torch.empty(
483
+ inference_params.max_batch_size,
484
+ inference_params.max_seqlen,
485
+ 2,
486
+ num_heads,
487
+ head_dim,
488
+ dtype=kv.dtype,
489
+ device=kv.device,
490
+ )
491
+
492
+ batch_start = inference_params.batch_size_offset
493
+ batch_end = batch_start + kv.shape[0]
494
+
495
+ sequence_start = inference_params.seqlen_offset
496
+ sequence_end = sequence_start + kv.shape[1]
497
+
498
+ # When the current sequence length is equal to or larger than the maximum sequence length,
499
+ # we need to concatenate the current `kv` with the cached `kv` to expand its length
500
+ if sequence_end >= inference_params.max_seqlen:
501
+ inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)
502
+
503
+ inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv
504
+ kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]
505
+
506
+ return kv
507
+
508
+
509
+ class MHA(nn.Module):
510
+ """Multi-head attention layer."""
511
+
512
+ def __init__(
513
+ self,
514
+ config: PretrainedConfig,
515
+ dtype: Optional[torch.dtype] = None,
516
+ device: Optional[str] = None,
517
+ rotary_dim: Optional[int] = None,
518
+ rotary_base: float = 10000.0,
519
+ rotary_scale_base: Optional[float] = None,
520
+ n_head: Optional[int] = None,
521
+ n_head_kv: Optional[int] = None,
522
+ head_dim: Optional[int] = None,
523
+ bias: bool = True,
524
+ causal: bool = True,
525
+ softmax_scale: Optional[float] = None,
526
+ layer_idx: Optional[int] = None,
527
+ return_residual: bool = False,
528
+ checkpointing: bool = False,
529
+ ) -> None:
530
+ super().__init__()
531
+
532
+ # Rotary embedding
533
+ self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
534
+ if self.rotary_dim > 0:
535
+ rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
536
+ if rotary_cls is None:
537
+ rotary_cls = RotaryEmbedding
538
+
539
+ rotary_kwargs = {}
540
+ if rotary_cls is RotaryEmbedding:
541
+ rotary_kwargs["max_position_embeddings"] = config.n_positions
542
+
543
+ self.rotary_emb = rotary_cls(
544
+ self.rotary_dim,
545
+ base=rotary_base,
546
+ scale_base=rotary_scale_base,
547
+ device=device,
548
+ **rotary_kwargs,
549
+ )
550
+
551
+ # MLP
552
+ self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
553
+ config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim
554
+ )
555
+ op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
556
+ hidden_size = config.n_embd
557
+
558
+ linear_cls = FusedDense if config.fused_dense else nn.Linear
559
+ if linear_cls is None:
560
+ linear_cls = nn.Linear
561
+
562
+ self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)
563
+ self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
564
+
565
+ # Attention
566
+ attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention
567
+ if attn_cls is None:
568
+ attn_cls = SelfAttention
569
+
570
+ cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention
571
+ if cross_attn_cls is None:
572
+ cross_attn_cls = CrossAttention
573
+
574
+ self.inner_attn = attn_cls(
575
+ causal=causal,
576
+ softmax_scale=softmax_scale,
577
+ attention_dropout=config.attn_pdrop,
578
+ )
579
+ self.inner_cross_attn = cross_attn_cls(
580
+ causal=causal,
581
+ softmax_scale=softmax_scale,
582
+ attention_dropout=config.attn_pdrop,
583
+ )
584
+
585
+ self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention
586
+ self.layer_idx = layer_idx
587
+ self.return_residual = return_residual
588
+ self.checkpointing = checkpointing
589
+
590
+ def _forward_self_attn(
591
+ self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
592
+ ) -> torch.FloatTensor:
593
+ qkv = self.Wqkv(x)
594
+ qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
595
+
596
+ if self.rotary_dim > 0:
597
+ qkv = self.rotary_emb(qkv)
598
+
599
+ if self.flash_attn:
600
+ batch_size, seqlen = qkv.shape[0], qkv.shape[1]
601
+
602
+ cu_seqlens, max_seqlen = None, None
603
+ if key_padding_mask is not None:
604
+ # If `key_padding_mask` is supplied, we need to unpad the input and retrieve
605
+ # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`
606
+ qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)
607
+
608
+ if self.checkpointing:
609
+ attn_output = torch.utils.checkpoint.checkpoint(
610
+ self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
611
+ )
612
+ else:
613
+ attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
614
+
615
+ # If `key_padding_mask` is supplied, we need to pad the output back to the original shape
616
+ return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
617
+
618
+ if self.checkpointing:
619
+ return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
620
+
621
+ return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
622
+
623
+ def _forward_cross_attn(
624
+ self,
625
+ x: torch.FloatTensor,
626
+ past_key_values: Optional[InferenceParams],
627
+ key_padding_mask: Optional[torch.BoolTensor],
628
+ ) -> torch.FloatTensor:
629
+ batch_size = x.shape[0]
630
+
631
+ qkv = self.Wqkv(x)
632
+
633
+ q = qkv[..., : self.n_head * self.head_dim]
634
+ q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
635
+
636
+ kv = qkv[..., self.n_head * self.head_dim :]
637
+ kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
638
+
639
+ seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0
640
+ causal = None if seqlen_offset == 0 else False
641
+ if self.rotary_dim > 0:
642
+ q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
643
+
644
+ if past_key_values is not None:
645
+ kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
646
+
647
+ if self.flash_attn:
648
+ batch_size, seqlen_q = q.shape[0], q.shape[1]
649
+ seqlen_k = kv.shape[1]
650
+
651
+ cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (
652
+ None,
653
+ None,
654
+ None,
655
+ None,
656
+ )
657
+ if key_padding_mask is not None:
658
+ kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)
659
+
660
+ if seqlen_q == 1:
661
+ key_padding_mask = torch.ones(batch_size, 1, device=q.device)
662
+ elif seqlen_q != seqlen_k:
663
+ key_padding_mask = key_padding_mask[:, -seqlen_q:]
664
+
665
+ q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)
666
+
667
+ if self.checkpointing:
668
+ attn_output = torch.utils.checkpoint.checkpoint(
669
+ self.inner_cross_attn,
670
+ q,
671
+ kv,
672
+ causal=causal,
673
+ cu_seqlens=cu_seqlens_q,
674
+ max_seqlen=max_seqlen_q,
675
+ cu_seqlens_k=cu_seqlens_k,
676
+ max_seqlen_k=max_seqlen_k,
677
+ )
678
+ else:
679
+ attn_output = self.inner_cross_attn(
680
+ q,
681
+ kv,
682
+ causal=causal,
683
+ cu_seqlens=cu_seqlens_q,
684
+ max_seqlen=max_seqlen_q,
685
+ cu_seqlens_k=cu_seqlens_k,
686
+ max_seqlen_k=max_seqlen_k,
687
+ )
688
+
689
+ return (
690
+ pad_input(attn_output, indices_q, batch_size, max_seqlen_q)
691
+ if key_padding_mask is not None
692
+ else attn_output
693
+ )
694
+
695
+ if self.checkpointing:
696
+ return torch.utils.checkpoint.checkpoint(
697
+ self.inner_cross_attn,
698
+ q,
699
+ kv,
700
+ key_padding_mask=key_padding_mask,
701
+ causal=causal,
702
+ )
703
+
704
+ return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
705
+
706
+ def forward(
707
+ self,
708
+ x: torch.FloatTensor,
709
+ past_key_values: Optional[InferenceParams] = None,
710
+ attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
711
+ **kwargs,
712
+ ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
713
+ if attention_mask is not None:
714
+ attention_mask = attention_mask.bool()
715
+ else:
716
+ attention_mask = None
717
+
718
+ # MHA
719
+ if self.n_head == self.n_head_kv:
720
+ if past_key_values is None:
721
+ # If `past_key_values` are not supplied, we run self-attention
722
+ attn_output = self._forward_self_attn(x, attention_mask)
723
+ else:
724
+ # If `past_key_values` are supplied, it means that we might have cached values and
725
+ # could take advantage of cross-attention
726
+ attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
727
+ # MQA / GQA
728
+ else:
729
+ # Regardless of `past_key_values` being supplied or not, it always use cross-attention
730
+ # because `q` and `kv` lengths might be different
731
+ attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
732
+
733
+ output = rearrange(attn_output, "... h d -> ... (h d)")
734
+ output = self.out_proj(output)
735
+
736
+ return output if not self.return_residual else (output, x)
737
+
738
+
739
+ class ParallelBlock(nn.Module):
740
+ """Parallel block.
741
+
742
+ This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
743
+
744
+ """
745
+
746
+ def __init__(
747
+ self,
748
+ config: PretrainedConfig,
749
+ block_idx: Optional[int] = None,
750
+ ) -> None:
751
+ super().__init__()
752
+
753
+ self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
754
+ self.resid_dropout = nn.Dropout(config.resid_pdrop)
755
+ self.block_idx = block_idx
756
+
757
+ self.mixer = MHA(config, layer_idx=block_idx)
758
+ self.mlp = MLP(config)
759
+
760
+ def forward(
761
+ self,
762
+ hidden_states: torch.FloatTensor,
763
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
764
+ attention_mask: Optional[torch.BoolTensor] = None,
765
+ **kwargs,
766
+ ) -> torch.FloatTensor:
767
+ residual = hidden_states
768
+ hidden_states = self.ln(hidden_states)
769
+
770
+ attn_outputs = self.mixer(
771
+ hidden_states,
772
+ past_key_values=past_key_values,
773
+ attention_mask=attention_mask,
774
+ )
775
+ if isinstance(attn_outputs, tuple):
776
+ attn_outputs = attn_outputs[0]
777
+
778
+ attn_outputs = self.resid_dropout(attn_outputs)
779
+ feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
780
+
781
+ hidden_states = attn_outputs + feed_forward_hidden_states + residual
782
+
783
+ return hidden_states
784
+
785
+
786
+ class CausalLMHead(nn.Module):
787
+ """Causal Language Modeling head.
788
+
789
+ Reference:
790
+ Improving Language Understanding by Generative Pre-Training.
791
+ https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
792
+
793
+ """
794
+
795
+ def __init__(self, config: PretrainedConfig) -> None:
796
+ super().__init__()
797
+
798
+ self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
799
+ self.linear = nn.Linear(config.n_embd, config.vocab_size)
800
+
801
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
802
+ hidden_states = self.ln(hidden_states)
803
+ logits = self.linear(hidden_states).to(torch.float32)
804
+
805
+ return logits
806
+
807
+
808
+ class CausalLMLoss(nn.Module):
809
+ """Causal Language Modeling loss.
810
+
811
+ Reference:
812
+ Improving Language Understanding by Generative Pre-Training.
813
+ https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
814
+
815
+ """
816
+
817
+ def __init__(self, shift_labels: bool = True) -> None:
818
+ super().__init__()
819
+
820
+ self.shift_labels = shift_labels
821
+ self.loss_fct = nn.CrossEntropyLoss()
822
+
823
+ def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:
824
+ if self.shift_labels:
825
+ logits = logits[..., :-1, :].contiguous()
826
+ labels = labels[..., 1:].contiguous()
827
+
828
+ loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
829
+
830
+ return loss
831
+
832
+
833
+ class PhiPreTrainedModel(PreTrainedModel):
834
+ """Phi pre-trained model."""
835
+
836
+ config_class = PhiConfig
837
+ base_model_prefix = "transformer"
838
+ supports_gradient_checkpointing = False
839
+ _no_split_modules = ["ParallelBlock"]
840
+
841
+ def __init__(self, *inputs, **kwargs) -> None:
842
+ super().__init__(*inputs, **kwargs)
843
+
844
+ def _init_weights(self, module: nn.Module) -> None:
845
+ if isinstance(module, (nn.Linear,)):
846
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
847
+ if module.bias is not None:
848
+ module.bias.data.zero_()
849
+ elif isinstance(module, nn.Embedding):
850
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
851
+ if module.padding_idx is not None:
852
+ module.weight.data[module.padding_idx].zero_()
853
+ elif isinstance(module, nn.LayerNorm):
854
+ if module.bias is not None:
855
+ module.bias.data.zero_()
856
+ module.weight.data.fill_(1.0)
857
+
858
+ def prepare_inputs_for_generation(
859
+ self,
860
+ input_ids: torch.LongTensor,
861
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
862
+ attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
863
+ **kwargs,
864
+ ) -> Dict[str, Any]:
865
+ if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
866
+ past_key_values = InferenceParams(
867
+ max_seqlen=self.config.n_positions,
868
+ max_batch_size=input_ids.shape[0],
869
+ seqlen_offset=0,
870
+ batch_size_offset=0,
871
+ key_value_memory_dict={},
872
+ lengths_per_sample=None,
873
+ )
874
+ else:
875
+ # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
876
+ past_key_values.seqlen_offset = input_ids.shape[1] - 1
877
+ input_ids = input_ids[:, -1].unsqueeze(-1)
878
+
879
+ return {
880
+ "input_ids": input_ids,
881
+ "past_key_values": past_key_values,
882
+ "attention_mask": attention_mask,
883
+ }
884
+
885
+
886
+ class PhiModel(PhiPreTrainedModel):
887
+ """Phi model."""
888
+
889
+ _keys_to_ignore_on_load_missing = [""]
890
+ _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
891
+
892
+ def __init__(self, config: PhiConfig) -> None:
893
+ super().__init__(config)
894
+
895
+ self.embd = Embedding(config)
896
+ self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
897
+ self.gradient_checkpointing = False
898
+ self.post_init()
899
+
900
+ def get_input_embeddings(self) -> nn.Embedding:
901
+ return self.embd.wte
902
+
903
+ def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
904
+ self.embd.wte = new_embeddings
905
+
906
+ def forward(
907
+ self,
908
+ input_ids: torch.LongTensor,
909
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
910
+ attention_mask: Optional[torch.BoolTensor] = None,
911
+ ) -> torch.FloatTensor:
912
+ hidden_states = self.embd(input_ids)
913
+
914
+ for layer in self.h:
915
+ hidden_states = layer(
916
+ hidden_states,
917
+ past_key_values=past_key_values,
918
+ attention_mask=attention_mask,
919
+ )
920
+
921
+ return hidden_states
922
+
923
+
924
+ class PhiForCausalLM(PhiPreTrainedModel):
925
+ """Phi for Causal Language Modeling."""
926
+
927
+ _keys_to_ignore_on_load_missing = [""]
928
+ _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
929
+
930
+ def __init__(self, config: PhiConfig) -> None:
931
+ super().__init__(config)
932
+
933
+ self.transformer = PhiModel(config)
934
+ self.lm_head = CausalLMHead(config)
935
+ self.loss = CausalLMLoss()
936
+
937
+ self.post_init()
938
+
939
+ def get_output_embeddings(self) -> nn.Linear:
940
+ return self.lm_head.linear
941
+
942
+ def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
943
+ self.lm_head.linear = new_embeddings
944
+
945
+ def forward(
946
+ self,
947
+ input_ids: torch.LongTensor,
948
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
949
+ attention_mask: Optional[torch.BoolTensor] = None,
950
+ labels: Optional[torch.LongTensor] = None,
951
+ **kwargs,
952
+ ) -> CausalLMOutputWithPast:
953
+ hidden_states = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask)
954
+ lm_logits = self.lm_head(hidden_states)
955
+
956
+ loss = None
957
+ if labels is not None:
958
+ loss = self.loss(lm_logits, labels)
959
+
960
+ return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values)
special_tokens_map.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<|endoftext|>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": "<|endoftext|>",
17
+ "unk_token": {
18
+ "content": "<|endoftext|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ }
24
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,324 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "added_tokens_decoder": {
4
+ "50256": {
5
+ "content": "<|endoftext|>",
6
+ "lstrip": false,
7
+ "normalized": false,
8
+ "rstrip": false,
9
+ "single_word": false,
10
+ "special": true
11
+ },
12
+ "50257": {
13
+ "content": " ",
14
+ "lstrip": false,
15
+ "normalized": true,
16
+ "rstrip": false,
17
+ "single_word": false,
18
+ "special": false
19
+ },
20
+ "50258": {
21
+ "content": " ",
22
+ "lstrip": false,
23
+ "normalized": true,
24
+ "rstrip": false,
25
+ "single_word": false,
26
+ "special": false
27
+ },
28
+ "50259": {
29
+ "content": " ",
30
+ "lstrip": false,
31
+ "normalized": true,
32
+ "rstrip": false,
33
+ "single_word": false,
34
+ "special": false
35
+ },
36
+ "50260": {
37
+ "content": " ",
38
+ "lstrip": false,
39
+ "normalized": true,
40
+ "rstrip": false,
41
+ "single_word": false,
42
+ "special": false
43
+ },
44
+ "50261": {
45
+ "content": " ",
46
+ "lstrip": false,
47
+ "normalized": true,
48
+ "rstrip": false,
49
+ "single_word": false,
50
+ "special": false
51
+ },
52
+ "50262": {
53
+ "content": " ",
54
+ "lstrip": false,
55
+ "normalized": true,
56
+ "rstrip": false,
57
+ "single_word": false,
58
+ "special": false
59
+ },
60
+ "50263": {
61
+ "content": " ",
62
+ "lstrip": false,
63
+ "normalized": true,
64
+ "rstrip": false,
65
+ "single_word": false,
66
+ "special": false
67
+ },
68
+ "50264": {
69
+ "content": " ",
70
+ "lstrip": false,
71
+ "normalized": true,
72
+ "rstrip": false,
73
+ "single_word": false,
74
+ "special": false
75
+ },
76
+ "50265": {
77
+ "content": " ",
78
+ "lstrip": false,
79
+ "normalized": true,
80
+ "rstrip": false,
81
+ "single_word": false,
82
+ "special": false
83
+ },
84
+ "50266": {
85
+ "content": " ",
86
+ "lstrip": false,
87
+ "normalized": true,
88
+ "rstrip": false,
89
+ "single_word": false,
90
+ "special": false
91
+ },
92
+ "50267": {
93
+ "content": " ",
94
+ "lstrip": false,
95
+ "normalized": true,
96
+ "rstrip": false,
97
+ "single_word": false,
98
+ "special": false
99
+ },
100
+ "50268": {
101
+ "content": " ",
102
+ "lstrip": false,
103
+ "normalized": true,
104
+ "rstrip": false,
105
+ "single_word": false,
106
+ "special": false
107
+ },
108
+ "50269": {
109
+ "content": " ",
110
+ "lstrip": false,
111
+ "normalized": true,
112
+ "rstrip": false,
113
+ "single_word": false,
114
+ "special": false
115
+ },
116
+ "50270": {
117
+ "content": " ",
118
+ "lstrip": false,
119
+ "normalized": true,
120
+ "rstrip": false,
121
+ "single_word": false,
122
+ "special": false
123
+ },
124
+ "50271": {
125
+ "content": " ",
126
+ "lstrip": false,
127
+ "normalized": true,
128
+ "rstrip": false,
129
+ "single_word": false,
130
+ "special": false
131
+ },
132
+ "50272": {
133
+ "content": " ",
134
+ "lstrip": false,
135
+ "normalized": true,
136
+ "rstrip": false,
137
+ "single_word": false,
138
+ "special": false
139
+ },
140
+ "50273": {
141
+ "content": " ",
142
+ "lstrip": false,
143
+ "normalized": true,
144
+ "rstrip": false,
145
+ "single_word": false,
146
+ "special": false
147
+ },
148
+ "50274": {
149
+ "content": " ",
150
+ "lstrip": false,
151
+ "normalized": true,
152
+ "rstrip": false,
153
+ "single_word": false,
154
+ "special": false
155
+ },
156
+ "50275": {
157
+ "content": " ",
158
+ "lstrip": false,
159
+ "normalized": true,
160
+ "rstrip": false,
161
+ "single_word": false,
162
+ "special": false
163
+ },
164
+ "50276": {
165
+ "content": " ",
166
+ "lstrip": false,
167
+ "normalized": true,
168
+ "rstrip": false,
169
+ "single_word": false,
170
+ "special": false
171
+ },
172
+ "50277": {
173
+ "content": " ",
174
+ "lstrip": false,
175
+ "normalized": true,
176
+ "rstrip": false,
177
+ "single_word": false,
178
+ "special": false
179
+ },
180
+ "50278": {
181
+ "content": " ",
182
+ "lstrip": false,
183
+ "normalized": true,
184
+ "rstrip": false,
185
+ "single_word": false,
186
+ "special": false
187
+ },
188
+ "50279": {
189
+ "content": " ",
190
+ "lstrip": false,
191
+ "normalized": true,
192
+ "rstrip": false,
193
+ "single_word": false,
194
+ "special": false
195
+ },
196
+ "50280": {
197
+ "content": " ",
198
+ "lstrip": false,
199
+ "normalized": true,
200
+ "rstrip": false,
201
+ "single_word": false,
202
+ "special": false
203
+ },
204
+ "50281": {
205
+ "content": " ",
206
+ "lstrip": false,
207
+ "normalized": true,
208
+ "rstrip": false,
209
+ "single_word": false,
210
+ "special": false
211
+ },
212
+ "50282": {
213
+ "content": " ",
214
+ "lstrip": false,
215
+ "normalized": true,
216
+ "rstrip": false,
217
+ "single_word": false,
218
+ "special": false
219
+ },
220
+ "50283": {
221
+ "content": " ",
222
+ "lstrip": false,
223
+ "normalized": true,
224
+ "rstrip": false,
225
+ "single_word": false,
226
+ "special": false
227
+ },
228
+ "50284": {
229
+ "content": " ",
230
+ "lstrip": false,
231
+ "normalized": true,
232
+ "rstrip": false,
233
+ "single_word": false,
234
+ "special": false
235
+ },
236
+ "50285": {
237
+ "content": " ",
238
+ "lstrip": false,
239
+ "normalized": true,
240
+ "rstrip": false,
241
+ "single_word": false,
242
+ "special": false
243
+ },
244
+ "50286": {
245
+ "content": " ",
246
+ "lstrip": false,
247
+ "normalized": true,
248
+ "rstrip": false,
249
+ "single_word": false,
250
+ "special": false
251
+ },
252
+ "50287": {
253
+ "content": "\t\t\t\t\t\t\t\t\t",
254
+ "lstrip": false,
255
+ "normalized": true,
256
+ "rstrip": false,
257
+ "single_word": false,
258
+ "special": false
259
+ },
260
+ "50288": {
261
+ "content": "\t\t\t\t\t\t\t\t",
262
+ "lstrip": false,
263
+ "normalized": true,
264
+ "rstrip": false,
265
+ "single_word": false,
266
+ "special": false
267
+ },
268
+ "50289": {
269
+ "content": "\t\t\t\t\t\t\t",
270
+ "lstrip": false,
271
+ "normalized": true,
272
+ "rstrip": false,
273
+ "single_word": false,
274
+ "special": false
275
+ },
276
+ "50290": {
277
+ "content": "\t\t\t\t\t\t",
278
+ "lstrip": false,
279
+ "normalized": true,
280
+ "rstrip": false,
281
+ "single_word": false,
282
+ "special": false
283
+ },
284
+ "50291": {
285
+ "content": "\t\t\t\t\t",
286
+ "lstrip": false,
287
+ "normalized": true,
288
+ "rstrip": false,
289
+ "single_word": false,
290
+ "special": false
291
+ },
292
+ "50292": {
293
+ "content": "\t\t\t\t",
294
+ "lstrip": false,
295
+ "normalized": true,
296
+ "rstrip": false,
297
+ "single_word": false,
298
+ "special": false
299
+ },
300
+ "50293": {
301
+ "content": "\t\t\t",
302
+ "lstrip": false,
303
+ "normalized": true,
304
+ "rstrip": false,
305
+ "single_word": false,
306
+ "special": false
307
+ },
308
+ "50294": {
309
+ "content": "\t\t",
310
+ "lstrip": false,
311
+ "normalized": true,
312
+ "rstrip": false,
313
+ "single_word": false,
314
+ "special": false
315
+ }
316
+ },
317
+ "bos_token": "<|endoftext|>",
318
+ "clean_up_tokenization_spaces": true,
319
+ "eos_token": "<|endoftext|>",
320
+ "model_max_length": 2048,
321
+ "pad_token": "<|endoftext|>",
322
+ "tokenizer_class": "CodeGenTokenizer",
323
+ "unk_token": "<|endoftext|>"
324
+ }
vocab.json ADDED
The diff for this file is too large to render. See raw diff