Upload folder using huggingface_hub
Browse files- .gitattributes +1 -0
- README.md +226 -3
- config.json +42 -0
- configuration_klear.py +232 -0
- generation_config.json +6 -0
- merges.txt +0 -0
- model-00001-of-00019.safetensors +3 -0
- model-00002-of-00019.safetensors +3 -0
- model-00003-of-00019.safetensors +3 -0
- model-00004-of-00019.safetensors +3 -0
- model-00005-of-00019.safetensors +3 -0
- model-00006-of-00019.safetensors +3 -0
- model-00007-of-00019.safetensors +3 -0
- model-00008-of-00019.safetensors +3 -0
- model-00009-of-00019.safetensors +3 -0
- model-00010-of-00019.safetensors +3 -0
- model-00011-of-00019.safetensors +3 -0
- model-00012-of-00019.safetensors +3 -0
- model-00013-of-00019.safetensors +3 -0
- model-00014-of-00019.safetensors +3 -0
- model-00015-of-00019.safetensors +3 -0
- model-00016-of-00019.safetensors +3 -0
- model-00017-of-00019.safetensors +3 -0
- model-00018-of-00019.safetensors +3 -0
- model-00019-of-00019.safetensors +3 -0
- model.safetensors.index.json +0 -0
- modeling_klear.py +763 -0
- tokenizer.json +3 -0
- tokenizer_config.json +239 -0
- vocab.json +0 -0
.gitattributes
CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
36 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
@@ -1,3 +1,226 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Klear
|
2 |
+
|
3 |
+
<div align="center">
|
4 |
+
<img src="figures/klear-logo-02.png" width="500"/>
|
5 |
+
<p>
|
6 |
+
🤗 <a href="https://huggingface.co/Kwai-Klear">Hugging Face</a> | 📑 <a href="">Technique Report</a>
|
7 |
+
<br>
|
8 |
+
🖥️ <a href="https://kml-dtmachine-15498-prod-1.kmlhb2az1l3-2.corp.kuaishou.com">Chat with Klear</a> | 💬 <a href="https://github.com/Kwai-Klear">Issues & Discussions</a>
|
9 |
+
</p>
|
10 |
+
</div>
|
11 |
+
|
12 |
+
|
13 |
+
## 🔥News
|
14 |
+
|
15 |
+
- 2025.09.05: We released `Klear-46B-A2.5B` series. Currently, Klear-46B-A2.5B offers two versions: `a base model` and an advanced version that includes `instruction tuned` model. Additionally, `an reasoning version is currently in training`. Please stay tuned for more updates.
|
16 |
+
|
17 |
+
|
18 |
+
## 1. Introduction
|
19 |
+
|
20 |
+
|
21 |
+
`Klear-46B-A2.5B` is a sparse Mixture-of-Experts (MoE) large language model developed by **the Kwai-Klear Team at Kuaishou**, designed to deliver both **high performance** and **inference efficiency**. It features **256 experts**, with only **8 activated** per forward pass, resulting in **46 billion total parameters** but just **2.5 billion active** — achieving dense-level performance at a fraction of the computational cost.
|
22 |
+
|
23 |
+
The model was trained on over **22 trillion tokens** using a **three-stage progressive curriculum**:
|
24 |
+
|
25 |
+
**1. Foundational Knowledge Learning (12T tokens):**
|
26 |
+
General-purpose datasets such as CommonCrawl were processed with stratified quality filters, following a curriculum learning strategy that progresses from lower to higher data quality.
|
27 |
+
|
28 |
+
**2. Data Complexity Enhancement (8T tokens):**
|
29 |
+
The proportion of mathematical, coding, and STEM-related data was gradually increased to strengthen the model's reasoning and problem-solving capabilities.
|
30 |
+
|
31 |
+
**3. Reasoning Enhancement and Longcontext Stage (2T tokens):**
|
32 |
+
Training focused on synthetic and reasoning-intensive data, combined with a fast learning rate annealing strategy to maximize data efficiency and optimize final performance.
|
33 |
+
|
34 |
+
As a result, Klear-46B-A2.5B-Base matches or surpasses the performance of dense models with several times more active parameters, while offering significantly better efficiency and cost-effectiveness for real-world deployment.
|
35 |
+
|
36 |
+
|
37 |
+
## Model Summary
|
38 |
+
|
39 |
+
this repo contains the base and instruction-tuned model**. which has the following architecture:
|
40 |
+
|
41 |
+
| **propoty** | **value** |
|
42 |
+
|---------------------------|------------------------------------------------------------------------|
|
43 |
+
| hidden_size | 2048 |
|
44 |
+
| moe_intermediate_size | 896 |
|
45 |
+
| n_shared_experts | 1 |
|
46 |
+
| num_attention_heads | 32 |
|
47 |
+
| num_experts | 256 |
|
48 |
+
| num_experts_per_tok | 8 |
|
49 |
+
| num_hidden_layers | 32 |
|
50 |
+
| num_key_value_heads | 4 |
|
51 |
+
| vocab_size | 151936 |
|
52 |
+
| tie_word_embeddings | false |
|
53 |
+
| context length | 65536 |
|
54 |
+
|
55 |
+
|
56 |
+
### Model Downloads
|
57 |
+
|
58 |
+
<div align="center">
|
59 |
+
|
60 |
+
| **Model** | **#Total Params** | **#Activated Params** | **Context Length** | **Download Link** |
|
61 |
+
| :------------: | :------------: | :------------: | :------------: | :------------: |
|
62 |
+
| Klear-46B-A2.5B-Base | 46B | 2.5B | 64K | [🤗 Hugging Face](https://huggingface.co/Kwai-Klear) |
|
63 |
+
| Klear-46B-A2.5B-Inst. | 46B | 2.5B | 64K | [🤗 Hugging Face](https://huggingface.co/Kwai-Klear) |
|
64 |
+
|
65 |
+
</div>
|
66 |
+
|
67 |
+
|
68 |
+
## 2. Benchmark Evaluation
|
69 |
+
### Klear-46B-A2.5B-Base Evaluation Results
|
70 |
+
| Ability | Benchmark | Klear-46B-A2.5B-Base | MiMO-7B-Base | Qwen3-8B-BASE | Qwen3-14B-BASE | Ling-lite-1.5-Base | Qwen3-30B-A3B-BASE |
|
71 |
+
| ----------- | ---------------------- | -------------------- | ------------ | ------------- | -------------- | ------------------ | ------------------ |
|
72 |
+
| | # Total Params | 46B | 7B | 8B | 14B | 16.8B | 30B |
|
73 |
+
| | # Activated Params | 2.5B | 7B | 8B | 14B | 2.75B | 3B |
|
74 |
+
| **Code** | HumanEval (0-shot*) | 89 | - | 84.1 | 87.8 | 83.5 | 90.9 |
|
75 |
+
| | MBPP (3-shot) | 76 | 55.2 | 69 | 74 | 66.6 | 75.6 |
|
76 |
+
| **Math** | MATH (4-shot, cot) | 55.7 | 36.78 | 58.4 | 57.1 | 56.98 | 57.6 |
|
77 |
+
| | CMATH (3-shot) | 87.8 | 78.5 | 88.3 | 90.7 | 85.7 | 89.7 |
|
78 |
+
| | GSM8K (4-shot, cot) | 87.3 | 78.47 | 89.4 | 90.3 | 87.6 | 91.1 |
|
79 |
+
| **General** | MMLU-Pro (5-shot, cot) | 57.6 | 43.1 | 55.2 | 58.1 | 49.9 | 58.8 |
|
80 |
+
| | MMLU (5-shot) | 80.5 | 69.24 | 77.1 | 80.6 | 73.7 | 80.4 |
|
81 |
+
| | CEval (5-shot) | 89.8 | 67.98 | 81.9 | 84.8 | 78.2 | 87.4 |
|
82 |
+
| | CMMLU (5-shot) | 88 | 70.79 | 82 | 85.6 | 81.2 | 87.1 |
|
83 |
+
| | GPQA (0-shot) | 35.3 | 31.03 | 33.9 | 35.7 | 30.1 | 35.5 |
|
84 |
+
| | AGIEval (0-shot) | 52.3 | 48.3* | 51.7 | 55.7 | 54.3 | 56 |
|
85 |
+
| | BBH (3-shot, cot) | 77.9 | 75.6 | 78.1 | 80.1 | 75.4 | 81.2 |
|
86 |
+
| **Others** | HellaSwag (0-shot) | 80.5 | 80* | 78.7 | 81.5 | 80 | 81.2 |
|
87 |
+
| | Winogrande (3-shot) | 78.8 | 78* | 73.6 | 78.5 | 72.1 | 77.9 |
|
88 |
+
| | Triviaqa (5-shot) | 69.6 | 60.8* | 56.3 | 62.1 | 60.9 | 65.6 |
|
89 |
+
| | Naturalqs (5-shot) | 37.5 | 23.46 | 25.7 | 29.1 | 28 | 30.7 |
|
90 |
+
| | PIQA (0-shot) | 81.6 | 80.14 | 79.5 | 81.9 | 82 | 80.7 |
|
91 |
+
| | SIQA (0-shot) | 67.9 | 51.74 | 56.2 | 58.4 | 56.3 | 56.3 |
|
92 |
+
| | OpenBookQA (0-shot) | 37.8 | 34.2 | 35 | 35.6 | 38.2 | 34.6 |
|
93 |
+
|
94 |
+
Note:
|
95 |
+
1. `*`During pretraining, we found that the HumanEval metric fluctuated significantly and was extremely sensitive to formatting. Therefore, we referred to the prompt from Ling-series paper to modify the original HumanEval. The results in the table are the evaluation metrics after this modification.
|
96 |
+
2. For Mimo-base-7B, the results marked with `*` are sourced from other public reports.
|
97 |
+
|
98 |
+
### Klear-46B-A2.5B-Inst. Evaluation Results
|
99 |
+
| Ability | Benchmark | Klear-46B-A2.5B-Inst. | MiniCPM4-8B | Qwen3-8B (NoThink) | gemma3-12b-it | Phi4-14B |
|
100 |
+
| ------------------------- | --------------------------- | --------------- | ----------- | ------------------ | ------------- | -------- |
|
101 |
+
| | # Total Params | 46B | 8B | 8B | 12B | 14B |
|
102 |
+
| | # Activated Params | 2.5B | 8B | 8B | 12B | 14B |
|
103 |
+
| **English Understanding** | MMLU-Redux | 82.23 | 77.63 | 79.32 | 78.39 | 83.09 |
|
104 |
+
| | MMLU-Pro | 64.82 | 54.69 | 63.8 | 60.69 | 67.25 |
|
105 |
+
| | GPQA-Diamoind | 49.49 | 38.51 | 51.77 | 39.02 | 59.47 |
|
106 |
+
| | SimpleQA | 5.94 | 3.51 | 5.5 | 6.22 | 3.28 |
|
107 |
+
| **Chinese Understanding** | CLUEWSC | 88.82 | 81.91 | 82.89 | 91.12 | 88.16 |
|
108 |
+
| | CEval | 84.29 | 81.78 | 81.66 | 60.81 | 64.79 |
|
109 |
+
| | C-SimpleQA | 42.03 | 23.13 | 37.07 | 28.97 | 24.77 |
|
110 |
+
| **Math & Reasoning** | MATH500 | 86.4 | 79.8 | 85 | 86.8 | 80.6 |
|
111 |
+
| | AIME24 | 30.42 | 22.92 | 28.33 | 23.96 | 15.83 |
|
112 |
+
| | AIME25 | 21.04 | 15.21 | 20.62 | 18.33 | 18.75 |
|
113 |
+
| | ZebraLogic | 46.4 | 8.5 | 25.7 | 18 | 30.3 |
|
114 |
+
| **Code** | HumanEval | 89.63 | 74.39 | 83.54 | 82.32 | 85.37 |
|
115 |
+
| | HumanEval+ | 87.2 | 70.12 | 76.83 | 75.61 | 83.54 |
|
116 |
+
| | MBPPEvalplus | 79.6 | 82 | 76.2 | 85.7 | 77.5 |
|
117 |
+
| | MBPPEvalplus++ | 68.5 | 69.3 | 66.1 | 74.1 | 66.7 |
|
118 |
+
| | LiveCodeBench v5(2408-2501) | 29.75 | 12.19 | 27.24 | 24.73 | 23.66 |
|
119 |
+
| **Instruction Following** | IF-Eval | 80.41 | 73.01 | 84.47 | 81.52 | 59.33 |
|
120 |
+
| | Multi-IF(en+zh) | 78.25 | 61.79 | 78.95 | 76.56 | 62.7 |
|
121 |
+
| **Comprehensive Ability** | MTBench | 8.03 | 6.875 | 8.21 | 8.675 | 8.625 |
|
122 |
+
| | MT-Eval | 8.1 | 6.7 | 8.18 | 8.45 | 8.12 |
|
123 |
+
| | Arena-Hard v2 | 19.8 | 2.2 | 19.8 | 50 | 9.6 |
|
124 |
+
| | AlignBench v1.1 | 6.8 | 5.99 | 6.95 | 6.3 | 6.33 |
|
125 |
+
| | LiveBench 1125 | 48.7 | 25.5 | 52.1 | 43.1 | 40 |
|
126 |
+
|
127 |
+
## 3. Quick start
|
128 |
+
|
129 |
+
### Inference with huggingface
|
130 |
+
|
131 |
+
You can now inference in Transformers starting from version `4.56.0`.
|
132 |
+
|
133 |
+
#### Klear-46B-A2.5B-Base
|
134 |
+
|
135 |
+
```python
|
136 |
+
import torch
|
137 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
138 |
+
|
139 |
+
model_name = "/path/to/Klear-Base"
|
140 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
141 |
+
|
142 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", dtype=torch.bfloat16, trust_remote_code=True)
|
143 |
+
|
144 |
+
text = "世界上最大的湖是"
|
145 |
+
inputs = tokenizer(text, return_tensors="pt")
|
146 |
+
outputs = model.generate(**inputs.to(model.device), max_new_tokens=256)
|
147 |
+
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
148 |
+
print(result)
|
149 |
+
```
|
150 |
+
|
151 |
+
#### Klear-46B-A2.5B-Inst.
|
152 |
+
|
153 |
+
```python
|
154 |
+
import torch
|
155 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
|
156 |
+
|
157 |
+
model_name = "/path/to/Klear-Inst."
|
158 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
159 |
+
|
160 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16)
|
161 |
+
|
162 |
+
messages = [
|
163 |
+
{"role": "user", "content": "帮我用 python 写一个计算器的代码吧。"}
|
164 |
+
]
|
165 |
+
input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
|
166 |
+
outputs = model.generate(input_tensor.to(model.device), max_new_tokens=1024)
|
167 |
+
|
168 |
+
result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
|
169 |
+
print(result)
|
170 |
+
```
|
171 |
+
|
172 |
+
### Inference with vllm
|
173 |
+
|
174 |
+
[vLLM](https://github.com/vllm-project/vllm) is a high-speed and memery-efficicent inference framework. We provide our own forked version of [vLLM](https://github.com/vllm-project/vllm) here.
|
175 |
+
|
176 |
+
```shell
|
177 |
+
git clone
|
178 |
+
cd vllm
|
179 |
+
pip install
|
180 |
+
vllm serve Klear-46B-A2.5B-inst --port 8000 --tensor-parallel-size 8 --trust-remote-code
|
181 |
+
```
|
182 |
+
|
183 |
+
An OpenAI-compatible API will be available at `http://localhost:8000/v1`.
|
184 |
+
|
185 |
+
Or you can refer to the following Python script for offline inference
|
186 |
+
```python
|
187 |
+
from vllm import LLM, SamplingParams
|
188 |
+
|
189 |
+
model_path = "/path/to/Klear"
|
190 |
+
llm = LLM(
|
191 |
+
model=model_path,
|
192 |
+
trust_remote_code=True,
|
193 |
+
num_speculative_tokens=1,
|
194 |
+
disable_log_stats=False
|
195 |
+
)
|
196 |
+
sampling_params = SamplingParams(temperature=0.2)
|
197 |
+
|
198 |
+
conversation = [
|
199 |
+
{
|
200 |
+
"role": "system",
|
201 |
+
"content": ""
|
202 |
+
},
|
203 |
+
{
|
204 |
+
"role": "user",
|
205 |
+
"content": "Please help me write a snake game code.",
|
206 |
+
},
|
207 |
+
]
|
208 |
+
|
209 |
+
outputs = llm.chat(conversation,
|
210 |
+
sampling_params=sampling_params,
|
211 |
+
use_tqdm=False)
|
212 |
+
|
213 |
+
for idx, output in enumerate(outputs):
|
214 |
+
prompt = output.prompt
|
215 |
+
generated_text = output.outputs[0].text
|
216 |
+
print(f"==== Response #{idx} ====")
|
217 |
+
print(f"Prompt: {prompt}, Generated text: {generated_text}")
|
218 |
+
|
219 |
+
```
|
220 |
+
|
221 |
+
## Citation
|
222 |
+
|
223 |
+
If you find `Klear-46B-A2.5B` is useful or want to use in your projects, please kindly cite our paper:
|
224 |
+
|
225 |
+
```
|
226 |
+
```
|
config.json
ADDED
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"KlearMoeForCausalLM"
|
4 |
+
],
|
5 |
+
"attention_bias": false,
|
6 |
+
"attention_dropout": 0.0,
|
7 |
+
"auto_map": {
|
8 |
+
"AutoConfig": "configuration_klear.KlearConfig",
|
9 |
+
"AutoModel": "modeling_klear.KlearModel",
|
10 |
+
"AutoModelForCausalLM": "modeling_klear.KlearMoeForCausalLM"
|
11 |
+
},
|
12 |
+
"decoder_sparse_step": 1,
|
13 |
+
"dtype": "bfloat16",
|
14 |
+
"hidden_act": "silu",
|
15 |
+
"hidden_size": 2048,
|
16 |
+
"initializer_range": 0.02,
|
17 |
+
"intermediate_size": 8064,
|
18 |
+
"max_position_embeddings": 65536,
|
19 |
+
"mlp_only_layers": [],
|
20 |
+
"model_type": "Klear",
|
21 |
+
"moe_aux_loss_coeff": 0.0001,
|
22 |
+
"moe_intermediate_size": 896,
|
23 |
+
"n_shared_experts": 1,
|
24 |
+
"norm_topk_prob": true,
|
25 |
+
"num_attention_heads": 32,
|
26 |
+
"num_experts": 256,
|
27 |
+
"num_experts_per_tok": 8,
|
28 |
+
"num_hidden_layers": 32,
|
29 |
+
"num_key_value_heads": 4,
|
30 |
+
"output_router_logits": false,
|
31 |
+
"rms_norm_eps": 1e-05,
|
32 |
+
"rope_scaling": null,
|
33 |
+
"rope_theta": 500000.0,
|
34 |
+
"routed_scaling_factor": 2.5,
|
35 |
+
"router_aux_loss_coef": 0.001,
|
36 |
+
"sliding_window": null,
|
37 |
+
"tie_word_embeddings": false,
|
38 |
+
"transformers_version": "4.56.0",
|
39 |
+
"use_cache": true,
|
40 |
+
"use_sliding_window": false,
|
41 |
+
"vocab_size": 151936
|
42 |
+
}
|
configuration_klear.py
ADDED
@@ -0,0 +1,232 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
2 |
+
# This file was automatically generated from src/transformers/models/Klear/modular_Klear.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_Klear.py file directly. One of our CI enforces this.
|
6 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
7 |
+
|
8 |
+
|
9 |
+
from transformers.configuration_utils import PretrainedConfig
|
10 |
+
from transformers.modeling_rope_utils import rope_config_validation
|
11 |
+
|
12 |
+
|
13 |
+
class KlearConfig(PretrainedConfig):
|
14 |
+
r"""
|
15 |
+
This is the configuration class to store the configuration of a [`KlearModel`]. It is used to instantiate a
|
16 |
+
Klear model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
17 |
+
with the defaults will yield a similar configuration to that of [Klear-kwaii/Klear-MoE](https://huggingface.co/Klear/Klear-MoE).
|
18 |
+
|
19 |
+
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
20 |
+
documentation from [`PretrainedConfig`] for more information.
|
21 |
+
|
22 |
+
|
23 |
+
Args:
|
24 |
+
vocab_size (`int`, *optional*, defaults to 151936):
|
25 |
+
Vocabulary size of the Klear model. Defines the number of different tokens that can be represented by the
|
26 |
+
`inputs_ids` passed when calling [`KlearModel`]
|
27 |
+
hidden_size (`int`, *optional*, defaults to 2048):
|
28 |
+
Dimension of the hidden representations.
|
29 |
+
intermediate_size (`int`, *optional*, defaults to 6144):
|
30 |
+
Dimension of the MLP representations.
|
31 |
+
num_hidden_layers (`int`, *optional*, defaults to 24):
|
32 |
+
Number of hidden layers in the Transformer encoder.
|
33 |
+
num_attention_heads (`int`, *optional*, defaults to 32):
|
34 |
+
Number of attention heads for each attention layer in the Transformer encoder.
|
35 |
+
num_key_value_heads (`int`, *optional*, defaults to 4):
|
36 |
+
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
37 |
+
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
38 |
+
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
39 |
+
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
40 |
+
by meanpooling all the original heads within that group. For more details, check out [this
|
41 |
+
paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
|
42 |
+
|
43 |
+
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
44 |
+
The non-linear activation function (function or string) in the decoder.
|
45 |
+
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
46 |
+
The maximum sequence length that this model might ever be used with.
|
47 |
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
48 |
+
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
49 |
+
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
50 |
+
The epsilon used by the rms normalization layers.
|
51 |
+
use_cache (`bool`, *optional*, defaults to `True`):
|
52 |
+
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
53 |
+
relevant if `config.is_decoder=True`.
|
54 |
+
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
55 |
+
Whether the model's input and output word embeddings should be tied.
|
56 |
+
rope_theta (`float`, *optional*, defaults to 10000.0):
|
57 |
+
The base period of the RoPE embeddings.
|
58 |
+
rope_scaling (`Dict`, *optional*):
|
59 |
+
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
60 |
+
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
61 |
+
accordingly.
|
62 |
+
Expected contents:
|
63 |
+
`rope_type` (`str`):
|
64 |
+
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
65 |
+
'llama3'], with 'default' being the original RoPE implementation.
|
66 |
+
`factor` (`float`, *optional*):
|
67 |
+
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
68 |
+
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
69 |
+
original maximum pre-trained length.
|
70 |
+
`original_max_position_embeddings` (`int`, *optional*):
|
71 |
+
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
72 |
+
pretraining.
|
73 |
+
`attention_factor` (`float`, *optional*):
|
74 |
+
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
75 |
+
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
76 |
+
`factor` field to infer the suggested value.
|
77 |
+
`beta_fast` (`float`, *optional*):
|
78 |
+
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
79 |
+
ramp function. If unspecified, it defaults to 32.
|
80 |
+
`beta_slow` (`float`, *optional*):
|
81 |
+
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
82 |
+
ramp function. If unspecified, it defaults to 1.
|
83 |
+
`short_factor` (`list[float]`, *optional*):
|
84 |
+
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
85 |
+
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
86 |
+
size divided by the number of attention heads divided by 2
|
87 |
+
`long_factor` (`list[float]`, *optional*):
|
88 |
+
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
89 |
+
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
90 |
+
size divided by the number of attention heads divided by 2
|
91 |
+
`low_freq_factor` (`float`, *optional*):
|
92 |
+
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
93 |
+
`high_freq_factor` (`float`, *optional*):
|
94 |
+
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
95 |
+
attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
|
96 |
+
Whether to use a bias in the query, key, value and output projection layers during self-attention.
|
97 |
+
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
98 |
+
Whether to use sliding window attention.
|
99 |
+
sliding_window (`int`, *optional*, defaults to 4096):
|
100 |
+
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
|
101 |
+
attention_dropout (`float`, *optional*, defaults to 0.0):
|
102 |
+
The dropout ratio for the attention probabilities.
|
103 |
+
decoder_sparse_step (`int`, *optional*, defaults to 1):
|
104 |
+
The frequency of the MoE layer.
|
105 |
+
moe_intermediate_size (`int`, *optional*, defaults to 768):
|
106 |
+
Intermediate size of the routed expert.
|
107 |
+
num_experts_per_tok (`int`, *optional*, defaults to 8):
|
108 |
+
Number of selected experts.
|
109 |
+
num_experts (`int`, *optional*, defaults to 128):
|
110 |
+
Number of routed experts.
|
111 |
+
norm_topk_prob (`bool`, *optional*, defaults to `False`):
|
112 |
+
Whether to normalize the topk probabilities.
|
113 |
+
output_router_logits (`bool`, *optional*, defaults to `False`):
|
114 |
+
Whether or not the router logits should be returned by the model. Enabling this will also
|
115 |
+
allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
|
116 |
+
router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
|
117 |
+
The aux loss factor for the total loss.
|
118 |
+
mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
|
119 |
+
Indicate which layers use KlearMLP rather than KlearSparseMoeBlock
|
120 |
+
The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
|
121 |
+
If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
|
122 |
+
|
123 |
+
```python
|
124 |
+
>>> from transformers import KlearModel, KlearConfig
|
125 |
+
|
126 |
+
>>> # Initializing a Klear style configuration
|
127 |
+
>>> configuration = KlearConfig()
|
128 |
+
|
129 |
+
>>> # Initializing a model from the Klear-MoE" style configuration
|
130 |
+
>>> model = KlearModel(configuration)
|
131 |
+
|
132 |
+
>>> # Accessing the model configuration
|
133 |
+
>>> configuration = model.config
|
134 |
+
```"""
|
135 |
+
|
136 |
+
model_type = "Klear"
|
137 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
138 |
+
|
139 |
+
# Default tensor parallel plan for base model `Klear`
|
140 |
+
base_model_tp_plan = {
|
141 |
+
"layers.*.self_attn.q_proj": "colwise",
|
142 |
+
"layers.*.self_attn.k_proj": "colwise",
|
143 |
+
"layers.*.self_attn.v_proj": "colwise",
|
144 |
+
"layers.*.self_attn.o_proj": "rowwise",
|
145 |
+
"layers.*.mlp.experts.*.gate_proj": "colwise",
|
146 |
+
"layers.*.mlp.experts.*.up_proj": "colwise",
|
147 |
+
"layers.*.mlp.experts.*.down_proj": "rowwise",
|
148 |
+
"layers.*.mlp.gate_proj": "colwise",
|
149 |
+
"layers.*.mlp.up_proj": "colwise",
|
150 |
+
"layers.*.mlp.down_proj": "rowwise",
|
151 |
+
}
|
152 |
+
base_model_pp_plan = {
|
153 |
+
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
154 |
+
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
155 |
+
"norm": (["hidden_states"], ["hidden_states"]),
|
156 |
+
}
|
157 |
+
|
158 |
+
def __init__(
|
159 |
+
self,
|
160 |
+
vocab_size=151936,
|
161 |
+
hidden_size=2048,
|
162 |
+
intermediate_size=6144,
|
163 |
+
num_hidden_layers=24,
|
164 |
+
num_attention_heads=32,
|
165 |
+
num_key_value_heads=4,
|
166 |
+
hidden_act="silu",
|
167 |
+
max_position_embeddings=32768,
|
168 |
+
initializer_range=0.02,
|
169 |
+
rms_norm_eps=1e-6,
|
170 |
+
use_cache=True,
|
171 |
+
tie_word_embeddings=False,
|
172 |
+
rope_theta=10000.0,
|
173 |
+
rope_scaling=None,
|
174 |
+
attention_bias=False,
|
175 |
+
use_sliding_window=False,
|
176 |
+
sliding_window=4096,
|
177 |
+
attention_dropout=0.0,
|
178 |
+
decoder_sparse_step=1,
|
179 |
+
moe_intermediate_size=768,
|
180 |
+
num_experts_per_tok=8,
|
181 |
+
num_experts=128,
|
182 |
+
norm_topk_prob=True,
|
183 |
+
output_router_logits=False,
|
184 |
+
router_aux_loss_coef=0.001,
|
185 |
+
mlp_only_layers=None,
|
186 |
+
routed_scaling_factor=2.5,
|
187 |
+
n_shared_experts=1,
|
188 |
+
**kwargs,
|
189 |
+
):
|
190 |
+
super().__init__(
|
191 |
+
tie_word_embeddings=tie_word_embeddings,
|
192 |
+
**kwargs,
|
193 |
+
)
|
194 |
+
self.vocab_size = vocab_size
|
195 |
+
self.max_position_embeddings = max_position_embeddings
|
196 |
+
self.hidden_size = hidden_size
|
197 |
+
self.intermediate_size = intermediate_size
|
198 |
+
self.num_hidden_layers = num_hidden_layers
|
199 |
+
self.num_attention_heads = num_attention_heads
|
200 |
+
self.use_sliding_window = use_sliding_window
|
201 |
+
self.sliding_window = sliding_window if use_sliding_window else None
|
202 |
+
|
203 |
+
self.num_key_value_heads = num_key_value_heads
|
204 |
+
self.hidden_act = hidden_act
|
205 |
+
self.initializer_range = initializer_range
|
206 |
+
self.rms_norm_eps = rms_norm_eps
|
207 |
+
self.use_cache = use_cache
|
208 |
+
self.rope_theta = rope_theta
|
209 |
+
self.rope_scaling = rope_scaling
|
210 |
+
self.attention_bias = attention_bias
|
211 |
+
self.attention_dropout = attention_dropout
|
212 |
+
# Validate the correctness of rotary position embeddings parameters
|
213 |
+
# BC: if there is a 'type' field, move it to 'rope_type'.
|
214 |
+
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
215 |
+
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
216 |
+
rope_config_validation(self)
|
217 |
+
|
218 |
+
# MoE arguments
|
219 |
+
self.decoder_sparse_step = decoder_sparse_step
|
220 |
+
self.moe_intermediate_size = moe_intermediate_size
|
221 |
+
self.num_experts_per_tok = num_experts_per_tok
|
222 |
+
self.num_experts = num_experts
|
223 |
+
self.norm_topk_prob = norm_topk_prob
|
224 |
+
self.output_router_logits = output_router_logits
|
225 |
+
self.router_aux_loss_coef = router_aux_loss_coef
|
226 |
+
self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
|
227 |
+
|
228 |
+
self.routed_scaling_factor = routed_scaling_factor
|
229 |
+
self.n_shared_experts = n_shared_experts
|
230 |
+
|
231 |
+
|
232 |
+
__all__ = ["KlearConfig"]
|
generation_config.json
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"bos_token_id": 151643,
|
3 |
+
"eos_token_id": 151643,
|
4 |
+
"max_new_tokens": 4096,
|
5 |
+
"transformers_version": "4.56.0"
|
6 |
+
}
|
merges.txt
ADDED
The diff for this file is too large to render.
See raw diff
|
|
model-00001-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:5abcb5985bf8c9458c0c7b09ec3ab1a2cf768ccdfe6f2ed2950464a788943077
|
3 |
+
size 4996631924
|
model-00002-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:34dc79fdf0169eda05906aa526489901a0478fb8a3e02c88d55ac8ce8b5b70ee
|
3 |
+
size 4998242320
|
model-00003-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:43bb174eeaafd3b828d78e3cd613b34b30542b37a3f8b13c5d431c6b6f77b872
|
3 |
+
size 4998242432
|
model-00004-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:616a619532c27e8bfa20ef98640c92dba8563818c03004410b8e416e9da9dd96
|
3 |
+
size 4996651044
|
model-00005-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:730e1e114edfacd7dfa3084a45ef92129db865627f49234909a66416bb1e3fc1
|
3 |
+
size 4998242232
|
model-00006-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a8d90ac87c7b8aecf88c24833d7c958a1bebd5824350163eb402390a0dc92cc5
|
3 |
+
size 4998242568
|
model-00007-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:9008f92d3d7e6bec6b4384f0f032759ed278295f20bdf37d155241efa9bc1caa
|
3 |
+
size 4998243808
|
model-00008-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:e10f16e09709fac19bd6b029b04a1e6b54b16be721e27345a87a3a2a8ccc8f35
|
3 |
+
size 4996652420
|
model-00009-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:99c8ec31a8b49dde25eb9e612fa699147dcdc4786e36b707e0060a076e490a78
|
3 |
+
size 4998243608
|
model-00010-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:8d8e6544ffb56bcc6066f62f1a4249019b98db9558f7d57e2e8beee384b8f933
|
3 |
+
size 4998243672
|
model-00011-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:01973a3b1b920364392af52ee39659c8260d3117f295de225f8f876e800a59d9
|
3 |
+
size 4998243808
|
model-00012-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:e1401c2fdc3677c9c281313d307b9a09e7db1f179e80fbd39b7275fdd34be784
|
3 |
+
size 4996652436
|
model-00013-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:48ae47d7f09144ba130c57facc60cf2c89ce55815defabb85871cf3e09aa479d
|
3 |
+
size 4998243608
|
model-00014-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:6fd105f4cd7647d0617b736975f9131608091bcdfd8e04252f7cecac1f12c629
|
3 |
+
size 4998243656
|
model-00015-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:acd4705ace9a526bfdc76fe94a3bcd16f70e0477f855318a83073ab61c180963
|
3 |
+
size 4998243808
|
model-00016-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:6e362d8add20c3581a0b7a0b42314e2fdd5f51aa15316ade8f8cc1d49241ac0f
|
3 |
+
size 4996652444
|
model-00017-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:1c1ed424862b26ac921af22f9e7aa48f7336153d4f0c503ccfa793f1fbae4703
|
3 |
+
size 4998243608
|
model-00018-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:7d3abdefe2a7e54d5e7f895ecb2577e8430c8b7b612dd7fb9873c7d0e83c47d0
|
3 |
+
size 4998243648
|
model-00019-of-00019.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:b1022f8ad9ed178de6c29c23a9341572872dfae455203f48ddac848811de443b
|
3 |
+
size 2472102012
|
model.safetensors.index.json
ADDED
The diff for this file is too large to render.
See raw diff
|
|
modeling_klear.py
ADDED
@@ -0,0 +1,763 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
2 |
+
# This file was automatically generated from src/transformers/models/klear/modular_klear.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_klear.py file directly. One of our CI enforces this.
|
6 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
7 |
+
|
8 |
+
from typing import Callable, Optional, Union
|
9 |
+
|
10 |
+
import torch
|
11 |
+
import torch.nn.functional as F
|
12 |
+
from torch import nn
|
13 |
+
|
14 |
+
from transformers.activations import ACT2FN
|
15 |
+
from transformers.cache_utils import Cache, DynamicCache
|
16 |
+
from transformers.generation import GenerationMixin
|
17 |
+
from transformers.integrations import use_kernel_forward_from_hub
|
18 |
+
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
|
19 |
+
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
20 |
+
from transformers.modeling_layers import (
|
21 |
+
GenericForQuestionAnswering,
|
22 |
+
GenericForSequenceClassification,
|
23 |
+
GenericForTokenClassification,
|
24 |
+
GradientCheckpointingLayer,
|
25 |
+
)
|
26 |
+
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
27 |
+
from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
|
28 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
29 |
+
from transformers.processing_utils import Unpack
|
30 |
+
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
|
31 |
+
from transformers.utils.generic import OutputRecorder, check_model_inputs
|
32 |
+
from .configuration_klear import KlearConfig
|
33 |
+
|
34 |
+
|
35 |
+
def rotate_half(x):
|
36 |
+
"""Rotates half the hidden dims of the input."""
|
37 |
+
x1 = x[..., : x.shape[-1] // 2]
|
38 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
39 |
+
return torch.cat((-x2, x1), dim=-1)
|
40 |
+
|
41 |
+
|
42 |
+
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
43 |
+
"""Applies Rotary Position Embedding to the query and key tensors.
|
44 |
+
|
45 |
+
Args:
|
46 |
+
q (`torch.Tensor`): The query tensor.
|
47 |
+
k (`torch.Tensor`): The key tensor.
|
48 |
+
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
49 |
+
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
50 |
+
position_ids (`torch.Tensor`, *optional*):
|
51 |
+
Deprecated and unused.
|
52 |
+
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
53 |
+
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
54 |
+
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
55 |
+
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
56 |
+
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
57 |
+
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
58 |
+
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
59 |
+
Returns:
|
60 |
+
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
61 |
+
"""
|
62 |
+
cos = cos.unsqueeze(unsqueeze_dim)
|
63 |
+
sin = sin.unsqueeze(unsqueeze_dim)
|
64 |
+
q_embed = (q * cos) + (rotate_half(q) * sin)
|
65 |
+
k_embed = (k * cos) + (rotate_half(k) * sin)
|
66 |
+
return q_embed, k_embed
|
67 |
+
|
68 |
+
|
69 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
70 |
+
"""
|
71 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
72 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
73 |
+
"""
|
74 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
75 |
+
if n_rep == 1:
|
76 |
+
return hidden_states
|
77 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
78 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
79 |
+
|
80 |
+
|
81 |
+
def eager_attention_forward(
|
82 |
+
module: nn.Module,
|
83 |
+
query: torch.Tensor,
|
84 |
+
key: torch.Tensor,
|
85 |
+
value: torch.Tensor,
|
86 |
+
attention_mask: Optional[torch.Tensor],
|
87 |
+
scaling: float,
|
88 |
+
dropout: float = 0.0,
|
89 |
+
**kwargs: Unpack[TransformersKwargs],
|
90 |
+
):
|
91 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
92 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
93 |
+
|
94 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
95 |
+
if attention_mask is not None:
|
96 |
+
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
97 |
+
attn_weights = attn_weights + causal_mask
|
98 |
+
|
99 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
100 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
101 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
102 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
103 |
+
|
104 |
+
return attn_output, attn_weights
|
105 |
+
|
106 |
+
|
107 |
+
class KlearAttention(nn.Module):
|
108 |
+
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
109 |
+
|
110 |
+
def __init__(self, config: KlearConfig, layer_idx: int):
|
111 |
+
super().__init__()
|
112 |
+
self.config = config
|
113 |
+
self.layer_idx = layer_idx
|
114 |
+
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
115 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
116 |
+
self.scaling = self.head_dim**-0.5
|
117 |
+
self.attention_dropout = config.attention_dropout
|
118 |
+
self.is_causal = True
|
119 |
+
|
120 |
+
self.q_proj = nn.Linear(
|
121 |
+
config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
|
122 |
+
)
|
123 |
+
self.k_proj = nn.Linear(
|
124 |
+
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
125 |
+
)
|
126 |
+
self.v_proj = nn.Linear(
|
127 |
+
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
128 |
+
)
|
129 |
+
self.o_proj = nn.Linear(
|
130 |
+
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
|
131 |
+
)
|
132 |
+
self.q_norm = KlearRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
|
133 |
+
self.k_norm = KlearRMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape
|
134 |
+
self.sliding_window = getattr(config, "sliding_window", None)
|
135 |
+
|
136 |
+
def forward(
|
137 |
+
self,
|
138 |
+
hidden_states: torch.Tensor,
|
139 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
140 |
+
attention_mask: Optional[torch.Tensor],
|
141 |
+
past_key_value: Optional[Cache] = None,
|
142 |
+
cache_position: Optional[torch.LongTensor] = None,
|
143 |
+
**kwargs: Unpack[FlashAttentionKwargs],
|
144 |
+
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
|
145 |
+
input_shape = hidden_states.shape[:-1]
|
146 |
+
hidden_shape = (*input_shape, -1, self.head_dim)
|
147 |
+
|
148 |
+
query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
|
149 |
+
key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
|
150 |
+
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
151 |
+
|
152 |
+
cos, sin = position_embeddings
|
153 |
+
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
154 |
+
|
155 |
+
if past_key_value is not None:
|
156 |
+
# sin and cos are specific to RoPE models; cache_position needed for the static cache
|
157 |
+
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
158 |
+
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
159 |
+
|
160 |
+
attention_interface: Callable = eager_attention_forward
|
161 |
+
if self.config._attn_implementation != "eager":
|
162 |
+
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
163 |
+
|
164 |
+
attn_output, attn_weights = attention_interface(
|
165 |
+
self,
|
166 |
+
query_states,
|
167 |
+
key_states,
|
168 |
+
value_states,
|
169 |
+
attention_mask,
|
170 |
+
dropout=0.0 if not self.training else self.attention_dropout,
|
171 |
+
scaling=self.scaling,
|
172 |
+
sliding_window=self.sliding_window, # diff with Llama
|
173 |
+
**kwargs,
|
174 |
+
)
|
175 |
+
|
176 |
+
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
177 |
+
attn_output = self.o_proj(attn_output)
|
178 |
+
return attn_output, attn_weights
|
179 |
+
|
180 |
+
|
181 |
+
class KlearMLP(nn.Module):
|
182 |
+
def __init__(self, config, intermediate_size=None):
|
183 |
+
super().__init__()
|
184 |
+
self.config = config
|
185 |
+
self.hidden_size = config.hidden_size
|
186 |
+
self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
|
187 |
+
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
188 |
+
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
189 |
+
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
190 |
+
self.act_fn = ACT2FN[config.hidden_act]
|
191 |
+
|
192 |
+
def forward(self, x):
|
193 |
+
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
194 |
+
return down_proj
|
195 |
+
|
196 |
+
|
197 |
+
class KlearSparseMoeBlock(nn.Module):
|
198 |
+
def __init__(self, config):
|
199 |
+
super().__init__()
|
200 |
+
self.config = config
|
201 |
+
self.num_experts = config.num_experts
|
202 |
+
self.top_k = config.num_experts_per_tok
|
203 |
+
self.norm_topk_prob = config.norm_topk_prob
|
204 |
+
|
205 |
+
# router
|
206 |
+
self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
|
207 |
+
self.experts = nn.ModuleList(
|
208 |
+
[KlearMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)]
|
209 |
+
)
|
210 |
+
self.shared_experts = KlearMLP(
|
211 |
+
config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
|
212 |
+
)
|
213 |
+
|
214 |
+
self.coefficient = nn.Linear(config.hidden_size, 2)
|
215 |
+
self.register_buffer("expert_bias", torch.zeros(self.num_experts, dtype=torch.float32))
|
216 |
+
|
217 |
+
def forward(self, hidden_states):
|
218 |
+
residuals = hidden_states
|
219 |
+
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
220 |
+
hidden_states = hidden_states.view(-1, hidden_dim)
|
221 |
+
# router_logits: (batch * sequence_length, n_experts)
|
222 |
+
router_logits = nn.functional.linear(hidden_states.to(torch.float32), self.gate.weight.to(torch.float32))
|
223 |
+
|
224 |
+
routing_weights = F.sigmoid(router_logits)
|
225 |
+
ori_routing_weights = routing_weights
|
226 |
+
|
227 |
+
# using bias
|
228 |
+
biasd_routing_weights = routing_weights + self.expert_bias.unsqueeze(0)
|
229 |
+
_, selected_experts = torch.topk(biasd_routing_weights, self.top_k, dim=-1)
|
230 |
+
|
231 |
+
# Extract corresponding original probabilities
|
232 |
+
ori_routing_weights = torch.gather(ori_routing_weights, dim=-1, index=selected_experts)
|
233 |
+
|
234 |
+
if self.norm_topk_prob: # only diff with mixtral sparse moe block!
|
235 |
+
ori_routing_weights /= ori_routing_weights.sum(dim=-1, keepdim=True)
|
236 |
+
# we cast back to the input dtype
|
237 |
+
ori_routing_weights = ori_routing_weights.to(hidden_states.dtype)
|
238 |
+
|
239 |
+
final_hidden_states = torch.zeros(
|
240 |
+
(batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
|
241 |
+
)
|
242 |
+
|
243 |
+
# One hot encode the selected experts to create an expert mask
|
244 |
+
# this will be used to easily index which expert is going to be sollicitated
|
245 |
+
expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
|
246 |
+
|
247 |
+
# Loop over all available experts in the model and perform the computation on each expert
|
248 |
+
expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
|
249 |
+
for expert_idx in expert_hitted:
|
250 |
+
expert_layer = self.experts[expert_idx]
|
251 |
+
idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
|
252 |
+
|
253 |
+
# Index the correct hidden states and compute the expert hidden state for
|
254 |
+
# the current expert. We need to make sure to multiply the output hidden
|
255 |
+
# states by `ori_routing_weights` on the corresponding tokens (top-1 and top-2)
|
256 |
+
current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
|
257 |
+
current_hidden_states = expert_layer(current_state) * ori_routing_weights[top_x, idx, None]
|
258 |
+
|
259 |
+
# However `index_add_` only support torch tensors for indexing so we'll use
|
260 |
+
# the `top_x` tensor here.
|
261 |
+
final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
|
262 |
+
final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
|
263 |
+
|
264 |
+
coef = self.coefficient(residuals).softmax(dim=-1)
|
265 |
+
final_hidden_states = final_hidden_states * coef[..., :1] + self.shared_experts(residuals) * coef[..., 1:]
|
266 |
+
|
267 |
+
return final_hidden_states, router_logits
|
268 |
+
|
269 |
+
|
270 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
271 |
+
class KlearRMSNorm(nn.Module):
|
272 |
+
def __init__(self, hidden_size, eps=1e-6):
|
273 |
+
"""
|
274 |
+
KlearRMSNorm is equivalent to T5LayerNorm
|
275 |
+
"""
|
276 |
+
super().__init__()
|
277 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
278 |
+
self.variance_epsilon = eps
|
279 |
+
|
280 |
+
def forward(self, hidden_states):
|
281 |
+
input_dtype = hidden_states.dtype
|
282 |
+
hidden_states = hidden_states.to(torch.float32)
|
283 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
284 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
285 |
+
return self.weight * hidden_states.to(input_dtype)
|
286 |
+
|
287 |
+
def extra_repr(self):
|
288 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
289 |
+
|
290 |
+
|
291 |
+
class KlearDecoderLayer(GradientCheckpointingLayer):
|
292 |
+
def __init__(self, config: KlearConfig, layer_idx: int):
|
293 |
+
super().__init__()
|
294 |
+
self.hidden_size = config.hidden_size
|
295 |
+
|
296 |
+
self.self_attn = KlearAttention(config, layer_idx)
|
297 |
+
|
298 |
+
if (layer_idx not in config.mlp_only_layers) and (
|
299 |
+
config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
|
300 |
+
):
|
301 |
+
self.mlp = KlearSparseMoeBlock(config)
|
302 |
+
else:
|
303 |
+
self.mlp = KlearMLP(config, intermediate_size=config.intermediate_size)
|
304 |
+
|
305 |
+
self.input_layernorm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
306 |
+
self.post_attention_layernorm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
307 |
+
|
308 |
+
def forward(
|
309 |
+
self,
|
310 |
+
hidden_states: torch.Tensor,
|
311 |
+
attention_mask: Optional[torch.Tensor] = None,
|
312 |
+
position_ids: Optional[torch.LongTensor] = None,
|
313 |
+
past_key_value: Optional[tuple[torch.Tensor]] = None,
|
314 |
+
output_attentions: Optional[bool] = False,
|
315 |
+
output_router_logits: Optional[bool] = False,
|
316 |
+
use_cache: Optional[bool] = False,
|
317 |
+
cache_position: Optional[torch.LongTensor] = None,
|
318 |
+
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
|
319 |
+
**kwargs: Unpack[FlashAttentionKwargs],
|
320 |
+
) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
|
321 |
+
"""
|
322 |
+
Args:
|
323 |
+
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
|
324 |
+
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
|
325 |
+
`(batch, sequence_length)` where padding elements are indicated by 0.
|
326 |
+
output_attentions (`bool`, *optional*):
|
327 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
328 |
+
returned tensors for more detail.
|
329 |
+
output_router_logits (`bool`, *optional*):
|
330 |
+
Whether or not to return the logits of all the routers. They are useful for computing the router loss,
|
331 |
+
and should not be returned during inference.
|
332 |
+
use_cache (`bool`, *optional*):
|
333 |
+
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
|
334 |
+
(see `past_key_values`).
|
335 |
+
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
|
336 |
+
cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
|
337 |
+
Indices depicting the position of the input sequence tokens in the sequence.
|
338 |
+
position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
|
339 |
+
Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
|
340 |
+
with `head_dim` being the embedding dimension of each attention head.
|
341 |
+
kwargs (`dict`, *optional*):
|
342 |
+
Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
|
343 |
+
into the model
|
344 |
+
"""
|
345 |
+
|
346 |
+
residual = hidden_states
|
347 |
+
hidden_states = self.input_layernorm(hidden_states)
|
348 |
+
|
349 |
+
# Self Attention
|
350 |
+
hidden_states, self_attn_weights = self.self_attn(
|
351 |
+
hidden_states=hidden_states,
|
352 |
+
attention_mask=attention_mask,
|
353 |
+
position_ids=position_ids,
|
354 |
+
past_key_value=past_key_value,
|
355 |
+
output_attentions=output_attentions,
|
356 |
+
use_cache=use_cache,
|
357 |
+
cache_position=cache_position,
|
358 |
+
position_embeddings=position_embeddings,
|
359 |
+
**kwargs,
|
360 |
+
)
|
361 |
+
hidden_states = residual + hidden_states
|
362 |
+
|
363 |
+
# Fully Connected
|
364 |
+
residual = hidden_states
|
365 |
+
hidden_states = self.post_attention_layernorm(hidden_states)
|
366 |
+
|
367 |
+
hidden_states = self.mlp(hidden_states)
|
368 |
+
if isinstance(hidden_states, tuple):
|
369 |
+
hidden_states, router_logits = hidden_states
|
370 |
+
else:
|
371 |
+
router_logits = None
|
372 |
+
|
373 |
+
hidden_states = residual + hidden_states
|
374 |
+
|
375 |
+
outputs = (hidden_states,)
|
376 |
+
|
377 |
+
if output_attentions:
|
378 |
+
outputs += (self_attn_weights,)
|
379 |
+
|
380 |
+
if output_router_logits:
|
381 |
+
outputs += (router_logits,)
|
382 |
+
|
383 |
+
return outputs
|
384 |
+
|
385 |
+
|
386 |
+
class KlearRotaryEmbedding(nn.Module):
|
387 |
+
def __init__(self, config: KlearConfig, device=None):
|
388 |
+
super().__init__()
|
389 |
+
# BC: "rope_type" was originally "type"
|
390 |
+
if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
|
391 |
+
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
|
392 |
+
else:
|
393 |
+
self.rope_type = "default"
|
394 |
+
self.max_seq_len_cached = config.max_position_embeddings
|
395 |
+
self.original_max_seq_len = config.max_position_embeddings
|
396 |
+
|
397 |
+
self.config = config
|
398 |
+
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
399 |
+
|
400 |
+
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
401 |
+
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
402 |
+
self.original_inv_freq = self.inv_freq
|
403 |
+
|
404 |
+
@torch.no_grad()
|
405 |
+
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
406 |
+
def forward(self, x, position_ids):
|
407 |
+
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
|
408 |
+
position_ids_expanded = position_ids[:, None, :].float()
|
409 |
+
|
410 |
+
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
411 |
+
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
412 |
+
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
413 |
+
emb = torch.cat((freqs, freqs), dim=-1)
|
414 |
+
cos = emb.cos() * self.attention_scaling
|
415 |
+
sin = emb.sin() * self.attention_scaling
|
416 |
+
|
417 |
+
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
418 |
+
|
419 |
+
|
420 |
+
@auto_docstring
|
421 |
+
class KlearPreTrainedModel(PreTrainedModel):
|
422 |
+
config: KlearConfig
|
423 |
+
base_model_prefix = "model"
|
424 |
+
supports_gradient_checkpointing = True
|
425 |
+
_no_split_modules = ["KlearDecoderLayer"]
|
426 |
+
_skip_keys_device_placement = ["past_key_values"]
|
427 |
+
_supports_flash_attn = True
|
428 |
+
_supports_sdpa = True
|
429 |
+
_supports_flex_attn = True
|
430 |
+
_can_compile_fullgraph = False # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
|
431 |
+
_supports_attention_backend = True
|
432 |
+
_can_record_outputs = {
|
433 |
+
"router_logits": OutputRecorder(KlearSparseMoeBlock, index=1),
|
434 |
+
"hidden_states": KlearDecoderLayer,
|
435 |
+
"attentions": KlearAttention,
|
436 |
+
}
|
437 |
+
|
438 |
+
|
439 |
+
@auto_docstring
|
440 |
+
class KlearModel(KlearPreTrainedModel):
|
441 |
+
def __init__(self, config: KlearConfig):
|
442 |
+
super().__init__(config)
|
443 |
+
self.padding_idx = config.pad_token_id
|
444 |
+
self.vocab_size = config.vocab_size
|
445 |
+
|
446 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
447 |
+
self.layers = nn.ModuleList(
|
448 |
+
[KlearDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
449 |
+
)
|
450 |
+
self.norm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
451 |
+
self.rotary_emb = KlearRotaryEmbedding(config=config)
|
452 |
+
self.gradient_checkpointing = False
|
453 |
+
|
454 |
+
# Initialize weights and apply final processing
|
455 |
+
self.post_init()
|
456 |
+
|
457 |
+
@check_model_inputs
|
458 |
+
@auto_docstring
|
459 |
+
def forward(
|
460 |
+
self,
|
461 |
+
input_ids: Optional[torch.LongTensor] = None,
|
462 |
+
attention_mask: Optional[torch.Tensor] = None,
|
463 |
+
position_ids: Optional[torch.LongTensor] = None,
|
464 |
+
past_key_values: Optional[list[torch.FloatTensor]] = None,
|
465 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
466 |
+
use_cache: Optional[bool] = None,
|
467 |
+
output_attentions: Optional[bool] = None,
|
468 |
+
output_hidden_states: Optional[bool] = None,
|
469 |
+
output_router_logits: Optional[bool] = None,
|
470 |
+
cache_position: Optional[torch.LongTensor] = None,
|
471 |
+
**kwargs: Unpack[TransformersKwargs],
|
472 |
+
) -> MoeModelOutputWithPast:
|
473 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
474 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
475 |
+
|
476 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
477 |
+
output_router_logits = (
|
478 |
+
output_router_logits if output_router_logits is not None else self.config.output_router_logits
|
479 |
+
)
|
480 |
+
output_hidden_states = (
|
481 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
482 |
+
)
|
483 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
484 |
+
|
485 |
+
if use_cache and past_key_values is None:
|
486 |
+
past_key_values = DynamicCache()
|
487 |
+
|
488 |
+
if inputs_embeds is None:
|
489 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
490 |
+
|
491 |
+
if cache_position is None:
|
492 |
+
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
493 |
+
cache_position = torch.arange(
|
494 |
+
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
495 |
+
)
|
496 |
+
if position_ids is None:
|
497 |
+
position_ids = cache_position.unsqueeze(0)
|
498 |
+
|
499 |
+
mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
|
500 |
+
causal_mask = mask_function(
|
501 |
+
config=self.config,
|
502 |
+
input_embeds=inputs_embeds,
|
503 |
+
attention_mask=attention_mask,
|
504 |
+
cache_position=cache_position,
|
505 |
+
past_key_values=past_key_values,
|
506 |
+
position_ids=position_ids,
|
507 |
+
)
|
508 |
+
|
509 |
+
hidden_states = inputs_embeds
|
510 |
+
|
511 |
+
# create position embeddings to be shared across the decoder layers
|
512 |
+
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
513 |
+
|
514 |
+
# decoder layers
|
515 |
+
all_hidden_states = () if output_hidden_states else None
|
516 |
+
all_self_attns = () if output_attentions else None
|
517 |
+
all_router_logits = () if output_router_logits else None
|
518 |
+
|
519 |
+
for decoder_layer in self.layers:
|
520 |
+
if output_hidden_states:
|
521 |
+
all_hidden_states += (hidden_states,)
|
522 |
+
|
523 |
+
layer_outputs = decoder_layer(
|
524 |
+
hidden_states,
|
525 |
+
attention_mask=causal_mask,
|
526 |
+
position_ids=position_ids,
|
527 |
+
past_key_value=past_key_values,
|
528 |
+
output_attentions=output_attentions,
|
529 |
+
output_router_logits=output_router_logits,
|
530 |
+
use_cache=use_cache,
|
531 |
+
cache_position=cache_position,
|
532 |
+
position_embeddings=position_embeddings,
|
533 |
+
**kwargs,
|
534 |
+
)
|
535 |
+
|
536 |
+
hidden_states = layer_outputs[0]
|
537 |
+
|
538 |
+
if output_attentions:
|
539 |
+
all_self_attns += (layer_outputs[1],)
|
540 |
+
|
541 |
+
if output_router_logits:
|
542 |
+
all_router_logits += (layer_outputs[-1],)
|
543 |
+
|
544 |
+
hidden_states = self.norm(hidden_states)
|
545 |
+
|
546 |
+
return MoeModelOutputWithPast(
|
547 |
+
last_hidden_state=hidden_states,
|
548 |
+
past_key_values=past_key_values,
|
549 |
+
hidden_states=all_hidden_states,
|
550 |
+
attentions=all_self_attns,
|
551 |
+
router_logits=all_router_logits,
|
552 |
+
)
|
553 |
+
|
554 |
+
|
555 |
+
def load_balancing_loss_func(
|
556 |
+
gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
|
557 |
+
num_experts: Optional[int] = None,
|
558 |
+
top_k: int = 2,
|
559 |
+
attention_mask: Optional[torch.Tensor] = None,
|
560 |
+
moe_aux_loss_coeff: float = 1,
|
561 |
+
) -> torch.Tensor:
|
562 |
+
"""
|
563 |
+
Computes sequence-level auxiliary load balancing loss for MoE gating.
|
564 |
+
|
565 |
+
Args:
|
566 |
+
gate_logits: Tensor of shape [batch_size, seq_len, num_experts]
|
567 |
+
or a tuple of such tensors (for multiple towers).
|
568 |
+
num_experts: Number of experts (inferred from gate_logits if None).
|
569 |
+
top_k: Number of top experts chosen per token.
|
570 |
+
attention_mask: Optional mask [batch_size, seq_len], 1 for valid tokens, 0 for padding.
|
571 |
+
moe_aux_loss_coeff: Scaling coefficient for the balancing loss.
|
572 |
+
|
573 |
+
Returns:
|
574 |
+
A scalar tensor representing the load balancing loss.
|
575 |
+
"""
|
576 |
+
# Merge towers if provided
|
577 |
+
if isinstance(gate_logits, tuple):
|
578 |
+
gate_logits = torch.cat(gate_logits, dim=0)
|
579 |
+
|
580 |
+
assert gate_logits is not None, "gate_logits must be provided"
|
581 |
+
batch_size, seq_len, n_experts = gate_logits.shape
|
582 |
+
num_experts = n_experts if num_experts is None else num_experts
|
583 |
+
assert num_experts == n_experts, f"num_experts ({num_experts}) != gate dimension ({n_experts})"
|
584 |
+
|
585 |
+
# Compute gating probabilities
|
586 |
+
gate_probs = F.softmax(gate_logits, dim=-1)
|
587 |
+
|
588 |
+
# Optionally mask padding tokens
|
589 |
+
if attention_mask is not None:
|
590 |
+
mask = attention_mask.float().unsqueeze(-1) # [batch, seq, 1]
|
591 |
+
else:
|
592 |
+
mask = torch.ones(batch_size, seq_len, 1, device=gate_logits.device)
|
593 |
+
|
594 |
+
# Select top_k experts per token
|
595 |
+
topk_vals, topk_idx = torch.topk(gate_probs, top_k, dim=-1) # both [batch, seq, top_k]
|
596 |
+
# Build one-hot mask of assignments
|
597 |
+
one_hot = F.one_hot(topk_idx, num_experts).float() # [batch, seq, top_k, num_experts]
|
598 |
+
# Sum along top_k to combine multiple choices
|
599 |
+
expert_mask = one_hot.sum(dim=2) # [batch, seq, num_experts]
|
600 |
+
|
601 |
+
# Apply token mask
|
602 |
+
expert_mask = expert_mask * mask # zeros out padding
|
603 |
+
gate_probs_masked = gate_probs * mask
|
604 |
+
|
605 |
+
# Normalizer: number of valid tokens per sample
|
606 |
+
tokens_per_sample = mask.sum(dim=1).clamp(min=1.0) # [batch, 1]
|
607 |
+
|
608 |
+
# Sequence-level tokens per expert: fraction of tokens routed to each expert per sample
|
609 |
+
tokens_per_expert = expert_mask.sum(dim=1).div_(tokens_per_sample * top_k / num_experts) # [batch, num_experts]
|
610 |
+
|
611 |
+
# Sequence-level average probability per expert per sample
|
612 |
+
router_prob_per_expert = gate_probs_masked.sum(dim=1).div(tokens_per_sample) # [batch, num_experts]
|
613 |
+
|
614 |
+
# Compute loss per sample: encourage uniform load
|
615 |
+
# Loss = sum_e (tokens_e * probs_e)
|
616 |
+
loss_per_sample = (tokens_per_expert * router_prob_per_expert).sum(dim=1) # [batch]
|
617 |
+
# Average across batch and scale
|
618 |
+
loss = moe_aux_loss_coeff * loss_per_sample.mean()
|
619 |
+
return loss
|
620 |
+
|
621 |
+
|
622 |
+
@auto_docstring
|
623 |
+
class KlearMoeForCausalLM(KlearPreTrainedModel, GenerationMixin):
|
624 |
+
_tied_weights_keys = ["lm_head.weight"]
|
625 |
+
_tp_plan = {"lm_head": "colwise_rep"}
|
626 |
+
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
627 |
+
|
628 |
+
def __init__(self, config):
|
629 |
+
super().__init__(config)
|
630 |
+
self.model = KlearModel(config)
|
631 |
+
self.vocab_size = config.vocab_size
|
632 |
+
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
633 |
+
self.router_aux_loss_coef = config.router_aux_loss_coef
|
634 |
+
self.num_experts = config.num_experts
|
635 |
+
self.num_experts_per_tok = config.num_experts_per_tok
|
636 |
+
self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 1.0)
|
637 |
+
|
638 |
+
# Initialize weights and apply final processing
|
639 |
+
self.post_init()
|
640 |
+
|
641 |
+
def set_decoder(self, decoder):
|
642 |
+
self.model = decoder
|
643 |
+
|
644 |
+
def get_decoder(self):
|
645 |
+
return self.model
|
646 |
+
|
647 |
+
@can_return_tuple
|
648 |
+
@auto_docstring
|
649 |
+
def forward(
|
650 |
+
self,
|
651 |
+
input_ids: Optional[torch.LongTensor] = None,
|
652 |
+
attention_mask: Optional[torch.Tensor] = None,
|
653 |
+
position_ids: Optional[torch.LongTensor] = None,
|
654 |
+
past_key_values: Optional[list[torch.FloatTensor]] = None,
|
655 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
656 |
+
labels: Optional[torch.LongTensor] = None,
|
657 |
+
use_cache: Optional[bool] = None,
|
658 |
+
output_attentions: Optional[bool] = None,
|
659 |
+
output_hidden_states: Optional[bool] = None,
|
660 |
+
output_router_logits: Optional[bool] = None,
|
661 |
+
cache_position: Optional[torch.LongTensor] = None,
|
662 |
+
logits_to_keep: Union[int, torch.Tensor] = 0,
|
663 |
+
**kwargs: Unpack[TransformersKwargs],
|
664 |
+
) -> MoeCausalLMOutputWithPast:
|
665 |
+
r"""
|
666 |
+
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
667 |
+
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
668 |
+
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
669 |
+
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
670 |
+
|
671 |
+
Example:
|
672 |
+
|
673 |
+
```python
|
674 |
+
>>> from transformers import AutoTokenizer, KlearMoeForCausalLM
|
675 |
+
|
676 |
+
>>> model = KlearMoeForCausalLM.from_pretrained("Klear-kwaii/Klear-MoE")
|
677 |
+
>>> tokenizer = AutoTokenizer.from_pretrained("Klear-kwaii/Klear-MoE")
|
678 |
+
|
679 |
+
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
680 |
+
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
681 |
+
|
682 |
+
>>> # Generate
|
683 |
+
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
684 |
+
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
685 |
+
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
686 |
+
```"""
|
687 |
+
|
688 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
689 |
+
output_router_logits = (
|
690 |
+
output_router_logits if output_router_logits is not None else self.config.output_router_logits
|
691 |
+
)
|
692 |
+
|
693 |
+
output_hidden_states = (
|
694 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
695 |
+
)
|
696 |
+
|
697 |
+
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
698 |
+
outputs: MoeModelOutputWithPast = self.model(
|
699 |
+
input_ids=input_ids,
|
700 |
+
attention_mask=attention_mask,
|
701 |
+
position_ids=position_ids,
|
702 |
+
past_key_values=past_key_values,
|
703 |
+
inputs_embeds=inputs_embeds,
|
704 |
+
use_cache=use_cache,
|
705 |
+
output_attentions=output_attentions,
|
706 |
+
output_hidden_states=output_hidden_states,
|
707 |
+
output_router_logits=output_router_logits,
|
708 |
+
cache_position=cache_position,
|
709 |
+
**kwargs,
|
710 |
+
)
|
711 |
+
|
712 |
+
hidden_states = outputs.last_hidden_state
|
713 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
714 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
715 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
716 |
+
|
717 |
+
loss = None
|
718 |
+
if labels is not None:
|
719 |
+
loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
|
720 |
+
|
721 |
+
aux_loss = None
|
722 |
+
if output_router_logits:
|
723 |
+
aux_loss = load_balancing_loss_func(
|
724 |
+
outputs.router_logits,
|
725 |
+
self.num_experts,
|
726 |
+
self.num_experts_per_tok,
|
727 |
+
attention_mask,
|
728 |
+
self.moe_aux_loss_coeff,
|
729 |
+
)
|
730 |
+
if labels is not None:
|
731 |
+
loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
|
732 |
+
|
733 |
+
return MoeCausalLMOutputWithPast(
|
734 |
+
loss=loss,
|
735 |
+
aux_loss=aux_loss,
|
736 |
+
logits=logits,
|
737 |
+
past_key_values=outputs.past_key_values,
|
738 |
+
hidden_states=outputs.hidden_states,
|
739 |
+
attentions=outputs.attentions,
|
740 |
+
router_logits=outputs.router_logits,
|
741 |
+
)
|
742 |
+
|
743 |
+
|
744 |
+
class KlearForSequenceClassification(GenericForSequenceClassification, KlearPreTrainedModel):
|
745 |
+
pass
|
746 |
+
|
747 |
+
|
748 |
+
class KlearForTokenClassification(GenericForTokenClassification, KlearPreTrainedModel):
|
749 |
+
pass
|
750 |
+
|
751 |
+
|
752 |
+
class KlearForQuestionAnswering(GenericForQuestionAnswering, KlearPreTrainedModel):
|
753 |
+
base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model`
|
754 |
+
|
755 |
+
|
756 |
+
__all__ = [
|
757 |
+
"KlearMoeForCausalLM",
|
758 |
+
"KlearForQuestionAnswering",
|
759 |
+
"KlearModel",
|
760 |
+
"KlearPreTrainedModel",
|
761 |
+
"KlearForSequenceClassification",
|
762 |
+
"KlearForTokenClassification",
|
763 |
+
]
|
tokenizer.json
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
|
3 |
+
size 11422654
|
tokenizer_config.json
ADDED
@@ -0,0 +1,239 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set content = message.content %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in message.content %}\n {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if loop.index0 > ns.last_query_index %}\n {%- if loop.last or (not loop.last and reasoning_content) %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- message.content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if enable_thinking is defined and enable_thinking is false %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- endif %}\n{%- endif %}",
|
231 |
+
"clean_up_tokenization_spaces": false,
|
232 |
+
"eos_token": "<|endoftext|>",
|
233 |
+
"errors": "replace",
|
234 |
+
"model_max_length": 131072,
|
235 |
+
"pad_token": "<|endoftext|>",
|
236 |
+
"split_special_tokens": false,
|
237 |
+
"tokenizer_class": "Qwen2Tokenizer",
|
238 |
+
"unk_token": null
|
239 |
+
}
|
vocab.json
ADDED
The diff for this file is too large to render.
See raw diff
|
|