BubbleQ commited on
Commit
0137742
·
verified ·
1 Parent(s): 9976de6

Upload folder using huggingface_hub

Browse files
.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
- license: apache-2.0
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