Judychoieee commited on
Commit
091400d
·
1 Parent(s): f2f133f

add model files

Browse files
added_tokens.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "</think>": 151668,
3
+ "</tool_call>": 151658,
4
+ "</tool_response>": 151666,
5
+ "<think>": 151667,
6
+ "<tool_call>": 151657,
7
+ "<tool_response>": 151665,
8
+ "<|box_end|>": 151649,
9
+ "<|box_start|>": 151648,
10
+ "<|endoftext|>": 151643,
11
+ "<|file_sep|>": 151664,
12
+ "<|fim_middle|>": 151660,
13
+ "<|fim_pad|>": 151662,
14
+ "<|fim_prefix|>": 151659,
15
+ "<|fim_suffix|>": 151661,
16
+ "<|im_end|>": 151645,
17
+ "<|im_start|>": 151644,
18
+ "<|image_pad|>": 151655,
19
+ "<|object_ref_end|>": 151647,
20
+ "<|object_ref_start|>": 151646,
21
+ "<|quad_end|>": 151651,
22
+ "<|quad_start|>": 151650,
23
+ "<|repo_name|>": 151663,
24
+ "<|video_pad|>": 151656,
25
+ "<|vision_end|>": 151653,
26
+ "<|vision_pad|>": 151654,
27
+ "<|vision_start|>": 151652
28
+ }
config.json ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Qwen3Model"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "modeling_qwen3_qr.Qwen3ConfigGating",
9
+ "AutoModel": "modeling_qwen3_qr.Qwen3Model"
10
+ },
11
+ "bos_token_id": 151643,
12
+ "eos_token_id": 151645,
13
+ "head_dim": 128,
14
+ "hidden_act": "silu",
15
+ "hidden_size": 2560,
16
+ "initializer_range": 0.02,
17
+ "intermediate_size": 9728,
18
+ "layer_types": [
19
+ "full_attention",
20
+ "full_attention",
21
+ "full_attention",
22
+ "full_attention",
23
+ "full_attention",
24
+ "full_attention",
25
+ "full_attention",
26
+ "full_attention",
27
+ "full_attention",
28
+ "full_attention",
29
+ "full_attention",
30
+ "full_attention",
31
+ "full_attention",
32
+ "full_attention",
33
+ "full_attention",
34
+ "full_attention",
35
+ "full_attention",
36
+ "full_attention",
37
+ "full_attention",
38
+ "full_attention",
39
+ "full_attention",
40
+ "full_attention",
41
+ "full_attention",
42
+ "full_attention",
43
+ "full_attention",
44
+ "full_attention",
45
+ "full_attention",
46
+ "full_attention",
47
+ "full_attention",
48
+ "full_attention",
49
+ "full_attention",
50
+ "full_attention",
51
+ "full_attention",
52
+ "full_attention",
53
+ "full_attention",
54
+ "full_attention"
55
+ ],
56
+ "max_position_embeddings": 262144,
57
+ "max_window_layers": 36,
58
+ "model_type": "qwen3",
59
+ "num_attention_heads": 32,
60
+ "num_hidden_layers": 36,
61
+ "num_key_value_heads": 8,
62
+ "qr_end_layer": 25,
63
+ "qr_head_list": [
64
+ [
65
+ 20,
66
+ 15
67
+ ],
68
+ [
69
+ 21,
70
+ 11
71
+ ],
72
+ [
73
+ 17,
74
+ 27
75
+ ],
76
+ [
77
+ 23,
78
+ 10
79
+ ],
80
+ [
81
+ 22,
82
+ 4
83
+ ],
84
+ [
85
+ 21,
86
+ 10
87
+ ],
88
+ [
89
+ 21,
90
+ 8
91
+ ],
92
+ [
93
+ 21,
94
+ 18
95
+ ],
96
+ [
97
+ 18,
98
+ 15
99
+ ],
100
+ [
101
+ 18,
102
+ 19
103
+ ],
104
+ [
105
+ 17,
106
+ 25
107
+ ],
108
+ [
109
+ 17,
110
+ 17
111
+ ],
112
+ [
113
+ 24,
114
+ 13
115
+ ],
116
+ [
117
+ 17,
118
+ 4
119
+ ],
120
+ [
121
+ 19,
122
+ 12
123
+ ],
124
+ [
125
+ 21,
126
+ 31
127
+ ]
128
+ ],
129
+ "qr_head_list_mapped": [
130
+ [
131
+ 3,
132
+ 15
133
+ ],
134
+ [
135
+ 4,
136
+ 11
137
+ ],
138
+ [
139
+ 0,
140
+ 27
141
+ ],
142
+ [
143
+ 6,
144
+ 10
145
+ ],
146
+ [
147
+ 5,
148
+ 4
149
+ ],
150
+ [
151
+ 4,
152
+ 10
153
+ ],
154
+ [
155
+ 4,
156
+ 8
157
+ ],
158
+ [
159
+ 4,
160
+ 18
161
+ ],
162
+ [
163
+ 1,
164
+ 15
165
+ ],
166
+ [
167
+ 1,
168
+ 19
169
+ ],
170
+ [
171
+ 0,
172
+ 25
173
+ ],
174
+ [
175
+ 0,
176
+ 17
177
+ ],
178
+ [
179
+ 7,
180
+ 13
181
+ ],
182
+ [
183
+ 0,
184
+ 4
185
+ ],
186
+ [
187
+ 2,
188
+ 12
189
+ ],
190
+ [
191
+ 4,
192
+ 31
193
+ ]
194
+ ],
195
+ "qr_start_layer": 17,
196
+ "rms_norm_eps": 1e-06,
197
+ "rope_scaling": null,
198
+ "rope_theta": 5000000,
199
+ "sliding_window": null,
200
+ "tie_word_embeddings": true,
201
+ "top_k_heads_per_layer": 2,
202
+ "torch_dtype": "bfloat16",
203
+ "transformers_version": "4.53.0",
204
+ "use_cache": true,
205
+ "use_sliding_window": false,
206
+ "vocab_size": 151936
207
+ }
custom_cache_new.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, Optional, Tuple
2
+ from transformers.cache_utils import DynamicCache
3
+ import torch
4
+
5
+
6
+ class DynamicCacheWithQuery(DynamicCache):
7
+ """
8
+ Cache class used for QRRetriever;
9
+ LJN: put the query states in the cache_kwargs to keep the same signature as DynamicCache
10
+ LJN: please take the query states from the cache_kwargs
11
+ """
12
+
13
+ def __init__(self, query_indices=[]) -> None:
14
+ super().__init__()
15
+ self._query_indices = query_indices # indices for query vectors to save
16
+ self.query_cache = []
17
+
18
+ def update(
19
+ self,
20
+ key_states: torch.Tensor,
21
+ value_states: torch.Tensor,
22
+ layer_idx: int,
23
+ cache_kwargs: Optional[Dict[str, Any]] = None,
24
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
25
+ """
26
+ Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
27
+
28
+ Parameters:
29
+ key_states (`torch.Tensor`):
30
+ The new key states to cache.
31
+ value_states (`torch.Tensor`):
32
+ The new value states to cache.
33
+ layer_idx (`int`):
34
+ The index of the layer to cache the states for.
35
+ cache_kwargs (`Dict[str, Any]`, `optional`):
36
+ Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`.
37
+
38
+ Return:
39
+ A tuple containing the updated key and value states.
40
+ """
41
+ # Update the number of seen tokens
42
+ if layer_idx == 0:
43
+ self._seen_tokens += key_states.shape[-2]
44
+
45
+ # Update the cache
46
+ if key_states is not None:
47
+ if len(self.key_cache) <= layer_idx:
48
+ # There may be skipped layers, fill them with empty lists
49
+ for _ in range(len(self.key_cache), layer_idx):
50
+ self.key_cache.append(torch.tensor([]))
51
+ self.value_cache.append(torch.tensor([]))
52
+ self.key_cache.append(key_states)
53
+ self.value_cache.append(value_states)
54
+ elif (
55
+ not self.key_cache[layer_idx].numel() # prefers not t.numel() to len(t) == 0 to export the model
56
+ ): # fills previously skipped layers; checking for tensor causes errors
57
+ self.key_cache[layer_idx] = key_states
58
+ self.value_cache[layer_idx] = value_states
59
+ else:
60
+ self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
61
+ self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
62
+
63
+ if cache_kwargs is not None:
64
+ query_states = cache_kwargs.get("query_states", None)
65
+ else:
66
+ query_states = None
67
+ if query_states is not None:
68
+ if len(self.query_cache) <= layer_idx:
69
+ self.query_cache.append(query_states)
70
+ else:
71
+ self.query_cache[layer_idx] = torch.cat([self.query_cache[layer_idx], query_states], dim=-2)
72
+ return self.key_cache[layer_idx], self.value_cache[layer_idx]
73
+
74
+ @classmethod
75
+ def from_legacy_cache_with_query_indices(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, query_indices = []) -> "DynamicCache":
76
+ """Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for
77
+ backward compatibility."""
78
+ cache = cls(query_indices=query_indices)
79
+ if past_key_values is not None:
80
+ for layer_idx in range(len(past_key_values)):
81
+ key_states, value_states = past_key_values[layer_idx]
82
+ cache.update(key_states, value_states, layer_idx)
83
+ return cache
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model-00001-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2a3ac26b8bc75eac6041c2f9f18ebfe452992897cd0397c1ff1c25ffa996e5e4
3
+ size 4967213984
model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bcaf871750830078b99143838c686a600de5deaa6190327ea715431c33c4f464
3
+ size 857269208
model.safetensors.index.json ADDED
@@ -0,0 +1,284 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_parameters": 2912226560,
4
+ "total_size": 5824453120
5
+ },
6
+ "weight_map": {
7
+ "embed_tokens.weight": "model-00001-of-00002.safetensors",
8
+ "layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
9
+ "layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
10
+ "layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
11
+ "layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
12
+ "layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
13
+ "layers.0.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
14
+ "layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
15
+ "layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
16
+ "layers.0.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
17
+ "layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
18
+ "layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
19
+ "layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
20
+ "layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
21
+ "layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
22
+ "layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
23
+ "layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
24
+ "layers.1.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
25
+ "layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
26
+ "layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
27
+ "layers.1.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
28
+ "layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
29
+ "layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
30
+ "layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
31
+ "layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
32
+ "layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
33
+ "layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
34
+ "layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
35
+ "layers.10.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
36
+ "layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
37
+ "layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
38
+ "layers.10.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
39
+ "layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
40
+ "layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
41
+ "layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
42
+ "layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
43
+ "layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
44
+ "layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
45
+ "layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
46
+ "layers.11.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
47
+ "layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
48
+ "layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
49
+ "layers.11.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
50
+ "layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
51
+ "layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
52
+ "layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
53
+ "layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
54
+ "layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
55
+ "layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
56
+ "layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
57
+ "layers.12.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
58
+ "layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
59
+ "layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
60
+ "layers.12.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
61
+ "layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
62
+ "layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
63
+ "layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
64
+ "layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
65
+ "layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
66
+ "layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
67
+ "layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
68
+ "layers.13.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
69
+ "layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
70
+ "layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
71
+ "layers.13.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
72
+ "layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
73
+ "layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
74
+ "layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
75
+ "layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
76
+ "layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
77
+ "layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
78
+ "layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
79
+ "layers.14.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
80
+ "layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
81
+ "layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
82
+ "layers.14.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
83
+ "layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
84
+ "layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
85
+ "layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
86
+ "layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
87
+ "layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
88
+ "layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
89
+ "layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
90
+ "layers.15.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
91
+ "layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
92
+ "layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
93
+ "layers.15.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
94
+ "layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
95
+ "layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
96
+ "layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
97
+ "layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
98
+ "layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
99
+ "layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
100
+ "layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
101
+ "layers.16.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
102
+ "layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
103
+ "layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
104
+ "layers.16.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
105
+ "layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
106
+ "layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
107
+ "layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
108
+ "layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
109
+ "layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
110
+ "layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
111
+ "layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
112
+ "layers.17.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
113
+ "layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
114
+ "layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
115
+ "layers.17.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
116
+ "layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
117
+ "layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
118
+ "layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
119
+ "layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
120
+ "layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
121
+ "layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
122
+ "layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
123
+ "layers.18.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
124
+ "layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
125
+ "layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
126
+ "layers.18.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
127
+ "layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
128
+ "layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
129
+ "layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
130
+ "layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
131
+ "layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
132
+ "layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
133
+ "layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
134
+ "layers.19.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
135
+ "layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
136
+ "layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
137
+ "layers.19.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
138
+ "layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
139
+ "layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
140
+ "layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
141
+ "layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
142
+ "layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
143
+ "layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
144
+ "layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
145
+ "layers.2.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
146
+ "layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
147
+ "layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
148
+ "layers.2.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
149
+ "layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
150
+ "layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
151
+ "layers.20.input_layernorm.weight": "model-00002-of-00002.safetensors",
152
+ "layers.20.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
153
+ "layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
154
+ "layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
155
+ "layers.20.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
156
+ "layers.20.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
157
+ "layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
158
+ "layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
159
+ "layers.20.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
160
+ "layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
161
+ "layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
162
+ "layers.21.input_layernorm.weight": "model-00002-of-00002.safetensors",
163
+ "layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
164
+ "layers.21.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
165
+ "layers.21.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
166
+ "layers.21.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
167
+ "layers.21.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
168
+ "layers.21.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
169
+ "layers.21.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
170
+ "layers.21.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
171
+ "layers.21.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
172
+ "layers.21.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
173
+ "layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
174
+ "layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
175
+ "layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
176
+ "layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
177
+ "layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
178
+ "layers.22.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
179
+ "layers.22.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
180
+ "layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
181
+ "layers.22.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
182
+ "layers.22.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
183
+ "layers.22.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
184
+ "layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
185
+ "layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
186
+ "layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
187
+ "layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
188
+ "layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
189
+ "layers.23.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
190
+ "layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
191
+ "layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
192
+ "layers.23.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
193
+ "layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
194
+ "layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
195
+ "layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
196
+ "layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
197
+ "layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
198
+ "layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
199
+ "layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
200
+ "layers.24.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
201
+ "layers.24.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
202
+ "layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
203
+ "layers.24.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
204
+ "layers.24.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
205
+ "layers.24.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
206
+ "layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
207
+ "layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
208
+ "layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
209
+ "layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
210
+ "layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
211
+ "layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
212
+ "layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
213
+ "layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
214
+ "layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
215
+ "layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
216
+ "layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
217
+ "layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
218
+ "layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
219
+ "layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
220
+ "layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
221
+ "layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
222
+ "layers.4.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
223
+ "layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
224
+ "layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
225
+ "layers.4.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
226
+ "layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
227
+ "layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
228
+ "layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
229
+ "layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
230
+ "layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
231
+ "layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
232
+ "layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
233
+ "layers.5.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
234
+ "layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
235
+ "layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
236
+ "layers.5.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
237
+ "layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
238
+ "layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
239
+ "layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
240
+ "layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
241
+ "layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
242
+ "layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
243
+ "layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
244
+ "layers.6.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
245
+ "layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
246
+ "layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
247
+ "layers.6.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
248
+ "layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
249
+ "layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
250
+ "layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
251
+ "layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
252
+ "layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
253
+ "layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
254
+ "layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
255
+ "layers.7.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
256
+ "layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
257
+ "layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
258
+ "layers.7.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
259
+ "layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
260
+ "layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
261
+ "layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
262
+ "layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
263
+ "layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
264
+ "layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
265
+ "layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
266
+ "layers.8.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
267
+ "layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
268
+ "layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
269
+ "layers.8.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
270
+ "layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
271
+ "layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
272
+ "layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
273
+ "layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
274
+ "layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
275
+ "layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
276
+ "layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
277
+ "layers.9.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
278
+ "layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
279
+ "layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
280
+ "layers.9.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
281
+ "layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
282
+ "layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors"
283
+ }
284
+ }
modeling_qwen3_qr.py ADDED
@@ -0,0 +1,825 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/qwen3/modular_qwen3.py.
3
+ # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
+ # the file from the modular. If any change should be done, please apply the change to the
5
+ # modular_qwen3.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # coding=utf-8
8
+ # Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
9
+ #
10
+ # Licensed under the Apache License, Version 2.0 (the "License");
11
+ # you may not use this file except in compliance with the License.
12
+ # You may obtain a copy of the License at
13
+ #
14
+ # http://www.apache.org/licenses/LICENSE-2.0
15
+ #
16
+ # Unless required by applicable law or agreed to in writing, software
17
+ # distributed under the License is distributed on an "AS IS" BASIS,
18
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ # See the License for the specific language governing permissions and
20
+ # limitations under the License.
21
+
22
+ from functools import partial
23
+ from typing import Callable, Optional, Tuple
24
+ # from typing import Union
25
+
26
+ import torch
27
+ from torch import nn
28
+
29
+ from transformers.activations import ACT2FN
30
+ from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
31
+
32
+ try:
33
+ from .custom_cache_new import DynamicCacheWithQuery
34
+ except ImportError:
35
+ from custom_cache_new import DynamicCacheWithQuery
36
+
37
+
38
+ # from transformers.generation import GenerationMixin
39
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
40
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
41
+ from transformers.modeling_outputs import (
42
+ BaseModelOutputWithPast,
43
+ # CausalLMOutputWithPast,
44
+ # QuestionAnsweringModelOutput,
45
+ # SequenceClassifierOutputWithPast,
46
+ # TokenClassifierOutput,
47
+ )
48
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
49
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
50
+ from transformers.processing_utils import Unpack
51
+ from transformers.utils import (
52
+ # LossKwargs,
53
+ # add_code_sample_docstrings,
54
+ add_start_docstrings,
55
+ add_start_docstrings_to_model_forward,
56
+ can_return_tuple,
57
+ logging,
58
+ # replace_return_docstrings,
59
+ )
60
+ # from transformers.utils.deprecation import deprecate_kwarg
61
+ from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
62
+
63
+
64
+ logger = logging.get_logger(__name__)
65
+
66
+ _CHECKPOINT_FOR_DOC = "Qwen/Qwen3-8B"
67
+ _CONFIG_FOR_DOC = "Qwen3Config"
68
+
69
+
70
+ class Qwen3ConfigGating(Qwen3Config):
71
+
72
+ def __init__(
73
+ self,
74
+ qr_head_list=None,
75
+ qr_start_layer=17,
76
+ qr_end_layer=25,
77
+ top_k_heads_per_layer=2,
78
+ **kwargs
79
+ ):
80
+ super().__init__(**kwargs)
81
+ if qr_head_list is None:
82
+ qr_head_list = [[20, 15], [21, 11], [17, 27], [23, 10],
83
+ [22, 4], [21, 10], [21, 8], [21, 18],
84
+ [18, 15], [18, 19], [17, 25], [17, 17],
85
+ [24, 13], [17, 4], [19, 12], [21, 31]]
86
+ self.qr_head_list = qr_head_list
87
+ # used for the variant
88
+ self.qr_start_layer = qr_start_layer
89
+ self.qr_end_layer = qr_end_layer
90
+ # used for the variant
91
+ self.top_k_heads_per_layer = top_k_heads_per_layer
92
+ self.qr_head_list_mapped = [[qr_layer[0] - self.qr_start_layer, qr_layer[1]] for qr_layer in self.qr_head_list]
93
+
94
+
95
+ class Qwen3RMSNorm(nn.Module):
96
+ def __init__(self, hidden_size, eps=1e-6):
97
+ """
98
+ Qwen3RMSNorm is equivalent to T5LayerNorm
99
+ """
100
+ super().__init__()
101
+ self.weight = nn.Parameter(torch.ones(hidden_size))
102
+ self.variance_epsilon = eps
103
+
104
+ def forward(self, hidden_states):
105
+ input_dtype = hidden_states.dtype
106
+ hidden_states = hidden_states.to(torch.float32)
107
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
108
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
109
+ return self.weight * hidden_states.to(input_dtype)
110
+
111
+ def extra_repr(self):
112
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
113
+
114
+
115
+ class Qwen3MLP(nn.Module):
116
+ def __init__(self, config):
117
+ super().__init__()
118
+ self.config = config
119
+ self.hidden_size = config.hidden_size
120
+ self.intermediate_size = config.intermediate_size
121
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
122
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
123
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
124
+ self.act_fn = ACT2FN[config.hidden_act]
125
+
126
+ def forward(self, x):
127
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
128
+ return down_proj
129
+
130
+
131
+ def rotate_half(x):
132
+ """Rotates half the hidden dims of the input."""
133
+ x1 = x[..., : x.shape[-1] // 2]
134
+ x2 = x[..., x.shape[-1] // 2 :]
135
+ return torch.cat((-x2, x1), dim=-1)
136
+
137
+
138
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
139
+ """Applies Rotary Position Embedding to the query and key tensors.
140
+
141
+ Args:
142
+ q (`torch.Tensor`): The query tensor.
143
+ k (`torch.Tensor`): The key tensor.
144
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
145
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
146
+ position_ids (`torch.Tensor`, *optional*):
147
+ Deprecated and unused.
148
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
149
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
150
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
151
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
152
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
153
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
154
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
155
+ Returns:
156
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
157
+ """
158
+ cos = cos.unsqueeze(unsqueeze_dim)
159
+ sin = sin.unsqueeze(unsqueeze_dim)
160
+ q_embed = (q * cos) + (rotate_half(q) * sin)
161
+ k_embed = (k * cos) + (rotate_half(k) * sin)
162
+ return q_embed, k_embed
163
+
164
+
165
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
166
+ """
167
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
168
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
169
+ """
170
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
171
+ if n_rep == 1:
172
+ return hidden_states
173
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
174
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
175
+
176
+
177
+ def eager_attention_forward(
178
+ module: nn.Module,
179
+ query: torch.Tensor,
180
+ key: torch.Tensor,
181
+ value: torch.Tensor,
182
+ attention_mask: Optional[torch.Tensor],
183
+ scaling: float,
184
+ dropout: float = 0.0,
185
+ **kwargs,
186
+ ):
187
+ key_states = repeat_kv(key, module.num_key_value_groups)
188
+ value_states = repeat_kv(value, module.num_key_value_groups)
189
+
190
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
191
+ if attention_mask is not None:
192
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
193
+ attn_weights = attn_weights + causal_mask
194
+
195
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
196
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
197
+ attn_output = torch.matmul(attn_weights, value_states)
198
+ attn_output = attn_output.transpose(1, 2).contiguous()
199
+
200
+ return attn_output, attn_weights
201
+
202
+
203
+ class Qwen3Attention(nn.Module):
204
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
205
+
206
+ def __init__(self, config: Qwen3Config, layer_idx: int):
207
+ super().__init__()
208
+ self.config = config
209
+ self.layer_idx = layer_idx
210
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
211
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
212
+ self.scaling = self.head_dim**-0.5
213
+ self.attention_dropout = config.attention_dropout
214
+ self.is_causal = True
215
+
216
+ self.q_proj = nn.Linear(
217
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
218
+ )
219
+ self.k_proj = nn.Linear(
220
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
221
+ )
222
+ self.v_proj = nn.Linear(
223
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
224
+ )
225
+ self.o_proj = nn.Linear(
226
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
227
+ )
228
+ self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
229
+ self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape
230
+ self.sliding_window = config.sliding_window
231
+ if not (
232
+ self.config.use_sliding_window
233
+ and getattr(self.config, "sliding_window", None) is not None
234
+ and self.layer_idx >= self.config.max_window_layers
235
+ ):
236
+ self.sliding_window = None
237
+
238
+ def forward(
239
+ self,
240
+ hidden_states: torch.Tensor,
241
+ position_embeddings: Tuple[torch.Tensor, torch.Tensor],
242
+ attention_mask: Optional[torch.Tensor],
243
+ past_key_value: Optional[Cache] = None,
244
+ cache_position: Optional[torch.LongTensor] = None,
245
+ **kwargs: Unpack[FlashAttentionKwargs],
246
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
247
+ input_shape = hidden_states.shape[:-1]
248
+ hidden_shape = (*input_shape, -1, self.head_dim)
249
+
250
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
251
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
252
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
253
+
254
+ cos, sin = position_embeddings
255
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
256
+
257
+ if past_key_value is not None:
258
+ if isinstance(past_key_value, DynamicCacheWithQuery):
259
+ # LJN: add query hidden states here
260
+ query_states_to_cache = query_states[:, :, past_key_value._query_indices, :]
261
+ cache_kwargs = {
262
+ "sin": sin, "cos": cos, "cache_position": cache_position, "query_states": query_states_to_cache}
263
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
264
+ else:
265
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
266
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
267
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
268
+
269
+ attention_interface: Callable = eager_attention_forward
270
+ if self.config._attn_implementation != "eager":
271
+ if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
272
+ logger.warning_once(
273
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
274
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
275
+ )
276
+ else:
277
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
278
+
279
+ attn_output, attn_weights = attention_interface(
280
+ self,
281
+ query_states,
282
+ key_states,
283
+ value_states,
284
+ attention_mask,
285
+ dropout=0.0 if not self.training else self.attention_dropout,
286
+ scaling=self.scaling,
287
+ sliding_window=self.sliding_window, # diff with Llama
288
+ **kwargs,
289
+ )
290
+
291
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
292
+ attn_output = self.o_proj(attn_output)
293
+ return attn_output, attn_weights
294
+
295
+
296
+ class Qwen3DecoderLayer(nn.Module):
297
+ def __init__(self, config: Qwen3Config, layer_idx: int):
298
+ super().__init__()
299
+ self.hidden_size = config.hidden_size
300
+ self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
301
+ self.mlp = Qwen3MLP(config)
302
+ self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
303
+ self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
304
+ if (
305
+ config.sliding_window and config._attn_implementation != "flash_attention_2"
306
+ ): # diff with Llama is this warning
307
+ logger.warning_once(
308
+ f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
309
+ "unexpected results may be encountered."
310
+ )
311
+
312
+ def forward(
313
+ self,
314
+ hidden_states: torch.Tensor,
315
+ attention_mask: Optional[torch.Tensor] = None,
316
+ position_ids: Optional[torch.LongTensor] = None,
317
+ past_key_value: Optional[Cache] = None,
318
+ output_attentions: Optional[bool] = False,
319
+ use_cache: Optional[bool] = False,
320
+ cache_position: Optional[torch.LongTensor] = None,
321
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
322
+ **kwargs: Unpack[FlashAttentionKwargs],
323
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
324
+ residual = hidden_states
325
+
326
+ hidden_states = self.input_layernorm(hidden_states)
327
+
328
+ # Self Attention
329
+ hidden_states, self_attn_weights = self.self_attn(
330
+ hidden_states=hidden_states,
331
+ attention_mask=attention_mask,
332
+ position_ids=position_ids,
333
+ past_key_value=past_key_value,
334
+ output_attentions=output_attentions,
335
+ use_cache=use_cache,
336
+ cache_position=cache_position,
337
+ position_embeddings=position_embeddings,
338
+ **kwargs,
339
+ )
340
+ hidden_states = residual + hidden_states
341
+
342
+ # Fully Connected
343
+ residual = hidden_states
344
+ hidden_states = self.post_attention_layernorm(hidden_states)
345
+ hidden_states = self.mlp(hidden_states)
346
+ hidden_states = residual + hidden_states
347
+
348
+ outputs = (hidden_states,)
349
+ if output_attentions:
350
+ outputs += (self_attn_weights,)
351
+
352
+ return outputs
353
+
354
+
355
+ class Qwen3RotaryEmbedding(nn.Module):
356
+ def __init__(self, config: Qwen3Config, device=None):
357
+ super().__init__()
358
+ # BC: "rope_type" was originally "type"
359
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
360
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
361
+ else:
362
+ self.rope_type = "default"
363
+ self.max_seq_len_cached = config.max_position_embeddings
364
+ self.original_max_seq_len = config.max_position_embeddings
365
+
366
+ self.config = config
367
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
368
+
369
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
370
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
371
+ self.original_inv_freq = self.inv_freq
372
+
373
+ @torch.no_grad()
374
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
375
+ def forward(self, x, position_ids):
376
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
377
+ position_ids_expanded = position_ids[:, None, :].float()
378
+
379
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
380
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
381
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
382
+ emb = torch.cat((freqs, freqs), dim=-1)
383
+ cos = emb.cos() * self.attention_scaling
384
+ sin = emb.sin() * self.attention_scaling
385
+
386
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
387
+
388
+
389
+ QWEN3_START_DOCSTRING = r"""
390
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
391
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
392
+ etc.)
393
+
394
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
395
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
396
+ and behavior.
397
+
398
+ Parameters:
399
+ config ([`Qwen3Config`]):
400
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
401
+ load the weights associated with the model, only the configuration. Check out the
402
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
403
+ """
404
+
405
+
406
+ @add_start_docstrings(
407
+ "The bare Qwen3 Model outputting raw hidden-states without any specific head on top.",
408
+ QWEN3_START_DOCSTRING,
409
+ )
410
+ class Qwen3PreTrainedModel(PreTrainedModel):
411
+ config_class = Qwen3ConfigGating
412
+ base_model_prefix = "model"
413
+ supports_gradient_checkpointing = True
414
+ _no_split_modules = ["Qwen3DecoderLayer"]
415
+ _skip_keys_device_placement = ["past_key_values"]
416
+ _supports_flash_attn_2 = True
417
+ _supports_sdpa = True
418
+ _supports_flex_attn = True
419
+ _supports_cache_class = True
420
+ _supports_quantized_cache = True
421
+ _supports_static_cache = True
422
+ _supports_attention_backend = True
423
+
424
+ def _init_weights(self, module):
425
+ std = self.config.initializer_range
426
+ if isinstance(module, nn.Linear):
427
+ module.weight.data.normal_(mean=0.0, std=std)
428
+ if module.bias is not None:
429
+ module.bias.data.zero_()
430
+ elif isinstance(module, nn.Embedding):
431
+ module.weight.data.normal_(mean=0.0, std=std)
432
+ if module.padding_idx is not None:
433
+ module.weight.data[module.padding_idx].zero_()
434
+
435
+
436
+ QWEN3_INPUTS_DOCSTRING = r"""
437
+ Args:
438
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
439
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
440
+ it.
441
+
442
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
443
+ [`PreTrainedTokenizer.__call__`] for details.
444
+
445
+ [What are input IDs?](../glossary#input-ids)
446
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
447
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
448
+
449
+ - 1 for tokens that are **not masked**,
450
+ - 0 for tokens that are **masked**.
451
+
452
+ [What are attention masks?](../glossary#attention-mask)
453
+
454
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
455
+ [`PreTrainedTokenizer.__call__`] for details.
456
+
457
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
458
+ `past_key_values`).
459
+
460
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
461
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
462
+ information on the default strategy.
463
+
464
+ - 1 indicates the head is **not masked**,
465
+ - 0 indicates the head is **masked**.
466
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
467
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
468
+ config.n_positions - 1]`.
469
+
470
+ [What are position IDs?](../glossary#position-ids)
471
+ past_key_values (`Cache`, *optional*):
472
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
473
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
474
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
475
+
476
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
477
+
478
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
479
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
480
+ of shape `(batch_size, sequence_length)`.
481
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
482
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
483
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
484
+ model's internal embedding lookup matrix.
485
+ use_cache (`bool`, *optional*):
486
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
487
+ `past_key_values`).
488
+ output_attentions (`bool`, *optional*):
489
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
490
+ tensors for more detail.
491
+ output_hidden_states (`bool`, *optional*):
492
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
493
+ more detail.
494
+ return_dict (`bool`, *optional*):
495
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
496
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
497
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
498
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
499
+ the complete sequence length.
500
+ """
501
+
502
+
503
+ @add_start_docstrings(
504
+ "The bare Qwen3 Model outputting raw hidden-states without any specific head on top.",
505
+ QWEN3_START_DOCSTRING,
506
+ )
507
+ class Qwen3Model(Qwen3PreTrainedModel):
508
+ """
509
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen3DecoderLayer`]
510
+
511
+ Args:
512
+ config: Qwen3Config
513
+ """
514
+
515
+ def __init__(self, config: Qwen3ConfigGating):
516
+ super().__init__(config)
517
+ self.padding_idx = config.pad_token_id
518
+ self.vocab_size = config.vocab_size
519
+
520
+ # qr_start_layer <= layer_idx < qr_end_layer: layers satisfying this condition MAY contain QR heads
521
+ self.qr_start_layer = config.qr_start_layer
522
+ self.qr_end_layer = config.qr_end_layer
523
+ self.qr_head_list = config.qr_head_list
524
+ self.qr_head_list_mapped = config.qr_head_list_mapped
525
+
526
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
527
+ # self.layers = nn.ModuleList(
528
+ # [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
529
+ # )
530
+ self.layers = nn.ModuleList(
531
+ [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(self.qr_end_layer)]
532
+ )
533
+ # no need to normalize the output of the last layer because it is not used
534
+ # self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
535
+ self.rotary_emb = Qwen3RotaryEmbedding(config=config)
536
+ self.gradient_checkpointing = False
537
+
538
+ # Initialize weights and apply final processing
539
+ self.post_init()
540
+
541
+ def get_input_embeddings(self):
542
+ return self.embed_tokens
543
+
544
+ def set_input_embeddings(self, value):
545
+ self.embed_tokens = value
546
+
547
+ @can_return_tuple
548
+ @add_start_docstrings_to_model_forward(QWEN3_INPUTS_DOCSTRING)
549
+ def forward(
550
+ self,
551
+ input_ids: Optional[torch.LongTensor] = None,
552
+ attention_mask: Optional[torch.Tensor] = None,
553
+ position_ids: Optional[torch.LongTensor] = None,
554
+ past_key_values: Optional[Cache] = None,
555
+ inputs_embeds: Optional[torch.FloatTensor] = None,
556
+ use_cache: Optional[bool] = None,
557
+ output_attentions: Optional[bool] = None,
558
+ output_hidden_states: Optional[bool] = None,
559
+ cache_position: Optional[torch.LongTensor] = None,
560
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
561
+ ) -> BaseModelOutputWithPast:
562
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
563
+ output_hidden_states = (
564
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
565
+ )
566
+ # this model only used for inference
567
+ # at this version, only batch_size=1 is supported
568
+ use_cache = True
569
+ self.training = False
570
+ self.gradient_checkpointing = False
571
+ # use_cache = use_cache if use_cache is not None else self.config.use_cache
572
+
573
+ if (input_ids is None) ^ (inputs_embeds is not None):
574
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
575
+
576
+ if self.gradient_checkpointing and self.training and use_cache:
577
+ logger.warning_once(
578
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
579
+ )
580
+ use_cache = False
581
+
582
+ # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
583
+ if not isinstance(past_key_values, (type(None), Cache)):
584
+ raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
585
+
586
+ if inputs_embeds is None:
587
+ inputs_embeds = self.embed_tokens(input_ids)
588
+
589
+ if use_cache and past_key_values is None:
590
+ past_key_values = DynamicCache()
591
+
592
+ if cache_position is None:
593
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
594
+ cache_position = torch.arange(
595
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
596
+ )
597
+
598
+ if position_ids is None:
599
+ position_ids = cache_position.unsqueeze(0)
600
+
601
+ causal_mask = self._update_causal_mask(
602
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
603
+ )
604
+
605
+ hidden_states = inputs_embeds
606
+
607
+ # create position embeddings to be shared across the decoder layers
608
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
609
+
610
+ # decoder layers
611
+ all_hidden_states = () if output_hidden_states else None
612
+ all_self_attns = () if output_attentions else None
613
+
614
+ # stop if upper layers contain no qr head
615
+ # for decoder_layer in self.layers[: self.config.num_hidden_layers]:
616
+ for decoder_layer in self.layers[: self.qr_end_layer]:
617
+ if output_hidden_states:
618
+ all_hidden_states += (hidden_states,)
619
+
620
+ if self.gradient_checkpointing and self.training:
621
+ layer_outputs = self._gradient_checkpointing_func(
622
+ partial(decoder_layer.__call__, **flash_attn_kwargs),
623
+ hidden_states,
624
+ causal_mask,
625
+ position_ids,
626
+ past_key_values,
627
+ output_attentions,
628
+ use_cache,
629
+ cache_position,
630
+ position_embeddings,
631
+ )
632
+ else:
633
+ layer_outputs = decoder_layer(
634
+ hidden_states,
635
+ attention_mask=causal_mask,
636
+ position_ids=position_ids,
637
+ past_key_value=past_key_values,
638
+ output_attentions=output_attentions,
639
+ use_cache=use_cache,
640
+ cache_position=cache_position,
641
+ position_embeddings=position_embeddings,
642
+ **flash_attn_kwargs,
643
+ )
644
+
645
+ hidden_states = layer_outputs[0]
646
+
647
+ if output_attentions:
648
+ all_self_attns += (layer_outputs[1],)
649
+
650
+ # not used
651
+ # hidden_states = self.norm(hidden_states)
652
+
653
+ # add hidden states from the last decoder layer
654
+ if output_hidden_states:
655
+ all_hidden_states += (hidden_states,)
656
+
657
+ return BaseModelOutputWithPast(
658
+ last_hidden_state=hidden_states,
659
+ past_key_values=past_key_values if use_cache else None,
660
+ hidden_states=all_hidden_states,
661
+ attentions=all_self_attns,
662
+ )
663
+
664
+ def _update_causal_mask(
665
+ self,
666
+ attention_mask: torch.Tensor,
667
+ input_tensor: torch.Tensor,
668
+ cache_position: torch.Tensor,
669
+ past_key_values: Cache,
670
+ output_attentions: bool = False,
671
+ ):
672
+ if self.config._attn_implementation == "flash_attention_2":
673
+ if attention_mask is not None and past_key_values is not None:
674
+ is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
675
+ if is_padding_right:
676
+ raise ValueError(
677
+ "You are attempting to perform batched generation with padding_side='right'"
678
+ " this may lead to unexpected behaviour for Flash Attention version of Qwen3. Make sure to "
679
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
680
+ )
681
+ if attention_mask is not None and 0.0 in attention_mask:
682
+ return attention_mask
683
+ return None
684
+
685
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
686
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
687
+ # to infer the attention mask.
688
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
689
+ using_static_cache = isinstance(past_key_values, StaticCache)
690
+ using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
691
+
692
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
693
+ if (
694
+ self.config._attn_implementation == "sdpa"
695
+ and not (using_static_cache or using_sliding_window_cache)
696
+ and not output_attentions
697
+ ):
698
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
699
+ attention_mask,
700
+ inputs_embeds=input_tensor,
701
+ past_key_values_length=past_seen_tokens,
702
+ sliding_window=self.config.sliding_window,
703
+ is_training=self.training,
704
+ ):
705
+ return None
706
+
707
+ dtype, device = input_tensor.dtype, input_tensor.device
708
+ min_dtype = torch.finfo(dtype).min
709
+ sequence_length = input_tensor.shape[1]
710
+ # SlidingWindowCache or StaticCache
711
+ if using_sliding_window_cache or using_static_cache:
712
+ target_length = past_key_values.get_max_cache_shape()
713
+ # DynamicCache or no cache
714
+ else:
715
+ target_length = (
716
+ attention_mask.shape[-1]
717
+ if isinstance(attention_mask, torch.Tensor)
718
+ else past_seen_tokens + sequence_length + 1
719
+ )
720
+
721
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
722
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
723
+ attention_mask,
724
+ sequence_length=sequence_length,
725
+ target_length=target_length,
726
+ dtype=dtype,
727
+ device=device,
728
+ cache_position=cache_position,
729
+ batch_size=input_tensor.shape[0],
730
+ config=self.config,
731
+ past_key_values=past_key_values,
732
+ )
733
+
734
+ if (
735
+ self.config._attn_implementation == "sdpa"
736
+ and attention_mask is not None
737
+ and attention_mask.device.type in ["cuda", "xpu"]
738
+ and not output_attentions
739
+ ):
740
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
741
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
742
+ # Details: https://github.com/pytorch/pytorch/issues/110213
743
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
744
+
745
+ return causal_mask
746
+
747
+ @staticmethod
748
+ def _prepare_4d_causal_attention_mask_with_cache_position(
749
+ attention_mask: torch.Tensor,
750
+ sequence_length: int,
751
+ target_length: int,
752
+ dtype: torch.dtype,
753
+ device: torch.device,
754
+ cache_position: torch.Tensor,
755
+ batch_size: int,
756
+ config: Qwen3Config,
757
+ past_key_values: Cache,
758
+ ):
759
+ """
760
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
761
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
762
+
763
+ Args:
764
+ attention_mask (`torch.Tensor`):
765
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
766
+ sequence_length (`int`):
767
+ The sequence length being processed.
768
+ target_length (`int`):
769
+ The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
770
+ dtype (`torch.dtype`):
771
+ The dtype to use for the 4D attention mask.
772
+ device (`torch.device`):
773
+ The device to place the 4D attention mask on.
774
+ cache_position (`torch.Tensor`):
775
+ Indices depicting the position of the input sequence tokens in the sequence.
776
+ batch_size (`torch.Tensor`):
777
+ Batch size.
778
+ config (`Qwen3Config`):
779
+ The model's configuration class
780
+ past_key_values (`Cache`):
781
+ The cache class that is being used currently to generate
782
+ """
783
+ if attention_mask is not None and attention_mask.dim() == 4:
784
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
785
+ causal_mask = attention_mask
786
+ else:
787
+ min_dtype = torch.finfo(dtype).min
788
+ causal_mask = torch.full(
789
+ (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
790
+ )
791
+ diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
792
+ if config.sliding_window is not None:
793
+ # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
794
+ # the check is needed to verify is current checkpoint was trained with sliding window or not
795
+ if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
796
+ sliding_attend_mask = torch.arange(target_length, device=device) <= (
797
+ cache_position.reshape(-1, 1) - config.sliding_window
798
+ )
799
+ diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
800
+ causal_mask *= diagonal_attend_mask
801
+ causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
802
+ if attention_mask is not None:
803
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
804
+ if attention_mask.shape[-1] > target_length:
805
+ attention_mask = attention_mask[:, :target_length]
806
+ mask_length = attention_mask.shape[-1]
807
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
808
+ causal_mask.device
809
+ )
810
+ padding_mask = padding_mask == 0
811
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
812
+ padding_mask, min_dtype
813
+ )
814
+ return causal_mask
815
+
816
+
817
+ Qwen3ConfigGating.register_for_auto_class()
818
+ Qwen3Model.register_for_auto_class("AutoModel")
819
+
820
+
821
+ __all__ = [
822
+ "Qwen3ConfigGating",
823
+ "Qwen3Model",
824
+ "Qwen3PreTrainedModel",
825
+ ]
special_tokens_map.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|object_ref_start|>",
6
+ "<|object_ref_end|>",
7
+ "<|box_start|>",
8
+ "<|box_end|>",
9
+ "<|quad_start|>",
10
+ "<|quad_end|>",
11
+ "<|vision_start|>",
12
+ "<|vision_end|>",
13
+ "<|vision_pad|>",
14
+ "<|image_pad|>",
15
+ "<|video_pad|>"
16
+ ],
17
+ "eos_token": {
18
+ "content": "<|im_end|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ },
24
+ "pad_token": {
25
+ "content": "<|endoftext|>",
26
+ "lstrip": false,
27
+ "normalized": false,
28
+ "rstrip": false,
29
+ "single_word": false
30
+ }
31
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3c3dfe474a8bbe89b0e83627fd9ff784ad71027f12fd8c618708c818e808789d
3
+ size 11422648
tokenizer_config.json ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "added_tokens_decoder": {
5
+ "151643": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "151644": {
14
+ "content": "<|im_start|>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "151645": {
22
+ "content": "<|im_end|>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "151646": {
30
+ "content": "<|object_ref_start|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "151647": {
38
+ "content": "<|object_ref_end|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "151648": {
46
+ "content": "<|box_start|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "151649": {
54
+ "content": "<|box_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "151650": {
62
+ "content": "<|quad_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "151651": {
70
+ "content": "<|quad_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "151652": {
78
+ "content": "<|vision_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ },
85
+ "151653": {
86
+ "content": "<|vision_end|>",
87
+ "lstrip": false,
88
+ "normalized": false,
89
+ "rstrip": false,
90
+ "single_word": false,
91
+ "special": true
92
+ },
93
+ "151654": {
94
+ "content": "<|vision_pad|>",
95
+ "lstrip": false,
96
+ "normalized": false,
97
+ "rstrip": false,
98
+ "single_word": false,
99
+ "special": true
100
+ },
101
+ "151655": {
102
+ "content": "<|image_pad|>",
103
+ "lstrip": false,
104
+ "normalized": false,
105
+ "rstrip": false,
106
+ "single_word": false,
107
+ "special": true
108
+ },
109
+ "151656": {
110
+ "content": "<|video_pad|>",
111
+ "lstrip": false,
112
+ "normalized": false,
113
+ "rstrip": false,
114
+ "single_word": false,
115
+ "special": true
116
+ },
117
+ "151657": {
118
+ "content": "<tool_call>",
119
+ "lstrip": false,
120
+ "normalized": false,
121
+ "rstrip": false,
122
+ "single_word": false,
123
+ "special": false
124
+ },
125
+ "151658": {
126
+ "content": "</tool_call>",
127
+ "lstrip": false,
128
+ "normalized": false,
129
+ "rstrip": false,
130
+ "single_word": false,
131
+ "special": false
132
+ },
133
+ "151659": {
134
+ "content": "<|fim_prefix|>",
135
+ "lstrip": false,
136
+ "normalized": false,
137
+ "rstrip": false,
138
+ "single_word": false,
139
+ "special": false
140
+ },
141
+ "151660": {
142
+ "content": "<|fim_middle|>",
143
+ "lstrip": false,
144
+ "normalized": false,
145
+ "rstrip": false,
146
+ "single_word": false,
147
+ "special": false
148
+ },
149
+ "151661": {
150
+ "content": "<|fim_suffix|>",
151
+ "lstrip": false,
152
+ "normalized": false,
153
+ "rstrip": false,
154
+ "single_word": false,
155
+ "special": false
156
+ },
157
+ "151662": {
158
+ "content": "<|fim_pad|>",
159
+ "lstrip": false,
160
+ "normalized": false,
161
+ "rstrip": false,
162
+ "single_word": false,
163
+ "special": false
164
+ },
165
+ "151663": {
166
+ "content": "<|repo_name|>",
167
+ "lstrip": false,
168
+ "normalized": false,
169
+ "rstrip": false,
170
+ "single_word": false,
171
+ "special": false
172
+ },
173
+ "151664": {
174
+ "content": "<|file_sep|>",
175
+ "lstrip": false,
176
+ "normalized": false,
177
+ "rstrip": false,
178
+ "single_word": false,
179
+ "special": false
180
+ },
181
+ "151665": {
182
+ "content": "<tool_response>",
183
+ "lstrip": false,
184
+ "normalized": false,
185
+ "rstrip": false,
186
+ "single_word": false,
187
+ "special": false
188
+ },
189
+ "151666": {
190
+ "content": "</tool_response>",
191
+ "lstrip": false,
192
+ "normalized": false,
193
+ "rstrip": false,
194
+ "single_word": false,
195
+ "special": false
196
+ },
197
+ "151667": {
198
+ "content": "<think>",
199
+ "lstrip": false,
200
+ "normalized": false,
201
+ "rstrip": false,
202
+ "single_word": false,
203
+ "special": false
204
+ },
205
+ "151668": {
206
+ "content": "</think>",
207
+ "lstrip": false,
208
+ "normalized": false,
209
+ "rstrip": false,
210
+ "single_word": false,
211
+ "special": false
212
+ }
213
+ },
214
+ "additional_special_tokens": [
215
+ "<|im_start|>",
216
+ "<|im_end|>",
217
+ "<|object_ref_start|>",
218
+ "<|object_ref_end|>",
219
+ "<|box_start|>",
220
+ "<|box_end|>",
221
+ "<|quad_start|>",
222
+ "<|quad_end|>",
223
+ "<|vision_start|>",
224
+ "<|vision_end|>",
225
+ "<|vision_pad|>",
226
+ "<|image_pad|>",
227
+ "<|video_pad|>"
228
+ ],
229
+ "bos_token": null,
230
+ "clean_up_tokenization_spaces": false,
231
+ "eos_token": "<|im_end|>",
232
+ "errors": "replace",
233
+ "extra_special_tokens": {},
234
+ "model_max_length": 1010000,
235
+ "pad_token": "<|endoftext|>",
236
+ "padding_side": "left",
237
+ "split_special_tokens": false,
238
+ "tokenizer_class": "Qwen2Tokenizer",
239
+ "unk_token": null
240
+ }
vocab.json ADDED
The diff for this file is too large to render. See raw diff