prince-canuma commited on
Commit
41149a6
·
verified ·
1 Parent(s): 1fd7ae4

Add files using upload-large-folder tool

Browse files
README.md ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: mlx
3
+ inference: false
4
+ extra_gated_description: To learn more about how we process your personal data, please
5
+ read our <a href="https://poolside.ai/legal/privacy">Privacy Policy</a>.
6
+ tags:
7
+ - laguna-xs.2
8
+ - mlx
9
+ license: apache-2.0
10
+ pipeline_tag: text-generation
11
+ base_model: poolside/Laguna-XS.2
12
+ ---
13
+
14
+ # mlx-community/Laguna-XS.2-bf16
15
+
16
+ This model [mlx-community/Laguna-XS.2-bf16](https://huggingface.co/mlx-community/Laguna-XS.2-bf16) was
17
+ converted to MLX format from [poolside/Laguna-XS.2](https://huggingface.co/poolside/Laguna-XS.2)
18
+ using mlx-lm version **0.31.3**.
19
+
20
+ ## Use with mlx
21
+
22
+ ```bash
23
+ pip install mlx-lm
24
+ ```
25
+
26
+ ```python
27
+ from mlx_lm import load, generate
28
+
29
+ model, tokenizer = load("mlx-community/Laguna-XS.2-bf16")
30
+
31
+ prompt = "hello"
32
+
33
+ if tokenizer.chat_template is not None:
34
+ messages = [{"role": "user", "content": prompt}]
35
+ prompt = tokenizer.apply_chat_template(
36
+ messages, add_generation_prompt=True, return_dict=False,
37
+ )
38
+
39
+ response = generate(model, tokenizer, prompt=prompt, verbose=True)
40
+ ```
chat_template.jinja ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {#- Iteration on laguna_glm_thinking_v5/chat_template.jinja -#}
2
+ {#- Adds a default system message (used when no system message is provided in `messages`). -#}
3
+ {{- "〈|EOS|〉" -}}
4
+ {%- set enable_thinking = enable_thinking | default(false) -%}
5
+ {%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
6
+ {%- set add_generation_prompt = add_generation_prompt | default(false) -%}
7
+
8
+ {#- ───── header (system message) ───── -#}
9
+ {%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
10
+ {%- if messages and messages[0].role == "system" -%}
11
+ {%- set system_message = messages[0].content -%}
12
+ {%- endif -%}
13
+
14
+ {%- if (system_message and system_message.strip()) or tools -%}
15
+ {{- "<system>\n" -}}
16
+
17
+ {%- if system_message and system_message.strip() -%}
18
+ {{- "\n" -}}
19
+ {{- system_message.rstrip() -}}
20
+ {%- endif -%}
21
+
22
+ {%- if tools -%}
23
+ {{- "\n\n### Tools\n\n" -}}
24
+ {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
25
+ ~ "All available function signatures are listed below:\n"
26
+ ~ "<available_tools>\n") -%}
27
+ {%- for tool in tools -%}
28
+ {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
29
+ {%- endfor -%}
30
+ {%- if enable_thinking -%}
31
+ {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
32
+ "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
33
+ "<think> your thoughts here </think>\n" ~
34
+ "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
35
+ "</tool_call>" -%}
36
+ {%- else -%}
37
+ {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
38
+ "For each function call, return an unescaped XML-like object " ~
39
+ "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
40
+ "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
41
+ "</tool_call>" -%}
42
+ {%- endif -%}
43
+ {{- tool_string -}}
44
+ {%- endif -%}
45
+
46
+ {{- "\n</system>\n" -}}
47
+ {%- endif -%}
48
+
49
+ {#- ───── main loop ───── -#}
50
+ {%- for message in messages -%}
51
+ {%- set content = message.content if message.content is string else "" -%}
52
+ {%- if message.role == "user" -%}
53
+ {{- "<user>\n" + content + "\n</user>\n" -}}
54
+ {%- elif message.role == "assistant" -%}
55
+ {%- generation -%}
56
+ {{- "<assistant>\n" -}}
57
+ {%- if render_assistant_messages_raw -%}
58
+ {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
59
+ {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
60
+ {#- Only prepend if content doesn't already start with it. -#}
61
+ {%- if enable_thinking -%}
62
+ {%- if not content.startswith('<think>') -%}
63
+ {{- '<think>' -}}
64
+ {%- endif -%}
65
+ {%- else -%}
66
+ {%- if not content.startswith('</think>') -%}
67
+ {{- '</think>' -}}
68
+ {%- endif -%}
69
+ {%- endif -%}
70
+ {{- content -}}
71
+ {#- Append closing tag if content doesn't already end with it. -#}
72
+ {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
73
+ {{- '\n</assistant>' -}}
74
+ {%- endif -%}
75
+ {{- "\n" -}}
76
+ {%- else -%}
77
+ {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
78
+ {%- set reasoning_content = '' %}
79
+ {%- if message.reasoning is string %}
80
+ {%- set reasoning_content = message.reasoning %}
81
+ {%- elif message.reasoning_content is string %}
82
+ {%- set reasoning_content = message.reasoning_content %}
83
+ {%- endif %}
84
+ {#- Always strip <think> tags from content if present to avoid duplication -#}
85
+ {%- if '</think>' in content %}
86
+ {%- if not reasoning_content %}
87
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
88
+ {%- endif %}
89
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
90
+ {%- endif %}
91
+ {#- Display reasoning content for all messages -#}
92
+ {%- if reasoning_content -%}
93
+ {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
94
+ {%- else -%}
95
+ {{- '</think>\n' -}}
96
+ {%- endif -%}
97
+ {#- Display main content -#}
98
+ {%- if content.strip() -%}
99
+ {{- content.strip() ~ "\n" -}}
100
+ {%- endif -%}
101
+ {%- if message.tool_calls -%}
102
+ {%- for tool_call in message.tool_calls -%}
103
+ {%- set function_data = tool_call.function -%}
104
+ {{- '<tool_call>' + function_data.name }}
105
+ {% set _args = function_data.arguments %}
106
+ {%- for k, v in _args.items() -%}
107
+ {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
108
+ {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
109
+ {%- endfor -%}
110
+ {{- "</tool_call>\n" -}}
111
+ {%- endfor -%}
112
+ {%- endif -%}
113
+ {{- "</assistant>\n" -}}
114
+ {%- endif -%}
115
+ {%- endgeneration -%}
116
+ {%- elif message.role == "tool" -%}
117
+ {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
118
+ {%- elif message.role == "system" and loop.index0 != 0 -%}
119
+ {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
120
+ {{- "<system>\n" + content + "\n</system>\n" -}}
121
+ {%- endif -%}
122
+ {%- endfor -%}
123
+ {#- ───── generation prompt ───── -#}
124
+ {%- if add_generation_prompt -%}
125
+ {{- "<assistant>\n" -}}
126
+ {#- ───── Include reasoning mode directive ───── -#}
127
+ {%- if not enable_thinking %}
128
+ {{- '</think>' -}}
129
+ {%- else %}
130
+ {{- '<think>' -}}
131
+ {%- endif %}
132
+ {%- endif -%}
config.json ADDED
@@ -0,0 +1,184 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LagunaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_laguna.LagunaConfig",
9
+ "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
10
+ },
11
+ "bos_token_id": 2,
12
+ "eos_token_id": [
13
+ 2,
14
+ 24
15
+ ],
16
+ "gating": true,
17
+ "head_dim": 128,
18
+ "hidden_size": 2048,
19
+ "intermediate_size": 8192,
20
+ "layer_types": [
21
+ "full_attention",
22
+ "sliding_attention",
23
+ "sliding_attention",
24
+ "sliding_attention",
25
+ "full_attention",
26
+ "sliding_attention",
27
+ "sliding_attention",
28
+ "sliding_attention",
29
+ "full_attention",
30
+ "sliding_attention",
31
+ "sliding_attention",
32
+ "sliding_attention",
33
+ "full_attention",
34
+ "sliding_attention",
35
+ "sliding_attention",
36
+ "sliding_attention",
37
+ "full_attention",
38
+ "sliding_attention",
39
+ "sliding_attention",
40
+ "sliding_attention",
41
+ "full_attention",
42
+ "sliding_attention",
43
+ "sliding_attention",
44
+ "sliding_attention",
45
+ "full_attention",
46
+ "sliding_attention",
47
+ "sliding_attention",
48
+ "sliding_attention",
49
+ "full_attention",
50
+ "sliding_attention",
51
+ "sliding_attention",
52
+ "sliding_attention",
53
+ "full_attention",
54
+ "sliding_attention",
55
+ "sliding_attention",
56
+ "sliding_attention",
57
+ "full_attention",
58
+ "sliding_attention",
59
+ "sliding_attention",
60
+ "sliding_attention"
61
+ ],
62
+ "max_position_embeddings": 131072,
63
+ "mlp_layer_types": [
64
+ "dense",
65
+ "sparse",
66
+ "sparse",
67
+ "sparse",
68
+ "sparse",
69
+ "sparse",
70
+ "sparse",
71
+ "sparse",
72
+ "sparse",
73
+ "sparse",
74
+ "sparse",
75
+ "sparse",
76
+ "sparse",
77
+ "sparse",
78
+ "sparse",
79
+ "sparse",
80
+ "sparse",
81
+ "sparse",
82
+ "sparse",
83
+ "sparse",
84
+ "sparse",
85
+ "sparse",
86
+ "sparse",
87
+ "sparse",
88
+ "sparse",
89
+ "sparse",
90
+ "sparse",
91
+ "sparse",
92
+ "sparse",
93
+ "sparse",
94
+ "sparse",
95
+ "sparse",
96
+ "sparse",
97
+ "sparse",
98
+ "sparse",
99
+ "sparse",
100
+ "sparse",
101
+ "sparse",
102
+ "sparse",
103
+ "sparse"
104
+ ],
105
+ "model_type": "laguna",
106
+ "moe_apply_router_weight_on_input": false,
107
+ "moe_intermediate_size": 512,
108
+ "moe_routed_scaling_factor": 2.5,
109
+ "num_attention_heads": 48,
110
+ "num_attention_heads_per_layer": [
111
+ 48,
112
+ 64,
113
+ 64,
114
+ 64,
115
+ 48,
116
+ 64,
117
+ 64,
118
+ 64,
119
+ 48,
120
+ 64,
121
+ 64,
122
+ 64,
123
+ 48,
124
+ 64,
125
+ 64,
126
+ 64,
127
+ 48,
128
+ 64,
129
+ 64,
130
+ 64,
131
+ 48,
132
+ 64,
133
+ 64,
134
+ 64,
135
+ 48,
136
+ 64,
137
+ 64,
138
+ 64,
139
+ 48,
140
+ 64,
141
+ 64,
142
+ 64,
143
+ 48,
144
+ 64,
145
+ 64,
146
+ 64,
147
+ 48,
148
+ 64,
149
+ 64,
150
+ 64
151
+ ],
152
+ "num_experts": 256,
153
+ "num_experts_per_tok": 8,
154
+ "num_hidden_layers": 40,
155
+ "num_key_value_heads": 8,
156
+ "pad_token_id": 9,
157
+ "partial_rotary_factor": 0.5,
158
+ "rms_norm_eps": 1e-06,
159
+ "rope_parameters": {
160
+ "full_attention": {
161
+ "rope_theta": 500000.0,
162
+ "rope_type": "yarn",
163
+ "factor": 32.0,
164
+ "original_max_position_embeddings": 4096,
165
+ "beta_slow": 1.0,
166
+ "beta_fast": 64.0,
167
+ "attention_factor": 1.0,
168
+ "partial_rotary_factor": 0.5
169
+ },
170
+ "sliding_attention": {
171
+ "rope_type": "default",
172
+ "rope_theta": 10000.0,
173
+ "partial_rotary_factor": 1.0
174
+ },
175
+ "original_max_position_embeddings": 4096
176
+ },
177
+ "router_aux_loss_coef": 0.0,
178
+ "shared_expert_intermediate_size": 512,
179
+ "sliding_window": 512,
180
+ "tie_word_embeddings": false,
181
+ "torch_dtype": "bfloat16",
182
+ "use_cache": true,
183
+ "vocab_size": 100352
184
+ }
configuration_laguna.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from typing import Any, Literal
15
+
16
+ from huggingface_hub.dataclasses import strict
17
+
18
+ from transformers.configuration_utils import PreTrainedConfig
19
+ from transformers.modeling_rope_utils import RopeParameters
20
+ from transformers.utils import auto_docstring
21
+
22
+
23
+ @auto_docstring(checkpoint="poolside/laguna-XS.2")
24
+ @strict
25
+ class LagunaConfig(PreTrainedConfig):
26
+ r"""
27
+ partial_rotary_factor (`float`, *optional*):
28
+ Fraction of ``head_dim`` to rotate. Folded into each ``rope_parameters[layer_type]``
29
+ entry by ``__post_init__``.
30
+ num_attention_heads_per_layer (`list[int]`, *optional*):
31
+ Per-layer override for ``num_attention_heads``. Length must equal ``num_hidden_layers``.
32
+ mlp_layer_types (`list[str]`, *optional*):
33
+ Per-layer MLP type — ``"dense"`` or ``"sparse"``. Length must equal
34
+ ``num_hidden_layers``. Defaults to first layer dense, rest sparse.
35
+ moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0):
36
+ Scalar applied to routed-expert output before combining with the shared-expert output.
37
+ moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`):
38
+ Whether to apply router weights to the MoE input rather than the output. Not supported
39
+ in transformers yet; ``True`` will raise a ``NotImplementedError`` for now.
40
+ moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0):
41
+ Scaling factor when applying tanh softcapping on the logits of the MoE router logits.
42
+
43
+ Example:
44
+
45
+ ```python
46
+ >>> from transformers import LagunaModel, LagunaConfig
47
+
48
+ >>> configuration = LagunaConfig()
49
+ >>> model = LagunaModel(configuration)
50
+ >>> configuration = model.config
51
+ ```
52
+ """
53
+
54
+ model_type = "laguna"
55
+ keys_to_ignore_at_inference = ["past_key_values"]
56
+ base_model_tp_plan = {
57
+ "layers.*.self_attn.q_proj": "colwise",
58
+ "layers.*.self_attn.k_proj": "colwise",
59
+ "layers.*.self_attn.v_proj": "colwise",
60
+ "layers.*.self_attn.g_proj": "colwise",
61
+ "layers.*.self_attn.o_proj": "rowwise",
62
+ "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
63
+ "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
64
+ "layers.*.mlp.gate_proj": "colwise",
65
+ "layers.*.mlp.up_proj": "colwise",
66
+ "layers.*.mlp.down_proj": "rowwise",
67
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
68
+ "layers.*.mlp.experts.down_proj": "rowwise",
69
+ "layers.*.mlp.experts": "moe_tp_experts",
70
+ "layers.*.mlp.shared_experts.gate_proj": "colwise",
71
+ "layers.*.mlp.shared_experts.up_proj": "colwise",
72
+ "layers.*.mlp.shared_experts.down_proj": "rowwise",
73
+ }
74
+ base_model_pp_plan = {
75
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
76
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
77
+ "norm": (["hidden_states"], ["hidden_states"]),
78
+ }
79
+
80
+ # Qwen2Moe-inherited defaults we want to override for Laguna's typical shape.
81
+ vocab_size: int = 100352
82
+ hidden_size: int = 2048
83
+ intermediate_size: int = 8192
84
+ num_hidden_layers: int = 40
85
+ num_attention_heads: int = 48
86
+ num_key_value_heads: int = 8
87
+ hidden_act: str = "silu"
88
+ max_position_embeddings: int = 131072
89
+ initializer_range: float = 0.02
90
+ rms_norm_eps: float = 1e-6
91
+ use_cache: bool = True
92
+ tie_word_embeddings: bool = False
93
+ rope_parameters: RopeParameters | dict | None = None
94
+ sliding_window: int | None = None
95
+ attention_dropout: float | int = 0.0
96
+ moe_intermediate_size: int = 512
97
+ shared_expert_intermediate_size: int = 512
98
+ num_experts_per_tok: int = 8
99
+ num_experts: int = 256
100
+ output_router_logits: bool = False
101
+ router_aux_loss_coef: float = 0.001
102
+ layer_types: list[str] | None = None
103
+ pad_token_id: int | None = None
104
+ bos_token_id: int | None = None
105
+ eos_token_id: int | list[int] | None = None
106
+
107
+ # Laguna-specific attention
108
+ head_dim: int = 128
109
+ attention_bias: bool = False
110
+ partial_rotary_factor: float | None = None
111
+ num_attention_heads_per_layer: list[int] | None = None
112
+ # Laguna-specific MoE
113
+ mlp_layer_types: list[str] | None = None
114
+ moe_routed_scaling_factor: float = 1.0
115
+ moe_apply_router_weight_on_input: bool = False
116
+ moe_router_logit_softcapping: float = 0.0
117
+
118
+ def __post_init__(self, **kwargs):
119
+ if self.layer_types is None:
120
+ self.layer_types = ["full_attention"] * self.num_hidden_layers
121
+ if self.mlp_layer_types is None:
122
+ self.mlp_layer_types = ["dense"] + ["sparse"] * (self.num_hidden_layers - 1)
123
+ if self.num_attention_heads_per_layer is None:
124
+ self.num_attention_heads_per_layer = [self.num_attention_heads] * self.num_hidden_layers
125
+
126
+ default_rope_params: dict[Literal["full_attention", "sliding_attention"], dict[str, Any]] = {
127
+ "full_attention": {"rope_type": "default", "rope_theta": 500000.0},
128
+ "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},
129
+ }
130
+ if self.rope_parameters is None:
131
+ self.rope_parameters = default_rope_params
132
+
133
+ self._normalize_rope_parameters()
134
+ # Skip ``Qwen2MoeConfig.__post_init__`` — it references ``mlp_only_layers`` /
135
+ # ``use_sliding_window`` / ``max_window_layers`` which Laguna drops above.
136
+ super().__post_init__(**kwargs)
137
+
138
+ def _normalize_rope_parameters(self):
139
+ """Coerce ``rope_parameters`` to the nested ``{layer_type: {...}}`` shape.
140
+
141
+ Accepts an already-nested dict as-is, or a flat dict that gets broadcast to every
142
+ layer type. A top-level ``partial_rotary_factor`` is folded into each sub-dict as
143
+ a default.
144
+ """
145
+ layer_types = set(self.layer_types)
146
+ rope_params = self.rope_parameters or {}
147
+ is_nested = isinstance(rope_params, dict) and any(k in layer_types for k in rope_params)
148
+ if is_nested:
149
+ nested = {lt: dict(rope_params.get(lt, {})) for lt in layer_types}
150
+ else:
151
+ nested = {lt: dict(rope_params) for lt in layer_types}
152
+
153
+ if self.partial_rotary_factor is not None:
154
+ for params in nested.values():
155
+ params.setdefault("partial_rotary_factor", self.partial_rotary_factor)
156
+
157
+ for params in nested.values():
158
+ params.setdefault("rope_type", "default")
159
+
160
+ self.rope_parameters = nested
161
+ # Null the top-level field now that its value lives in each sub-dict — otherwise
162
+ # ``standardize_rope_params`` would overwrite per-type values with the global one.
163
+ self.partial_rotary_factor = None
164
+
165
+ def convert_rope_params_to_dict(self, **kwargs):
166
+ # No need to handle BC for new models, because they have no old-format `rope_scaling`
167
+ return kwargs
168
+
169
+ def _validate_yarn_rope_parameters(self, rope_parameters: dict, ignore_keys=None):
170
+ """Override: parent reads ``self.rope_parameters["original_max_position_embeddings"]``
171
+ for its post-hoc factor sanity-check, which works for flat rope configs but raises
172
+ ``KeyError`` when ``self.rope_parameters`` is the Laguna/Gemma3-style per-layer-type
173
+ map (its keys are layer types like ``"full_attention"``). Fix locally by reading
174
+ from the per-call ``rope_parameters`` dict that ``validate_rope`` already passes in.
175
+ """
176
+ # Delegate to parent for the shared checks by temporarily swapping in a flat
177
+ # ``self.rope_parameters`` that has the key the parent expects. Cheapest way to
178
+ # share the parent's logic without reimplementing it here.
179
+ flat = getattr(self, "rope_parameters", None)
180
+ self.rope_parameters = rope_parameters
181
+ try:
182
+ super()._validate_yarn_rope_parameters(rope_parameters, ignore_keys=ignore_keys)
183
+ finally:
184
+ self.rope_parameters = flat
185
+
186
+ def validate_architecture(self):
187
+ """Part of ``@strict``-powered validation."""
188
+ if self.moe_apply_router_weight_on_input:
189
+ raise NotImplementedError(
190
+ "moe_apply_router_weight_on_input=True is not yet supported in the "
191
+ "transformers implementation of Laguna."
192
+ )
193
+ if (
194
+ self.num_attention_heads_per_layer is not None
195
+ and len(self.num_attention_heads_per_layer) != self.num_hidden_layers
196
+ ):
197
+ raise ValueError(
198
+ f"num_attention_heads_per_layer length ({len(self.num_attention_heads_per_layer)}) "
199
+ f"must equal num_hidden_layers ({self.num_hidden_layers})."
200
+ )
201
+ if len(self.layer_types) != self.num_hidden_layers:
202
+ raise ValueError(
203
+ f"layer_types length ({len(self.layer_types)}) "
204
+ f"must equal num_hidden_layers ({self.num_hidden_layers})."
205
+ )
206
+ if len(self.mlp_layer_types) != self.num_hidden_layers:
207
+ raise ValueError(
208
+ f"mlp_layer_types length ({len(self.mlp_layer_types)}) "
209
+ f"must equal num_hidden_layers ({self.num_hidden_layers})."
210
+ )
211
+
212
+
213
+ __all__ = ["LagunaConfig"]
generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 2,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 2,
6
+ 24
7
+ ],
8
+ "max_new_tokens": 2048,
9
+ "pad_token_id": 9,
10
+ "temperature": 0.7,
11
+ "top_p": 0.9
12
+ }
model-00001-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f9fe4500c6fa1674cc0cb4af6ee39c771fe3c4f5a50a04cd4bd484be57e70ac7
3
+ size 5108631760
model-00002-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bec7bc2b1f85933f19e24fc36b63ffd08d9f8491c9223938ce54155bcbef2807
3
+ size 5064328168
model-00003-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1c7c2a3a2bff6128e3142c784b734137753fcdb229ba5bf5fea1cef616f56609
3
+ size 5064328152
model-00004-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7ac36bbdb16cc0f3aacc1f18d31ab2b357be892d309b24d26799c7b5eb92d158
3
+ size 5064328201
model-00005-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0808b34a8cc8335826c82f39ca584766ecb99c1d327f75c2ef3c8f35de1d498d
3
+ size 5081170971
model-00006-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e62695f20640e3c7004b62f9cfe5c9aaf2ab25ecce08b8ff5be6d4e75f585162
3
+ size 5064328219
model-00007-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fd3e10d357f725365bd09705ce4af92c11008b0adb92681702601fb2fea2d372
3
+ size 5064328167
model-00008-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e802c51a04c31410a0fb69a495cfcfaa9b2c134ddbca888f9de295d5463f6a74
3
+ size 5064328179
model-00009-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aeaa943ff2cbb4c08e3106eb143430e212d8290273bdd8b2c631d42f15778e12
3
+ size 5081170967
model-00010-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ab0a5e67a8c1b0c6234ad83008615dcac627288feeba4329f2cb58f989c2af35
3
+ size 5064328173
model-00011-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d25feaac0fa88400e53c6c84db94018813d0343b6f0ae4dc415a096b4c4fb32b
3
+ size 5064328207
model-00012-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4e7d3eddfae6108bd949fb9044b0853e341fff9290e18dcbd5dfe1a8ebb5a394
3
+ size 5064328203
model-00013-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5c04663e58796c20df7aa858bc96e749658a74be28e1c33b1fb63ebb230f26f3
3
+ size 5081170929
model-00014-of-00014.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:030a59809f2e98c4429b87ffeae3758011e54e2a194810c8f0dd77d34e397e93
3
+ size 954217395
model.safetensors.index.json ADDED
@@ -0,0 +1,686 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 66885234176,
4
+ "total_parameters": 33442607104
5
+ },
6
+ "weight_map": {
7
+ "lm_head.weight": "model-00014-of-00014.safetensors",
8
+ "model.embed_tokens.weight": "model-00001-of-00014.safetensors",
9
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00014.safetensors",
10
+ "model.layers.0.mlp.down_proj.weight": "model-00001-of-00014.safetensors",
11
+ "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
12
+ "model.layers.0.mlp.up_proj.weight": "model-00001-of-00014.safetensors",
13
+ "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
14
+ "model.layers.0.self_attn.g_proj.weight": "model-00001-of-00014.safetensors",
15
+ "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00014.safetensors",
16
+ "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
17
+ "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
18
+ "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00014.safetensors",
19
+ "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
20
+ "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
21
+ "model.layers.1.input_layernorm.weight": "model-00001-of-00014.safetensors",
22
+ "model.layers.1.mlp.gate.e_score_correction_bias": "model-00001-of-00014.safetensors",
23
+ "model.layers.1.mlp.gate.proj.weight": "model-00001-of-00014.safetensors",
24
+ "model.layers.1.mlp.shared_expert.down_proj.weight": "model-00001-of-00014.safetensors",
25
+ "model.layers.1.mlp.shared_expert.gate_proj.weight": "model-00001-of-00014.safetensors",
26
+ "model.layers.1.mlp.shared_expert.up_proj.weight": "model-00001-of-00014.safetensors",
27
+ "model.layers.1.mlp.switch_mlp.down_proj.weight": "model-00001-of-00014.safetensors",
28
+ "model.layers.1.mlp.switch_mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
29
+ "model.layers.1.mlp.switch_mlp.up_proj.weight": "model-00001-of-00014.safetensors",
30
+ "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
31
+ "model.layers.1.self_attn.g_proj.weight": "model-00001-of-00014.safetensors",
32
+ "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00014.safetensors",
33
+ "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
34
+ "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
35
+ "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00014.safetensors",
36
+ "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
37
+ "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
38
+ "model.layers.10.input_layernorm.weight": "model-00004-of-00014.safetensors",
39
+ "model.layers.10.mlp.gate.e_score_correction_bias": "model-00004-of-00014.safetensors",
40
+ "model.layers.10.mlp.gate.proj.weight": "model-00004-of-00014.safetensors",
41
+ "model.layers.10.mlp.shared_expert.down_proj.weight": "model-00004-of-00014.safetensors",
42
+ "model.layers.10.mlp.shared_expert.gate_proj.weight": "model-00004-of-00014.safetensors",
43
+ "model.layers.10.mlp.shared_expert.up_proj.weight": "model-00004-of-00014.safetensors",
44
+ "model.layers.10.mlp.switch_mlp.down_proj.weight": "model-00004-of-00014.safetensors",
45
+ "model.layers.10.mlp.switch_mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
46
+ "model.layers.10.mlp.switch_mlp.up_proj.weight": "model-00004-of-00014.safetensors",
47
+ "model.layers.10.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
48
+ "model.layers.10.self_attn.g_proj.weight": "model-00004-of-00014.safetensors",
49
+ "model.layers.10.self_attn.k_norm.weight": "model-00004-of-00014.safetensors",
50
+ "model.layers.10.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
51
+ "model.layers.10.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
52
+ "model.layers.10.self_attn.q_norm.weight": "model-00004-of-00014.safetensors",
53
+ "model.layers.10.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
54
+ "model.layers.10.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
55
+ "model.layers.11.input_layernorm.weight": "model-00004-of-00014.safetensors",
56
+ "model.layers.11.mlp.gate.e_score_correction_bias": "model-00004-of-00014.safetensors",
57
+ "model.layers.11.mlp.gate.proj.weight": "model-00004-of-00014.safetensors",
58
+ "model.layers.11.mlp.shared_expert.down_proj.weight": "model-00004-of-00014.safetensors",
59
+ "model.layers.11.mlp.shared_expert.gate_proj.weight": "model-00004-of-00014.safetensors",
60
+ "model.layers.11.mlp.shared_expert.up_proj.weight": "model-00004-of-00014.safetensors",
61
+ "model.layers.11.mlp.switch_mlp.down_proj.weight": "model-00004-of-00014.safetensors",
62
+ "model.layers.11.mlp.switch_mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
63
+ "model.layers.11.mlp.switch_mlp.up_proj.weight": "model-00004-of-00014.safetensors",
64
+ "model.layers.11.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
65
+ "model.layers.11.self_attn.g_proj.weight": "model-00004-of-00014.safetensors",
66
+ "model.layers.11.self_attn.k_norm.weight": "model-00004-of-00014.safetensors",
67
+ "model.layers.11.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
68
+ "model.layers.11.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
69
+ "model.layers.11.self_attn.q_norm.weight": "model-00004-of-00014.safetensors",
70
+ "model.layers.11.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
71
+ "model.layers.11.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
72
+ "model.layers.12.input_layernorm.weight": "model-00005-of-00014.safetensors",
73
+ "model.layers.12.mlp.gate.e_score_correction_bias": "model-00004-of-00014.safetensors",
74
+ "model.layers.12.mlp.gate.proj.weight": "model-00004-of-00014.safetensors",
75
+ "model.layers.12.mlp.shared_expert.down_proj.weight": "model-00005-of-00014.safetensors",
76
+ "model.layers.12.mlp.shared_expert.gate_proj.weight": "model-00005-of-00014.safetensors",
77
+ "model.layers.12.mlp.shared_expert.up_proj.weight": "model-00005-of-00014.safetensors",
78
+ "model.layers.12.mlp.switch_mlp.down_proj.weight": "model-00005-of-00014.safetensors",
79
+ "model.layers.12.mlp.switch_mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
80
+ "model.layers.12.mlp.switch_mlp.up_proj.weight": "model-00004-of-00014.safetensors",
81
+ "model.layers.12.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
82
+ "model.layers.12.self_attn.g_proj.weight": "model-00004-of-00014.safetensors",
83
+ "model.layers.12.self_attn.k_norm.weight": "model-00004-of-00014.safetensors",
84
+ "model.layers.12.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
85
+ "model.layers.12.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
86
+ "model.layers.12.self_attn.q_norm.weight": "model-00004-of-00014.safetensors",
87
+ "model.layers.12.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
88
+ "model.layers.12.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
89
+ "model.layers.13.input_layernorm.weight": "model-00005-of-00014.safetensors",
90
+ "model.layers.13.mlp.gate.e_score_correction_bias": "model-00005-of-00014.safetensors",
91
+ "model.layers.13.mlp.gate.proj.weight": "model-00005-of-00014.safetensors",
92
+ "model.layers.13.mlp.shared_expert.down_proj.weight": "model-00005-of-00014.safetensors",
93
+ "model.layers.13.mlp.shared_expert.gate_proj.weight": "model-00005-of-00014.safetensors",
94
+ "model.layers.13.mlp.shared_expert.up_proj.weight": "model-00005-of-00014.safetensors",
95
+ "model.layers.13.mlp.switch_mlp.down_proj.weight": "model-00005-of-00014.safetensors",
96
+ "model.layers.13.mlp.switch_mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
97
+ "model.layers.13.mlp.switch_mlp.up_proj.weight": "model-00005-of-00014.safetensors",
98
+ "model.layers.13.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
99
+ "model.layers.13.self_attn.g_proj.weight": "model-00005-of-00014.safetensors",
100
+ "model.layers.13.self_attn.k_norm.weight": "model-00005-of-00014.safetensors",
101
+ "model.layers.13.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
102
+ "model.layers.13.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
103
+ "model.layers.13.self_attn.q_norm.weight": "model-00005-of-00014.safetensors",
104
+ "model.layers.13.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
105
+ "model.layers.13.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
106
+ "model.layers.14.input_layernorm.weight": "model-00005-of-00014.safetensors",
107
+ "model.layers.14.mlp.gate.e_score_correction_bias": "model-00005-of-00014.safetensors",
108
+ "model.layers.14.mlp.gate.proj.weight": "model-00005-of-00014.safetensors",
109
+ "model.layers.14.mlp.shared_expert.down_proj.weight": "model-00005-of-00014.safetensors",
110
+ "model.layers.14.mlp.shared_expert.gate_proj.weight": "model-00005-of-00014.safetensors",
111
+ "model.layers.14.mlp.shared_expert.up_proj.weight": "model-00005-of-00014.safetensors",
112
+ "model.layers.14.mlp.switch_mlp.down_proj.weight": "model-00005-of-00014.safetensors",
113
+ "model.layers.14.mlp.switch_mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
114
+ "model.layers.14.mlp.switch_mlp.up_proj.weight": "model-00005-of-00014.safetensors",
115
+ "model.layers.14.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
116
+ "model.layers.14.self_attn.g_proj.weight": "model-00005-of-00014.safetensors",
117
+ "model.layers.14.self_attn.k_norm.weight": "model-00005-of-00014.safetensors",
118
+ "model.layers.14.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
119
+ "model.layers.14.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
120
+ "model.layers.14.self_attn.q_norm.weight": "model-00005-of-00014.safetensors",
121
+ "model.layers.14.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
122
+ "model.layers.14.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
123
+ "model.layers.15.input_layernorm.weight": "model-00006-of-00014.safetensors",
124
+ "model.layers.15.mlp.gate.e_score_correction_bias": "model-00005-of-00014.safetensors",
125
+ "model.layers.15.mlp.gate.proj.weight": "model-00005-of-00014.safetensors",
126
+ "model.layers.15.mlp.shared_expert.down_proj.weight": "model-00006-of-00014.safetensors",
127
+ "model.layers.15.mlp.shared_expert.gate_proj.weight": "model-00006-of-00014.safetensors",
128
+ "model.layers.15.mlp.shared_expert.up_proj.weight": "model-00006-of-00014.safetensors",
129
+ "model.layers.15.mlp.switch_mlp.down_proj.weight": "model-00006-of-00014.safetensors",
130
+ "model.layers.15.mlp.switch_mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
131
+ "model.layers.15.mlp.switch_mlp.up_proj.weight": "model-00005-of-00014.safetensors",
132
+ "model.layers.15.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
133
+ "model.layers.15.self_attn.g_proj.weight": "model-00005-of-00014.safetensors",
134
+ "model.layers.15.self_attn.k_norm.weight": "model-00005-of-00014.safetensors",
135
+ "model.layers.15.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
136
+ "model.layers.15.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
137
+ "model.layers.15.self_attn.q_norm.weight": "model-00005-of-00014.safetensors",
138
+ "model.layers.15.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
139
+ "model.layers.15.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
140
+ "model.layers.16.input_layernorm.weight": "model-00006-of-00014.safetensors",
141
+ "model.layers.16.mlp.gate.e_score_correction_bias": "model-00006-of-00014.safetensors",
142
+ "model.layers.16.mlp.gate.proj.weight": "model-00006-of-00014.safetensors",
143
+ "model.layers.16.mlp.shared_expert.down_proj.weight": "model-00006-of-00014.safetensors",
144
+ "model.layers.16.mlp.shared_expert.gate_proj.weight": "model-00006-of-00014.safetensors",
145
+ "model.layers.16.mlp.shared_expert.up_proj.weight": "model-00006-of-00014.safetensors",
146
+ "model.layers.16.mlp.switch_mlp.down_proj.weight": "model-00006-of-00014.safetensors",
147
+ "model.layers.16.mlp.switch_mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
148
+ "model.layers.16.mlp.switch_mlp.up_proj.weight": "model-00006-of-00014.safetensors",
149
+ "model.layers.16.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
150
+ "model.layers.16.self_attn.g_proj.weight": "model-00006-of-00014.safetensors",
151
+ "model.layers.16.self_attn.k_norm.weight": "model-00006-of-00014.safetensors",
152
+ "model.layers.16.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
153
+ "model.layers.16.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
154
+ "model.layers.16.self_attn.q_norm.weight": "model-00006-of-00014.safetensors",
155
+ "model.layers.16.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
156
+ "model.layers.16.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
157
+ "model.layers.17.input_layernorm.weight": "model-00006-of-00014.safetensors",
158
+ "model.layers.17.mlp.gate.e_score_correction_bias": "model-00006-of-00014.safetensors",
159
+ "model.layers.17.mlp.gate.proj.weight": "model-00006-of-00014.safetensors",
160
+ "model.layers.17.mlp.shared_expert.down_proj.weight": "model-00006-of-00014.safetensors",
161
+ "model.layers.17.mlp.shared_expert.gate_proj.weight": "model-00006-of-00014.safetensors",
162
+ "model.layers.17.mlp.shared_expert.up_proj.weight": "model-00006-of-00014.safetensors",
163
+ "model.layers.17.mlp.switch_mlp.down_proj.weight": "model-00006-of-00014.safetensors",
164
+ "model.layers.17.mlp.switch_mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
165
+ "model.layers.17.mlp.switch_mlp.up_proj.weight": "model-00006-of-00014.safetensors",
166
+ "model.layers.17.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
167
+ "model.layers.17.self_attn.g_proj.weight": "model-00006-of-00014.safetensors",
168
+ "model.layers.17.self_attn.k_norm.weight": "model-00006-of-00014.safetensors",
169
+ "model.layers.17.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
170
+ "model.layers.17.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
171
+ "model.layers.17.self_attn.q_norm.weight": "model-00006-of-00014.safetensors",
172
+ "model.layers.17.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
173
+ "model.layers.17.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
174
+ "model.layers.18.input_layernorm.weight": "model-00007-of-00014.safetensors",
175
+ "model.layers.18.mlp.gate.e_score_correction_bias": "model-00006-of-00014.safetensors",
176
+ "model.layers.18.mlp.gate.proj.weight": "model-00006-of-00014.safetensors",
177
+ "model.layers.18.mlp.shared_expert.down_proj.weight": "model-00007-of-00014.safetensors",
178
+ "model.layers.18.mlp.shared_expert.gate_proj.weight": "model-00007-of-00014.safetensors",
179
+ "model.layers.18.mlp.shared_expert.up_proj.weight": "model-00007-of-00014.safetensors",
180
+ "model.layers.18.mlp.switch_mlp.down_proj.weight": "model-00007-of-00014.safetensors",
181
+ "model.layers.18.mlp.switch_mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
182
+ "model.layers.18.mlp.switch_mlp.up_proj.weight": "model-00006-of-00014.safetensors",
183
+ "model.layers.18.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
184
+ "model.layers.18.self_attn.g_proj.weight": "model-00006-of-00014.safetensors",
185
+ "model.layers.18.self_attn.k_norm.weight": "model-00006-of-00014.safetensors",
186
+ "model.layers.18.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
187
+ "model.layers.18.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
188
+ "model.layers.18.self_attn.q_norm.weight": "model-00006-of-00014.safetensors",
189
+ "model.layers.18.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
190
+ "model.layers.18.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
191
+ "model.layers.19.input_layernorm.weight": "model-00007-of-00014.safetensors",
192
+ "model.layers.19.mlp.gate.e_score_correction_bias": "model-00007-of-00014.safetensors",
193
+ "model.layers.19.mlp.gate.proj.weight": "model-00007-of-00014.safetensors",
194
+ "model.layers.19.mlp.shared_expert.down_proj.weight": "model-00007-of-00014.safetensors",
195
+ "model.layers.19.mlp.shared_expert.gate_proj.weight": "model-00007-of-00014.safetensors",
196
+ "model.layers.19.mlp.shared_expert.up_proj.weight": "model-00007-of-00014.safetensors",
197
+ "model.layers.19.mlp.switch_mlp.down_proj.weight": "model-00007-of-00014.safetensors",
198
+ "model.layers.19.mlp.switch_mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
199
+ "model.layers.19.mlp.switch_mlp.up_proj.weight": "model-00007-of-00014.safetensors",
200
+ "model.layers.19.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
201
+ "model.layers.19.self_attn.g_proj.weight": "model-00007-of-00014.safetensors",
202
+ "model.layers.19.self_attn.k_norm.weight": "model-00007-of-00014.safetensors",
203
+ "model.layers.19.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
204
+ "model.layers.19.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
205
+ "model.layers.19.self_attn.q_norm.weight": "model-00007-of-00014.safetensors",
206
+ "model.layers.19.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
207
+ "model.layers.19.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
208
+ "model.layers.2.input_layernorm.weight": "model-00001-of-00014.safetensors",
209
+ "model.layers.2.mlp.gate.e_score_correction_bias": "model-00001-of-00014.safetensors",
210
+ "model.layers.2.mlp.gate.proj.weight": "model-00001-of-00014.safetensors",
211
+ "model.layers.2.mlp.shared_expert.down_proj.weight": "model-00001-of-00014.safetensors",
212
+ "model.layers.2.mlp.shared_expert.gate_proj.weight": "model-00001-of-00014.safetensors",
213
+ "model.layers.2.mlp.shared_expert.up_proj.weight": "model-00001-of-00014.safetensors",
214
+ "model.layers.2.mlp.switch_mlp.down_proj.weight": "model-00001-of-00014.safetensors",
215
+ "model.layers.2.mlp.switch_mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
216
+ "model.layers.2.mlp.switch_mlp.up_proj.weight": "model-00001-of-00014.safetensors",
217
+ "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
218
+ "model.layers.2.self_attn.g_proj.weight": "model-00001-of-00014.safetensors",
219
+ "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00014.safetensors",
220
+ "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
221
+ "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
222
+ "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00014.safetensors",
223
+ "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
224
+ "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
225
+ "model.layers.20.input_layernorm.weight": "model-00007-of-00014.safetensors",
226
+ "model.layers.20.mlp.gate.e_score_correction_bias": "model-00007-of-00014.safetensors",
227
+ "model.layers.20.mlp.gate.proj.weight": "model-00007-of-00014.safetensors",
228
+ "model.layers.20.mlp.shared_expert.down_proj.weight": "model-00007-of-00014.safetensors",
229
+ "model.layers.20.mlp.shared_expert.gate_proj.weight": "model-00007-of-00014.safetensors",
230
+ "model.layers.20.mlp.shared_expert.up_proj.weight": "model-00007-of-00014.safetensors",
231
+ "model.layers.20.mlp.switch_mlp.down_proj.weight": "model-00007-of-00014.safetensors",
232
+ "model.layers.20.mlp.switch_mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
233
+ "model.layers.20.mlp.switch_mlp.up_proj.weight": "model-00007-of-00014.safetensors",
234
+ "model.layers.20.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
235
+ "model.layers.20.self_attn.g_proj.weight": "model-00007-of-00014.safetensors",
236
+ "model.layers.20.self_attn.k_norm.weight": "model-00007-of-00014.safetensors",
237
+ "model.layers.20.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
238
+ "model.layers.20.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
239
+ "model.layers.20.self_attn.q_norm.weight": "model-00007-of-00014.safetensors",
240
+ "model.layers.20.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
241
+ "model.layers.20.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
242
+ "model.layers.21.input_layernorm.weight": "model-00008-of-00014.safetensors",
243
+ "model.layers.21.mlp.gate.e_score_correction_bias": "model-00007-of-00014.safetensors",
244
+ "model.layers.21.mlp.gate.proj.weight": "model-00007-of-00014.safetensors",
245
+ "model.layers.21.mlp.shared_expert.down_proj.weight": "model-00008-of-00014.safetensors",
246
+ "model.layers.21.mlp.shared_expert.gate_proj.weight": "model-00008-of-00014.safetensors",
247
+ "model.layers.21.mlp.shared_expert.up_proj.weight": "model-00008-of-00014.safetensors",
248
+ "model.layers.21.mlp.switch_mlp.down_proj.weight": "model-00008-of-00014.safetensors",
249
+ "model.layers.21.mlp.switch_mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
250
+ "model.layers.21.mlp.switch_mlp.up_proj.weight": "model-00007-of-00014.safetensors",
251
+ "model.layers.21.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
252
+ "model.layers.21.self_attn.g_proj.weight": "model-00007-of-00014.safetensors",
253
+ "model.layers.21.self_attn.k_norm.weight": "model-00007-of-00014.safetensors",
254
+ "model.layers.21.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
255
+ "model.layers.21.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
256
+ "model.layers.21.self_attn.q_norm.weight": "model-00007-of-00014.safetensors",
257
+ "model.layers.21.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
258
+ "model.layers.21.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
259
+ "model.layers.22.input_layernorm.weight": "model-00008-of-00014.safetensors",
260
+ "model.layers.22.mlp.gate.e_score_correction_bias": "model-00008-of-00014.safetensors",
261
+ "model.layers.22.mlp.gate.proj.weight": "model-00008-of-00014.safetensors",
262
+ "model.layers.22.mlp.shared_expert.down_proj.weight": "model-00008-of-00014.safetensors",
263
+ "model.layers.22.mlp.shared_expert.gate_proj.weight": "model-00008-of-00014.safetensors",
264
+ "model.layers.22.mlp.shared_expert.up_proj.weight": "model-00008-of-00014.safetensors",
265
+ "model.layers.22.mlp.switch_mlp.down_proj.weight": "model-00008-of-00014.safetensors",
266
+ "model.layers.22.mlp.switch_mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
267
+ "model.layers.22.mlp.switch_mlp.up_proj.weight": "model-00008-of-00014.safetensors",
268
+ "model.layers.22.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
269
+ "model.layers.22.self_attn.g_proj.weight": "model-00008-of-00014.safetensors",
270
+ "model.layers.22.self_attn.k_norm.weight": "model-00008-of-00014.safetensors",
271
+ "model.layers.22.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
272
+ "model.layers.22.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
273
+ "model.layers.22.self_attn.q_norm.weight": "model-00008-of-00014.safetensors",
274
+ "model.layers.22.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
275
+ "model.layers.22.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
276
+ "model.layers.23.input_layernorm.weight": "model-00008-of-00014.safetensors",
277
+ "model.layers.23.mlp.gate.e_score_correction_bias": "model-00008-of-00014.safetensors",
278
+ "model.layers.23.mlp.gate.proj.weight": "model-00008-of-00014.safetensors",
279
+ "model.layers.23.mlp.shared_expert.down_proj.weight": "model-00008-of-00014.safetensors",
280
+ "model.layers.23.mlp.shared_expert.gate_proj.weight": "model-00008-of-00014.safetensors",
281
+ "model.layers.23.mlp.shared_expert.up_proj.weight": "model-00008-of-00014.safetensors",
282
+ "model.layers.23.mlp.switch_mlp.down_proj.weight": "model-00008-of-00014.safetensors",
283
+ "model.layers.23.mlp.switch_mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
284
+ "model.layers.23.mlp.switch_mlp.up_proj.weight": "model-00008-of-00014.safetensors",
285
+ "model.layers.23.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
286
+ "model.layers.23.self_attn.g_proj.weight": "model-00008-of-00014.safetensors",
287
+ "model.layers.23.self_attn.k_norm.weight": "model-00008-of-00014.safetensors",
288
+ "model.layers.23.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
289
+ "model.layers.23.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
290
+ "model.layers.23.self_attn.q_norm.weight": "model-00008-of-00014.safetensors",
291
+ "model.layers.23.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
292
+ "model.layers.23.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
293
+ "model.layers.24.input_layernorm.weight": "model-00009-of-00014.safetensors",
294
+ "model.layers.24.mlp.gate.e_score_correction_bias": "model-00008-of-00014.safetensors",
295
+ "model.layers.24.mlp.gate.proj.weight": "model-00008-of-00014.safetensors",
296
+ "model.layers.24.mlp.shared_expert.down_proj.weight": "model-00009-of-00014.safetensors",
297
+ "model.layers.24.mlp.shared_expert.gate_proj.weight": "model-00009-of-00014.safetensors",
298
+ "model.layers.24.mlp.shared_expert.up_proj.weight": "model-00009-of-00014.safetensors",
299
+ "model.layers.24.mlp.switch_mlp.down_proj.weight": "model-00009-of-00014.safetensors",
300
+ "model.layers.24.mlp.switch_mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
301
+ "model.layers.24.mlp.switch_mlp.up_proj.weight": "model-00008-of-00014.safetensors",
302
+ "model.layers.24.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
303
+ "model.layers.24.self_attn.g_proj.weight": "model-00008-of-00014.safetensors",
304
+ "model.layers.24.self_attn.k_norm.weight": "model-00008-of-00014.safetensors",
305
+ "model.layers.24.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
306
+ "model.layers.24.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
307
+ "model.layers.24.self_attn.q_norm.weight": "model-00008-of-00014.safetensors",
308
+ "model.layers.24.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
309
+ "model.layers.24.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
310
+ "model.layers.25.input_layernorm.weight": "model-00009-of-00014.safetensors",
311
+ "model.layers.25.mlp.gate.e_score_correction_bias": "model-00009-of-00014.safetensors",
312
+ "model.layers.25.mlp.gate.proj.weight": "model-00009-of-00014.safetensors",
313
+ "model.layers.25.mlp.shared_expert.down_proj.weight": "model-00009-of-00014.safetensors",
314
+ "model.layers.25.mlp.shared_expert.gate_proj.weight": "model-00009-of-00014.safetensors",
315
+ "model.layers.25.mlp.shared_expert.up_proj.weight": "model-00009-of-00014.safetensors",
316
+ "model.layers.25.mlp.switch_mlp.down_proj.weight": "model-00009-of-00014.safetensors",
317
+ "model.layers.25.mlp.switch_mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
318
+ "model.layers.25.mlp.switch_mlp.up_proj.weight": "model-00009-of-00014.safetensors",
319
+ "model.layers.25.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
320
+ "model.layers.25.self_attn.g_proj.weight": "model-00009-of-00014.safetensors",
321
+ "model.layers.25.self_attn.k_norm.weight": "model-00009-of-00014.safetensors",
322
+ "model.layers.25.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
323
+ "model.layers.25.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
324
+ "model.layers.25.self_attn.q_norm.weight": "model-00009-of-00014.safetensors",
325
+ "model.layers.25.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
326
+ "model.layers.25.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
327
+ "model.layers.26.input_layernorm.weight": "model-00009-of-00014.safetensors",
328
+ "model.layers.26.mlp.gate.e_score_correction_bias": "model-00009-of-00014.safetensors",
329
+ "model.layers.26.mlp.gate.proj.weight": "model-00009-of-00014.safetensors",
330
+ "model.layers.26.mlp.shared_expert.down_proj.weight": "model-00009-of-00014.safetensors",
331
+ "model.layers.26.mlp.shared_expert.gate_proj.weight": "model-00009-of-00014.safetensors",
332
+ "model.layers.26.mlp.shared_expert.up_proj.weight": "model-00009-of-00014.safetensors",
333
+ "model.layers.26.mlp.switch_mlp.down_proj.weight": "model-00009-of-00014.safetensors",
334
+ "model.layers.26.mlp.switch_mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
335
+ "model.layers.26.mlp.switch_mlp.up_proj.weight": "model-00009-of-00014.safetensors",
336
+ "model.layers.26.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
337
+ "model.layers.26.self_attn.g_proj.weight": "model-00009-of-00014.safetensors",
338
+ "model.layers.26.self_attn.k_norm.weight": "model-00009-of-00014.safetensors",
339
+ "model.layers.26.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
340
+ "model.layers.26.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
341
+ "model.layers.26.self_attn.q_norm.weight": "model-00009-of-00014.safetensors",
342
+ "model.layers.26.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
343
+ "model.layers.26.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
344
+ "model.layers.27.input_layernorm.weight": "model-00010-of-00014.safetensors",
345
+ "model.layers.27.mlp.gate.e_score_correction_bias": "model-00009-of-00014.safetensors",
346
+ "model.layers.27.mlp.gate.proj.weight": "model-00009-of-00014.safetensors",
347
+ "model.layers.27.mlp.shared_expert.down_proj.weight": "model-00010-of-00014.safetensors",
348
+ "model.layers.27.mlp.shared_expert.gate_proj.weight": "model-00010-of-00014.safetensors",
349
+ "model.layers.27.mlp.shared_expert.up_proj.weight": "model-00010-of-00014.safetensors",
350
+ "model.layers.27.mlp.switch_mlp.down_proj.weight": "model-00010-of-00014.safetensors",
351
+ "model.layers.27.mlp.switch_mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
352
+ "model.layers.27.mlp.switch_mlp.up_proj.weight": "model-00009-of-00014.safetensors",
353
+ "model.layers.27.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
354
+ "model.layers.27.self_attn.g_proj.weight": "model-00009-of-00014.safetensors",
355
+ "model.layers.27.self_attn.k_norm.weight": "model-00009-of-00014.safetensors",
356
+ "model.layers.27.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
357
+ "model.layers.27.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
358
+ "model.layers.27.self_attn.q_norm.weight": "model-00009-of-00014.safetensors",
359
+ "model.layers.27.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
360
+ "model.layers.27.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
361
+ "model.layers.28.input_layernorm.weight": "model-00010-of-00014.safetensors",
362
+ "model.layers.28.mlp.gate.e_score_correction_bias": "model-00010-of-00014.safetensors",
363
+ "model.layers.28.mlp.gate.proj.weight": "model-00010-of-00014.safetensors",
364
+ "model.layers.28.mlp.shared_expert.down_proj.weight": "model-00010-of-00014.safetensors",
365
+ "model.layers.28.mlp.shared_expert.gate_proj.weight": "model-00010-of-00014.safetensors",
366
+ "model.layers.28.mlp.shared_expert.up_proj.weight": "model-00010-of-00014.safetensors",
367
+ "model.layers.28.mlp.switch_mlp.down_proj.weight": "model-00010-of-00014.safetensors",
368
+ "model.layers.28.mlp.switch_mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
369
+ "model.layers.28.mlp.switch_mlp.up_proj.weight": "model-00010-of-00014.safetensors",
370
+ "model.layers.28.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
371
+ "model.layers.28.self_attn.g_proj.weight": "model-00010-of-00014.safetensors",
372
+ "model.layers.28.self_attn.k_norm.weight": "model-00010-of-00014.safetensors",
373
+ "model.layers.28.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
374
+ "model.layers.28.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
375
+ "model.layers.28.self_attn.q_norm.weight": "model-00010-of-00014.safetensors",
376
+ "model.layers.28.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
377
+ "model.layers.28.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
378
+ "model.layers.29.input_layernorm.weight": "model-00010-of-00014.safetensors",
379
+ "model.layers.29.mlp.gate.e_score_correction_bias": "model-00010-of-00014.safetensors",
380
+ "model.layers.29.mlp.gate.proj.weight": "model-00010-of-00014.safetensors",
381
+ "model.layers.29.mlp.shared_expert.down_proj.weight": "model-00010-of-00014.safetensors",
382
+ "model.layers.29.mlp.shared_expert.gate_proj.weight": "model-00010-of-00014.safetensors",
383
+ "model.layers.29.mlp.shared_expert.up_proj.weight": "model-00010-of-00014.safetensors",
384
+ "model.layers.29.mlp.switch_mlp.down_proj.weight": "model-00010-of-00014.safetensors",
385
+ "model.layers.29.mlp.switch_mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
386
+ "model.layers.29.mlp.switch_mlp.up_proj.weight": "model-00010-of-00014.safetensors",
387
+ "model.layers.29.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
388
+ "model.layers.29.self_attn.g_proj.weight": "model-00010-of-00014.safetensors",
389
+ "model.layers.29.self_attn.k_norm.weight": "model-00010-of-00014.safetensors",
390
+ "model.layers.29.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
391
+ "model.layers.29.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
392
+ "model.layers.29.self_attn.q_norm.weight": "model-00010-of-00014.safetensors",
393
+ "model.layers.29.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
394
+ "model.layers.29.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
395
+ "model.layers.3.input_layernorm.weight": "model-00002-of-00014.safetensors",
396
+ "model.layers.3.mlp.gate.e_score_correction_bias": "model-00001-of-00014.safetensors",
397
+ "model.layers.3.mlp.gate.proj.weight": "model-00001-of-00014.safetensors",
398
+ "model.layers.3.mlp.shared_expert.down_proj.weight": "model-00002-of-00014.safetensors",
399
+ "model.layers.3.mlp.shared_expert.gate_proj.weight": "model-00002-of-00014.safetensors",
400
+ "model.layers.3.mlp.shared_expert.up_proj.weight": "model-00002-of-00014.safetensors",
401
+ "model.layers.3.mlp.switch_mlp.down_proj.weight": "model-00002-of-00014.safetensors",
402
+ "model.layers.3.mlp.switch_mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
403
+ "model.layers.3.mlp.switch_mlp.up_proj.weight": "model-00001-of-00014.safetensors",
404
+ "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
405
+ "model.layers.3.self_attn.g_proj.weight": "model-00001-of-00014.safetensors",
406
+ "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00014.safetensors",
407
+ "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
408
+ "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
409
+ "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00014.safetensors",
410
+ "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
411
+ "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
412
+ "model.layers.30.input_layernorm.weight": "model-00011-of-00014.safetensors",
413
+ "model.layers.30.mlp.gate.e_score_correction_bias": "model-00010-of-00014.safetensors",
414
+ "model.layers.30.mlp.gate.proj.weight": "model-00010-of-00014.safetensors",
415
+ "model.layers.30.mlp.shared_expert.down_proj.weight": "model-00011-of-00014.safetensors",
416
+ "model.layers.30.mlp.shared_expert.gate_proj.weight": "model-00011-of-00014.safetensors",
417
+ "model.layers.30.mlp.shared_expert.up_proj.weight": "model-00011-of-00014.safetensors",
418
+ "model.layers.30.mlp.switch_mlp.down_proj.weight": "model-00011-of-00014.safetensors",
419
+ "model.layers.30.mlp.switch_mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
420
+ "model.layers.30.mlp.switch_mlp.up_proj.weight": "model-00010-of-00014.safetensors",
421
+ "model.layers.30.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
422
+ "model.layers.30.self_attn.g_proj.weight": "model-00010-of-00014.safetensors",
423
+ "model.layers.30.self_attn.k_norm.weight": "model-00010-of-00014.safetensors",
424
+ "model.layers.30.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
425
+ "model.layers.30.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
426
+ "model.layers.30.self_attn.q_norm.weight": "model-00010-of-00014.safetensors",
427
+ "model.layers.30.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
428
+ "model.layers.30.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
429
+ "model.layers.31.input_layernorm.weight": "model-00011-of-00014.safetensors",
430
+ "model.layers.31.mlp.gate.e_score_correction_bias": "model-00011-of-00014.safetensors",
431
+ "model.layers.31.mlp.gate.proj.weight": "model-00011-of-00014.safetensors",
432
+ "model.layers.31.mlp.shared_expert.down_proj.weight": "model-00011-of-00014.safetensors",
433
+ "model.layers.31.mlp.shared_expert.gate_proj.weight": "model-00011-of-00014.safetensors",
434
+ "model.layers.31.mlp.shared_expert.up_proj.weight": "model-00011-of-00014.safetensors",
435
+ "model.layers.31.mlp.switch_mlp.down_proj.weight": "model-00011-of-00014.safetensors",
436
+ "model.layers.31.mlp.switch_mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
437
+ "model.layers.31.mlp.switch_mlp.up_proj.weight": "model-00011-of-00014.safetensors",
438
+ "model.layers.31.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
439
+ "model.layers.31.self_attn.g_proj.weight": "model-00011-of-00014.safetensors",
440
+ "model.layers.31.self_attn.k_norm.weight": "model-00011-of-00014.safetensors",
441
+ "model.layers.31.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
442
+ "model.layers.31.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
443
+ "model.layers.31.self_attn.q_norm.weight": "model-00011-of-00014.safetensors",
444
+ "model.layers.31.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
445
+ "model.layers.31.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
446
+ "model.layers.32.input_layernorm.weight": "model-00011-of-00014.safetensors",
447
+ "model.layers.32.mlp.gate.e_score_correction_bias": "model-00011-of-00014.safetensors",
448
+ "model.layers.32.mlp.gate.proj.weight": "model-00011-of-00014.safetensors",
449
+ "model.layers.32.mlp.shared_expert.down_proj.weight": "model-00011-of-00014.safetensors",
450
+ "model.layers.32.mlp.shared_expert.gate_proj.weight": "model-00011-of-00014.safetensors",
451
+ "model.layers.32.mlp.shared_expert.up_proj.weight": "model-00011-of-00014.safetensors",
452
+ "model.layers.32.mlp.switch_mlp.down_proj.weight": "model-00011-of-00014.safetensors",
453
+ "model.layers.32.mlp.switch_mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
454
+ "model.layers.32.mlp.switch_mlp.up_proj.weight": "model-00011-of-00014.safetensors",
455
+ "model.layers.32.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
456
+ "model.layers.32.self_attn.g_proj.weight": "model-00011-of-00014.safetensors",
457
+ "model.layers.32.self_attn.k_norm.weight": "model-00011-of-00014.safetensors",
458
+ "model.layers.32.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
459
+ "model.layers.32.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
460
+ "model.layers.32.self_attn.q_norm.weight": "model-00011-of-00014.safetensors",
461
+ "model.layers.32.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
462
+ "model.layers.32.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
463
+ "model.layers.33.input_layernorm.weight": "model-00012-of-00014.safetensors",
464
+ "model.layers.33.mlp.gate.e_score_correction_bias": "model-00011-of-00014.safetensors",
465
+ "model.layers.33.mlp.gate.proj.weight": "model-00011-of-00014.safetensors",
466
+ "model.layers.33.mlp.shared_expert.down_proj.weight": "model-00012-of-00014.safetensors",
467
+ "model.layers.33.mlp.shared_expert.gate_proj.weight": "model-00012-of-00014.safetensors",
468
+ "model.layers.33.mlp.shared_expert.up_proj.weight": "model-00012-of-00014.safetensors",
469
+ "model.layers.33.mlp.switch_mlp.down_proj.weight": "model-00012-of-00014.safetensors",
470
+ "model.layers.33.mlp.switch_mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
471
+ "model.layers.33.mlp.switch_mlp.up_proj.weight": "model-00011-of-00014.safetensors",
472
+ "model.layers.33.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
473
+ "model.layers.33.self_attn.g_proj.weight": "model-00011-of-00014.safetensors",
474
+ "model.layers.33.self_attn.k_norm.weight": "model-00011-of-00014.safetensors",
475
+ "model.layers.33.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
476
+ "model.layers.33.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
477
+ "model.layers.33.self_attn.q_norm.weight": "model-00011-of-00014.safetensors",
478
+ "model.layers.33.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
479
+ "model.layers.33.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
480
+ "model.layers.34.input_layernorm.weight": "model-00012-of-00014.safetensors",
481
+ "model.layers.34.mlp.gate.e_score_correction_bias": "model-00012-of-00014.safetensors",
482
+ "model.layers.34.mlp.gate.proj.weight": "model-00012-of-00014.safetensors",
483
+ "model.layers.34.mlp.shared_expert.down_proj.weight": "model-00012-of-00014.safetensors",
484
+ "model.layers.34.mlp.shared_expert.gate_proj.weight": "model-00012-of-00014.safetensors",
485
+ "model.layers.34.mlp.shared_expert.up_proj.weight": "model-00012-of-00014.safetensors",
486
+ "model.layers.34.mlp.switch_mlp.down_proj.weight": "model-00012-of-00014.safetensors",
487
+ "model.layers.34.mlp.switch_mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
488
+ "model.layers.34.mlp.switch_mlp.up_proj.weight": "model-00012-of-00014.safetensors",
489
+ "model.layers.34.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
490
+ "model.layers.34.self_attn.g_proj.weight": "model-00012-of-00014.safetensors",
491
+ "model.layers.34.self_attn.k_norm.weight": "model-00012-of-00014.safetensors",
492
+ "model.layers.34.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
493
+ "model.layers.34.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
494
+ "model.layers.34.self_attn.q_norm.weight": "model-00012-of-00014.safetensors",
495
+ "model.layers.34.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
496
+ "model.layers.34.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
497
+ "model.layers.35.input_layernorm.weight": "model-00012-of-00014.safetensors",
498
+ "model.layers.35.mlp.gate.e_score_correction_bias": "model-00012-of-00014.safetensors",
499
+ "model.layers.35.mlp.gate.proj.weight": "model-00012-of-00014.safetensors",
500
+ "model.layers.35.mlp.shared_expert.down_proj.weight": "model-00012-of-00014.safetensors",
501
+ "model.layers.35.mlp.shared_expert.gate_proj.weight": "model-00012-of-00014.safetensors",
502
+ "model.layers.35.mlp.shared_expert.up_proj.weight": "model-00012-of-00014.safetensors",
503
+ "model.layers.35.mlp.switch_mlp.down_proj.weight": "model-00012-of-00014.safetensors",
504
+ "model.layers.35.mlp.switch_mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
505
+ "model.layers.35.mlp.switch_mlp.up_proj.weight": "model-00012-of-00014.safetensors",
506
+ "model.layers.35.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
507
+ "model.layers.35.self_attn.g_proj.weight": "model-00012-of-00014.safetensors",
508
+ "model.layers.35.self_attn.k_norm.weight": "model-00012-of-00014.safetensors",
509
+ "model.layers.35.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
510
+ "model.layers.35.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
511
+ "model.layers.35.self_attn.q_norm.weight": "model-00012-of-00014.safetensors",
512
+ "model.layers.35.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
513
+ "model.layers.35.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
514
+ "model.layers.36.input_layernorm.weight": "model-00013-of-00014.safetensors",
515
+ "model.layers.36.mlp.gate.e_score_correction_bias": "model-00012-of-00014.safetensors",
516
+ "model.layers.36.mlp.gate.proj.weight": "model-00012-of-00014.safetensors",
517
+ "model.layers.36.mlp.shared_expert.down_proj.weight": "model-00013-of-00014.safetensors",
518
+ "model.layers.36.mlp.shared_expert.gate_proj.weight": "model-00013-of-00014.safetensors",
519
+ "model.layers.36.mlp.shared_expert.up_proj.weight": "model-00013-of-00014.safetensors",
520
+ "model.layers.36.mlp.switch_mlp.down_proj.weight": "model-00013-of-00014.safetensors",
521
+ "model.layers.36.mlp.switch_mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
522
+ "model.layers.36.mlp.switch_mlp.up_proj.weight": "model-00012-of-00014.safetensors",
523
+ "model.layers.36.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
524
+ "model.layers.36.self_attn.g_proj.weight": "model-00012-of-00014.safetensors",
525
+ "model.layers.36.self_attn.k_norm.weight": "model-00012-of-00014.safetensors",
526
+ "model.layers.36.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
527
+ "model.layers.36.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
528
+ "model.layers.36.self_attn.q_norm.weight": "model-00012-of-00014.safetensors",
529
+ "model.layers.36.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
530
+ "model.layers.36.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
531
+ "model.layers.37.input_layernorm.weight": "model-00013-of-00014.safetensors",
532
+ "model.layers.37.mlp.gate.e_score_correction_bias": "model-00013-of-00014.safetensors",
533
+ "model.layers.37.mlp.gate.proj.weight": "model-00013-of-00014.safetensors",
534
+ "model.layers.37.mlp.shared_expert.down_proj.weight": "model-00013-of-00014.safetensors",
535
+ "model.layers.37.mlp.shared_expert.gate_proj.weight": "model-00013-of-00014.safetensors",
536
+ "model.layers.37.mlp.shared_expert.up_proj.weight": "model-00013-of-00014.safetensors",
537
+ "model.layers.37.mlp.switch_mlp.down_proj.weight": "model-00013-of-00014.safetensors",
538
+ "model.layers.37.mlp.switch_mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
539
+ "model.layers.37.mlp.switch_mlp.up_proj.weight": "model-00013-of-00014.safetensors",
540
+ "model.layers.37.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
541
+ "model.layers.37.self_attn.g_proj.weight": "model-00013-of-00014.safetensors",
542
+ "model.layers.37.self_attn.k_norm.weight": "model-00013-of-00014.safetensors",
543
+ "model.layers.37.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
544
+ "model.layers.37.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
545
+ "model.layers.37.self_attn.q_norm.weight": "model-00013-of-00014.safetensors",
546
+ "model.layers.37.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
547
+ "model.layers.37.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
548
+ "model.layers.38.input_layernorm.weight": "model-00013-of-00014.safetensors",
549
+ "model.layers.38.mlp.gate.e_score_correction_bias": "model-00013-of-00014.safetensors",
550
+ "model.layers.38.mlp.gate.proj.weight": "model-00013-of-00014.safetensors",
551
+ "model.layers.38.mlp.shared_expert.down_proj.weight": "model-00013-of-00014.safetensors",
552
+ "model.layers.38.mlp.shared_expert.gate_proj.weight": "model-00013-of-00014.safetensors",
553
+ "model.layers.38.mlp.shared_expert.up_proj.weight": "model-00013-of-00014.safetensors",
554
+ "model.layers.38.mlp.switch_mlp.down_proj.weight": "model-00013-of-00014.safetensors",
555
+ "model.layers.38.mlp.switch_mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
556
+ "model.layers.38.mlp.switch_mlp.up_proj.weight": "model-00013-of-00014.safetensors",
557
+ "model.layers.38.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
558
+ "model.layers.38.self_attn.g_proj.weight": "model-00013-of-00014.safetensors",
559
+ "model.layers.38.self_attn.k_norm.weight": "model-00013-of-00014.safetensors",
560
+ "model.layers.38.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
561
+ "model.layers.38.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
562
+ "model.layers.38.self_attn.q_norm.weight": "model-00013-of-00014.safetensors",
563
+ "model.layers.38.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
564
+ "model.layers.38.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
565
+ "model.layers.39.input_layernorm.weight": "model-00014-of-00014.safetensors",
566
+ "model.layers.39.mlp.gate.e_score_correction_bias": "model-00013-of-00014.safetensors",
567
+ "model.layers.39.mlp.gate.proj.weight": "model-00013-of-00014.safetensors",
568
+ "model.layers.39.mlp.shared_expert.down_proj.weight": "model-00014-of-00014.safetensors",
569
+ "model.layers.39.mlp.shared_expert.gate_proj.weight": "model-00014-of-00014.safetensors",
570
+ "model.layers.39.mlp.shared_expert.up_proj.weight": "model-00014-of-00014.safetensors",
571
+ "model.layers.39.mlp.switch_mlp.down_proj.weight": "model-00014-of-00014.safetensors",
572
+ "model.layers.39.mlp.switch_mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
573
+ "model.layers.39.mlp.switch_mlp.up_proj.weight": "model-00013-of-00014.safetensors",
574
+ "model.layers.39.post_attention_layernorm.weight": "model-00014-of-00014.safetensors",
575
+ "model.layers.39.self_attn.g_proj.weight": "model-00013-of-00014.safetensors",
576
+ "model.layers.39.self_attn.k_norm.weight": "model-00013-of-00014.safetensors",
577
+ "model.layers.39.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
578
+ "model.layers.39.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
579
+ "model.layers.39.self_attn.q_norm.weight": "model-00013-of-00014.safetensors",
580
+ "model.layers.39.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
581
+ "model.layers.39.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
582
+ "model.layers.4.input_layernorm.weight": "model-00002-of-00014.safetensors",
583
+ "model.layers.4.mlp.gate.e_score_correction_bias": "model-00002-of-00014.safetensors",
584
+ "model.layers.4.mlp.gate.proj.weight": "model-00002-of-00014.safetensors",
585
+ "model.layers.4.mlp.shared_expert.down_proj.weight": "model-00002-of-00014.safetensors",
586
+ "model.layers.4.mlp.shared_expert.gate_proj.weight": "model-00002-of-00014.safetensors",
587
+ "model.layers.4.mlp.shared_expert.up_proj.weight": "model-00002-of-00014.safetensors",
588
+ "model.layers.4.mlp.switch_mlp.down_proj.weight": "model-00002-of-00014.safetensors",
589
+ "model.layers.4.mlp.switch_mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
590
+ "model.layers.4.mlp.switch_mlp.up_proj.weight": "model-00002-of-00014.safetensors",
591
+ "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
592
+ "model.layers.4.self_attn.g_proj.weight": "model-00002-of-00014.safetensors",
593
+ "model.layers.4.self_attn.k_norm.weight": "model-00002-of-00014.safetensors",
594
+ "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
595
+ "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
596
+ "model.layers.4.self_attn.q_norm.weight": "model-00002-of-00014.safetensors",
597
+ "model.layers.4.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
598
+ "model.layers.4.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
599
+ "model.layers.5.input_layernorm.weight": "model-00002-of-00014.safetensors",
600
+ "model.layers.5.mlp.gate.e_score_correction_bias": "model-00002-of-00014.safetensors",
601
+ "model.layers.5.mlp.gate.proj.weight": "model-00002-of-00014.safetensors",
602
+ "model.layers.5.mlp.shared_expert.down_proj.weight": "model-00002-of-00014.safetensors",
603
+ "model.layers.5.mlp.shared_expert.gate_proj.weight": "model-00002-of-00014.safetensors",
604
+ "model.layers.5.mlp.shared_expert.up_proj.weight": "model-00002-of-00014.safetensors",
605
+ "model.layers.5.mlp.switch_mlp.down_proj.weight": "model-00002-of-00014.safetensors",
606
+ "model.layers.5.mlp.switch_mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
607
+ "model.layers.5.mlp.switch_mlp.up_proj.weight": "model-00002-of-00014.safetensors",
608
+ "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
609
+ "model.layers.5.self_attn.g_proj.weight": "model-00002-of-00014.safetensors",
610
+ "model.layers.5.self_attn.k_norm.weight": "model-00002-of-00014.safetensors",
611
+ "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
612
+ "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
613
+ "model.layers.5.self_attn.q_norm.weight": "model-00002-of-00014.safetensors",
614
+ "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
615
+ "model.layers.5.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
616
+ "model.layers.6.input_layernorm.weight": "model-00003-of-00014.safetensors",
617
+ "model.layers.6.mlp.gate.e_score_correction_bias": "model-00002-of-00014.safetensors",
618
+ "model.layers.6.mlp.gate.proj.weight": "model-00002-of-00014.safetensors",
619
+ "model.layers.6.mlp.shared_expert.down_proj.weight": "model-00003-of-00014.safetensors",
620
+ "model.layers.6.mlp.shared_expert.gate_proj.weight": "model-00003-of-00014.safetensors",
621
+ "model.layers.6.mlp.shared_expert.up_proj.weight": "model-00003-of-00014.safetensors",
622
+ "model.layers.6.mlp.switch_mlp.down_proj.weight": "model-00003-of-00014.safetensors",
623
+ "model.layers.6.mlp.switch_mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
624
+ "model.layers.6.mlp.switch_mlp.up_proj.weight": "model-00002-of-00014.safetensors",
625
+ "model.layers.6.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
626
+ "model.layers.6.self_attn.g_proj.weight": "model-00002-of-00014.safetensors",
627
+ "model.layers.6.self_attn.k_norm.weight": "model-00002-of-00014.safetensors",
628
+ "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
629
+ "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
630
+ "model.layers.6.self_attn.q_norm.weight": "model-00002-of-00014.safetensors",
631
+ "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
632
+ "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
633
+ "model.layers.7.input_layernorm.weight": "model-00003-of-00014.safetensors",
634
+ "model.layers.7.mlp.gate.e_score_correction_bias": "model-00003-of-00014.safetensors",
635
+ "model.layers.7.mlp.gate.proj.weight": "model-00003-of-00014.safetensors",
636
+ "model.layers.7.mlp.shared_expert.down_proj.weight": "model-00003-of-00014.safetensors",
637
+ "model.layers.7.mlp.shared_expert.gate_proj.weight": "model-00003-of-00014.safetensors",
638
+ "model.layers.7.mlp.shared_expert.up_proj.weight": "model-00003-of-00014.safetensors",
639
+ "model.layers.7.mlp.switch_mlp.down_proj.weight": "model-00003-of-00014.safetensors",
640
+ "model.layers.7.mlp.switch_mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
641
+ "model.layers.7.mlp.switch_mlp.up_proj.weight": "model-00003-of-00014.safetensors",
642
+ "model.layers.7.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
643
+ "model.layers.7.self_attn.g_proj.weight": "model-00003-of-00014.safetensors",
644
+ "model.layers.7.self_attn.k_norm.weight": "model-00003-of-00014.safetensors",
645
+ "model.layers.7.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
646
+ "model.layers.7.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
647
+ "model.layers.7.self_attn.q_norm.weight": "model-00003-of-00014.safetensors",
648
+ "model.layers.7.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
649
+ "model.layers.7.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
650
+ "model.layers.8.input_layernorm.weight": "model-00003-of-00014.safetensors",
651
+ "model.layers.8.mlp.gate.e_score_correction_bias": "model-00003-of-00014.safetensors",
652
+ "model.layers.8.mlp.gate.proj.weight": "model-00003-of-00014.safetensors",
653
+ "model.layers.8.mlp.shared_expert.down_proj.weight": "model-00003-of-00014.safetensors",
654
+ "model.layers.8.mlp.shared_expert.gate_proj.weight": "model-00003-of-00014.safetensors",
655
+ "model.layers.8.mlp.shared_expert.up_proj.weight": "model-00003-of-00014.safetensors",
656
+ "model.layers.8.mlp.switch_mlp.down_proj.weight": "model-00003-of-00014.safetensors",
657
+ "model.layers.8.mlp.switch_mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
658
+ "model.layers.8.mlp.switch_mlp.up_proj.weight": "model-00003-of-00014.safetensors",
659
+ "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
660
+ "model.layers.8.self_attn.g_proj.weight": "model-00003-of-00014.safetensors",
661
+ "model.layers.8.self_attn.k_norm.weight": "model-00003-of-00014.safetensors",
662
+ "model.layers.8.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
663
+ "model.layers.8.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
664
+ "model.layers.8.self_attn.q_norm.weight": "model-00003-of-00014.safetensors",
665
+ "model.layers.8.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
666
+ "model.layers.8.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
667
+ "model.layers.9.input_layernorm.weight": "model-00004-of-00014.safetensors",
668
+ "model.layers.9.mlp.gate.e_score_correction_bias": "model-00003-of-00014.safetensors",
669
+ "model.layers.9.mlp.gate.proj.weight": "model-00003-of-00014.safetensors",
670
+ "model.layers.9.mlp.shared_expert.down_proj.weight": "model-00004-of-00014.safetensors",
671
+ "model.layers.9.mlp.shared_expert.gate_proj.weight": "model-00004-of-00014.safetensors",
672
+ "model.layers.9.mlp.shared_expert.up_proj.weight": "model-00004-of-00014.safetensors",
673
+ "model.layers.9.mlp.switch_mlp.down_proj.weight": "model-00004-of-00014.safetensors",
674
+ "model.layers.9.mlp.switch_mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
675
+ "model.layers.9.mlp.switch_mlp.up_proj.weight": "model-00003-of-00014.safetensors",
676
+ "model.layers.9.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
677
+ "model.layers.9.self_attn.g_proj.weight": "model-00003-of-00014.safetensors",
678
+ "model.layers.9.self_attn.k_norm.weight": "model-00003-of-00014.safetensors",
679
+ "model.layers.9.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
680
+ "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
681
+ "model.layers.9.self_attn.q_norm.weight": "model-00003-of-00014.safetensors",
682
+ "model.layers.9.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
683
+ "model.layers.9.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
684
+ "model.norm.weight": "model-00014-of-00014.safetensors"
685
+ }
686
+ }
modeling_laguna.py ADDED
@@ -0,0 +1,755 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from collections.abc import Callable
16
+ from typing import Optional
17
+
18
+ import torch
19
+ import torch.nn.functional as F
20
+ from torch import nn
21
+
22
+ from transformers import initialization as init
23
+ from transformers.activations import ACT2FN
24
+ from transformers.cache_utils import Cache, DynamicCache
25
+ from transformers.generation import GenerationMixin
26
+ from transformers.integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
27
+ from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
28
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
29
+ from transformers.modeling_layers import GradientCheckpointingLayer
30
+ from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
31
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
32
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
33
+ from transformers.processing_utils import Unpack
34
+ from transformers.utils import auto_docstring, can_return_tuple
35
+ from transformers.utils.generic import TransformersKwargs, maybe_autocast
36
+ from transformers.utils.output_capturing import OutputRecorder, capture_outputs
37
+ from .configuration_laguna import LagunaConfig
38
+
39
+
40
+ @use_kernel_forward_from_hub("RMSNorm")
41
+ class LagunaRMSNorm(nn.Module):
42
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
43
+ """
44
+ LagunaRMSNorm is equivalent to T5LayerNorm
45
+ """
46
+ super().__init__()
47
+ self.weight = nn.Parameter(torch.ones(hidden_size))
48
+ self.variance_epsilon = eps
49
+
50
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
51
+ input_dtype = hidden_states.dtype
52
+ hidden_states = hidden_states.to(torch.float32)
53
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
54
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
55
+ return self.weight * hidden_states.to(input_dtype)
56
+
57
+ def extra_repr(self):
58
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
59
+
60
+
61
+ class LagunaRotaryEmbedding(nn.Module):
62
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
63
+
64
+ def __init__(self, config: LagunaConfig, device=None, layer_type=None):
65
+ super().__init__()
66
+ self.max_seq_len_cached = config.max_position_embeddings
67
+ self.original_max_seq_len = config.max_position_embeddings
68
+
69
+ self.config = config
70
+
71
+ self.layer_types = list(set(config.layer_types))
72
+ self.rope_type = {}
73
+ for layer_type in self.layer_types:
74
+ rope_params = self.config.rope_parameters[layer_type]
75
+ if rope_params is None:
76
+ continue
77
+
78
+ self.rope_type[layer_type] = rope_params["rope_type"]
79
+ rope_init_fn: Callable = self.compute_default_rope_parameters
80
+ if self.rope_type[layer_type] != "default":
81
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]
82
+ curr_inv_freq, curr_attention_scaling = rope_init_fn(self.config, device, layer_type=layer_type)
83
+ self.register_buffer(f"{layer_type}_inv_freq", curr_inv_freq, persistent=False)
84
+ self.register_buffer(f"{layer_type}_original_inv_freq", curr_inv_freq.clone(), persistent=False)
85
+ setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
86
+
87
+ @staticmethod
88
+ def compute_default_rope_parameters(
89
+ config: LagunaConfig | None = None,
90
+ device: Optional["torch.device"] = None,
91
+ seq_len: int | None = None,
92
+ layer_type: str | None = None,
93
+ ) -> tuple["torch.Tensor", float]:
94
+ """
95
+ Computes the inverse frequencies according to the original RoPE implementation
96
+ Args:
97
+ config ([`~transformers.PreTrainedConfig`]):
98
+ The model configuration.
99
+ device (`torch.device`):
100
+ The device to use for initialization of the inverse frequencies.
101
+ seq_len (`int`, *optional*):
102
+ The current sequence length. Unused for this type of RoPE.
103
+ layer_type (`str`, *optional*):
104
+ The current layer type if the model has different RoPE parameters per type.
105
+ Should not be used unless `config.layer_types is not None`
106
+ Returns:
107
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
108
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
109
+ """
110
+ base = config.rope_parameters[layer_type]["rope_theta"]
111
+ # key difference to gemma3: partial rope
112
+ partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
113
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
114
+ dim = int(head_dim * partial_rotary_factor)
115
+
116
+ attention_factor = 1.0 # Unused in this type of RoPE
117
+
118
+ # Compute the inverse frequencies
119
+ inv_freq = 1.0 / (
120
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
121
+ )
122
+ return inv_freq, attention_factor
123
+
124
+ @torch.no_grad()
125
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
126
+ def forward(self, x, position_ids, layer_type=None):
127
+ inv_freq = getattr(self, f"{layer_type}_inv_freq")
128
+ attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
129
+
130
+ inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
131
+ position_ids_expanded = position_ids[:, None, :].float()
132
+
133
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
134
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
135
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
136
+ emb = torch.cat((freqs, freqs), dim=-1)
137
+ cos = emb.cos() * attention_scaling
138
+ sin = emb.sin() * attention_scaling
139
+
140
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
141
+
142
+
143
+ class LagunaMLP(nn.Module):
144
+ def __init__(self, config, intermediate_size=None):
145
+ super().__init__()
146
+ self.config = config
147
+ self.hidden_size = config.hidden_size
148
+ self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
149
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
150
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
151
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
152
+ self.act_fn = ACT2FN[config.hidden_act]
153
+
154
+ def forward(self, x):
155
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
156
+ return down_proj
157
+
158
+
159
+ class LagunaTopKRouter(nn.Module):
160
+ def __init__(self, config):
161
+ super().__init__()
162
+ self.top_k = config.num_experts_per_tok
163
+ self.num_experts = config.num_experts
164
+ self.hidden_dim = config.hidden_size
165
+ self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
166
+ self.e_score_correction_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
167
+ self.router_logit_softcapping = config.moe_router_logit_softcapping
168
+
169
+ def forward(
170
+ self,
171
+ hidden_states: torch.Tensor,
172
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
173
+ hidden_states = hidden_states.reshape(-1, self.hidden_dim)
174
+ router_logits = F.linear(hidden_states, self.weight).float()
175
+ # Optional logits softcapping
176
+ if self.router_logit_softcapping > 0.0:
177
+ router_logits = torch.tanh(router_logits / self.router_logit_softcapping) * self.router_logit_softcapping
178
+ # Sigmoid instead of softmax normalization
179
+ routing_scores = torch.sigmoid(router_logits)
180
+
181
+ scores_for_selection = routing_scores + self.e_score_correction_bias.to(routing_scores.dtype)
182
+ _, selected_experts = torch.topk(scores_for_selection, self.top_k, dim=-1)
183
+ routing_weights = routing_scores.gather(-1, selected_experts)
184
+ routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
185
+ routing_weights = routing_weights.to(hidden_states.dtype)
186
+
187
+ return router_logits, routing_weights, selected_experts
188
+
189
+
190
+ @use_experts_implementation
191
+ class LagunaExperts(nn.Module):
192
+ """Collection of expert weights stored as 3D tensors."""
193
+
194
+ def __init__(self, config):
195
+ super().__init__()
196
+ self.num_experts = config.num_experts
197
+ self.hidden_dim = config.hidden_size
198
+ self.intermediate_dim = config.moe_intermediate_size
199
+ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
200
+ self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
201
+ self.act_fn = ACT2FN[config.hidden_act]
202
+
203
+ def forward(
204
+ self,
205
+ hidden_states: torch.Tensor,
206
+ top_k_index: torch.Tensor,
207
+ top_k_weights: torch.Tensor,
208
+ ) -> torch.Tensor:
209
+ final_hidden_states = torch.zeros_like(hidden_states)
210
+ with torch.no_grad():
211
+ expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
212
+ expert_mask = expert_mask.permute(2, 1, 0)
213
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
214
+
215
+ for expert_idx in expert_hit:
216
+ expert_idx = expert_idx[0]
217
+ if expert_idx == self.num_experts:
218
+ continue
219
+ top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
220
+ current_state = hidden_states[token_idx]
221
+ gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
222
+ current_hidden_states = self.act_fn(gate) * up
223
+ current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
224
+ current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
225
+ final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
226
+
227
+ return final_hidden_states
228
+
229
+
230
+ class LagunaSparseMoeBlock(nn.Module):
231
+ def __init__(self, config: LagunaConfig):
232
+ super().__init__()
233
+ self.experts = LagunaExperts(config)
234
+ self.gate = LagunaTopKRouter(config)
235
+ self.shared_experts = LagunaMLP(config, intermediate_size=config.shared_expert_intermediate_size)
236
+ self.routed_scaling_factor = config.moe_routed_scaling_factor
237
+
238
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
239
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
240
+ hidden_states = hidden_states.view(-1, hidden_dim)
241
+ shared_output = self.shared_experts(hidden_states)
242
+
243
+ _, routing_weights, selected_experts = self.gate(hidden_states)
244
+ hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
245
+ # Additional scaling
246
+ hidden_states = hidden_states * self.routed_scaling_factor
247
+ hidden_states = hidden_states + shared_output
248
+
249
+ hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim)
250
+ return hidden_states
251
+
252
+
253
+ def rotate_half(x):
254
+ """Rotates half the hidden dims of the input."""
255
+ x1 = x[..., : x.shape[-1] // 2]
256
+ x2 = x[..., x.shape[-1] // 2 :]
257
+ return torch.cat((-x2, x1), dim=-1)
258
+
259
+
260
+ # Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
261
+ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
262
+ """Applies Rotary Position Embedding to the query and key tensors.
263
+
264
+ Removes the interleaving of cos and sin from GLM
265
+
266
+ Args:
267
+ q (`torch.Tensor`): The query tensor.
268
+ k (`torch.Tensor`): The key tensor.
269
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
270
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
271
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
272
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
273
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
274
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
275
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
276
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
277
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
278
+ Returns:
279
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
280
+ """
281
+ cos = cos.unsqueeze(unsqueeze_dim)
282
+ sin = sin.unsqueeze(unsqueeze_dim)
283
+
284
+ # Keep half or full tensor for later concatenation
285
+ rotary_dim = cos.shape[-1]
286
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
287
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
288
+
289
+ # Apply rotary embeddings on the first half or full tensor
290
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
291
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
292
+
293
+ # Concatenate back to full shape
294
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
295
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
296
+ return q_embed, k_embed
297
+
298
+
299
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
300
+ """
301
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
302
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
303
+ """
304
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
305
+ if n_rep == 1:
306
+ return hidden_states
307
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
308
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
309
+
310
+
311
+ def eager_attention_forward(
312
+ module: nn.Module,
313
+ query: torch.Tensor,
314
+ key: torch.Tensor,
315
+ value: torch.Tensor,
316
+ attention_mask: torch.Tensor | None,
317
+ scaling: float,
318
+ dropout: float = 0.0,
319
+ **kwargs: Unpack[TransformersKwargs],
320
+ ):
321
+ key_states = repeat_kv(key, module.num_key_value_groups)
322
+ value_states = repeat_kv(value, module.num_key_value_groups)
323
+
324
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
325
+ if attention_mask is not None:
326
+ attn_weights = attn_weights + attention_mask
327
+
328
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
329
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
330
+ attn_output = torch.matmul(attn_weights, value_states)
331
+ attn_output = attn_output.transpose(1, 2).contiguous()
332
+
333
+ return attn_output, attn_weights
334
+
335
+
336
+ @use_kernelized_func(apply_rotary_pos_emb)
337
+ class LagunaAttention(nn.Module):
338
+ """Afmoe-style SWA/GQA attention with Laguna-specific gating and per-layer head count."""
339
+
340
+ def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int):
341
+ super().__init__()
342
+ # Number of heads is controlled via `config.num_attention_heads_per_layer` which is passed from the parent for the specific layer
343
+ self.num_heads = num_heads
344
+ self.config = config
345
+ self.layer_idx = layer_idx
346
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
347
+ self.num_key_value_groups = self.num_heads // config.num_key_value_heads
348
+ self.scaling = self.head_dim**-0.5
349
+ self.attention_dropout = config.attention_dropout
350
+ self.is_causal = True
351
+
352
+ # Per-layer head count: rebuild q_proj and o_proj using self.num_heads (parent uses config.num_attention_heads).
353
+ self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
354
+ self.k_proj = nn.Linear(
355
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
356
+ )
357
+ self.v_proj = nn.Linear(
358
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
359
+ )
360
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
361
+ # Parent LlamaAttention already sets: layer_idx, num_heads, num_key_value_heads, num_key_value_groups, head_dim
362
+ # We only add Laguna-specific attributes
363
+ self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
364
+ self.sliding_window = config.sliding_window if self.is_local_attention else None
365
+
366
+ self.q_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
367
+ self.k_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
368
+ self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
369
+
370
+ def forward(
371
+ self,
372
+ hidden_states: torch.Tensor,
373
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
374
+ attention_mask: torch.Tensor | None,
375
+ past_key_values: Cache | None = None,
376
+ **kwargs: Unpack[FlashAttentionKwargs],
377
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
378
+ input_shape = hidden_states.shape[:-1]
379
+ hidden_shape = (*input_shape, -1, self.head_dim)
380
+
381
+ query_states = self.q_proj(hidden_states).view(hidden_shape)
382
+ key_states = self.k_proj(hidden_states).view(hidden_shape)
383
+ value_states = self.v_proj(hidden_states).view(hidden_shape)
384
+
385
+ query_states = self.q_norm(query_states).transpose(1, 2)
386
+ key_states = self.k_norm(key_states).transpose(1, 2)
387
+ value_states = value_states.transpose(1, 2)
388
+
389
+ cos, sin = position_embeddings
390
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
391
+
392
+ if past_key_values is not None:
393
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
394
+
395
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
396
+ self.config._attn_implementation, eager_attention_forward
397
+ )
398
+ attn_output, attn_weights = attention_interface(
399
+ self,
400
+ query_states,
401
+ key_states,
402
+ value_states,
403
+ attention_mask,
404
+ dropout=0.0 if not self.training else self.attention_dropout,
405
+ scaling=self.scaling,
406
+ sliding_window=self.sliding_window,
407
+ **kwargs,
408
+ )
409
+
410
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
411
+
412
+ gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
413
+ attn_output = (attn_output.view(*input_shape, -1, self.head_dim) * gate.unsqueeze(-1)).view(*input_shape, -1)
414
+
415
+ attn_output = self.o_proj(attn_output)
416
+ return attn_output, attn_weights
417
+
418
+
419
+ class LagunaDecoderLayer(GradientCheckpointingLayer):
420
+ def __init__(self, config: LagunaConfig, layer_idx: int):
421
+ super().__init__()
422
+ self.hidden_size = config.hidden_size
423
+ self.self_attn = LagunaAttention(config, layer_idx, config.num_attention_heads_per_layer[layer_idx])
424
+ if config.mlp_layer_types[layer_idx] == "sparse":
425
+ self.mlp = LagunaSparseMoeBlock(config)
426
+ else:
427
+ self.mlp = LagunaMLP(config, intermediate_size=config.intermediate_size)
428
+ self.input_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
429
+ self.post_attention_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
430
+
431
+ def forward(
432
+ self,
433
+ hidden_states: torch.Tensor,
434
+ attention_mask: torch.Tensor | None = None,
435
+ position_ids: torch.LongTensor | None = None,
436
+ past_key_values: Cache | None = None,
437
+ use_cache: bool | None = False,
438
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
439
+ **kwargs: Unpack[TransformersKwargs],
440
+ ) -> torch.Tensor:
441
+ residual = hidden_states
442
+ hidden_states = self.input_layernorm(hidden_states)
443
+ # Self Attention
444
+ hidden_states, _ = self.self_attn(
445
+ hidden_states=hidden_states,
446
+ attention_mask=attention_mask,
447
+ position_ids=position_ids,
448
+ past_key_values=past_key_values,
449
+ use_cache=use_cache,
450
+ position_embeddings=position_embeddings,
451
+ **kwargs,
452
+ )
453
+ hidden_states = residual + hidden_states
454
+
455
+ # Fully Connected
456
+ residual = hidden_states
457
+ hidden_states = self.post_attention_layernorm(hidden_states)
458
+ hidden_states = self.mlp(hidden_states)
459
+ hidden_states = residual + hidden_states
460
+ return hidden_states
461
+
462
+
463
+ @auto_docstring
464
+ class LagunaPreTrainedModel(PreTrainedModel):
465
+ config: LagunaConfig
466
+ base_model_prefix = "model"
467
+ supports_gradient_checkpointing = True
468
+ _no_split_modules = ["LagunaDecoderLayer"]
469
+ _skip_keys_device_placement = ["past_key_values"]
470
+ _supports_flash_attn = True
471
+ _supports_sdpa = True
472
+ _supports_flex_attn = True
473
+
474
+ _can_compile_fullgraph = True
475
+ _supports_attention_backend = True
476
+ _can_record_outputs = {
477
+ "router_logits": OutputRecorder(LagunaTopKRouter, index=0),
478
+ "hidden_states": LagunaDecoderLayer,
479
+ "attentions": LagunaAttention,
480
+ }
481
+
482
+ @torch.no_grad()
483
+ def _init_weights(self, module):
484
+ super()._init_weights(module)
485
+ std = self.config.initializer_range
486
+ if isinstance(module, LagunaExperts):
487
+ init.normal_(module.gate_up_proj, mean=0.0, std=std)
488
+ init.normal_(module.down_proj, mean=0.0, std=std)
489
+ elif isinstance(module, LagunaTopKRouter):
490
+ init.normal_(module.weight, mean=0.0, std=std)
491
+ if isinstance(module, LagunaTopKRouter):
492
+ torch.nn.init.zeros_(module.e_score_correction_bias)
493
+ elif isinstance(module, LagunaRotaryEmbedding):
494
+ for layer_type in module.layer_types:
495
+ rope_init_fn = module.compute_default_rope_parameters
496
+ if module.rope_type[layer_type] != "default":
497
+ rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]
498
+ curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
499
+ init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
500
+ init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
501
+
502
+
503
+ @auto_docstring
504
+ class LagunaModel(LagunaPreTrainedModel):
505
+ def __init__(self, config: LagunaConfig):
506
+ super().__init__(config)
507
+ self.padding_idx = config.pad_token_id
508
+ self.vocab_size = config.vocab_size
509
+
510
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
511
+ self.layers = nn.ModuleList(
512
+ [LagunaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
513
+ )
514
+ self.norm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
515
+ self.rotary_emb = LagunaRotaryEmbedding(config=config)
516
+ self.gradient_checkpointing = False
517
+
518
+ # Initialize weights and apply final processing
519
+ self.post_init()
520
+
521
+ @capture_outputs
522
+ @auto_docstring
523
+ def forward(
524
+ self,
525
+ input_ids: torch.LongTensor | None = None,
526
+ attention_mask: torch.Tensor | None = None,
527
+ position_ids: torch.LongTensor | None = None,
528
+ past_key_values: Cache | None = None,
529
+ inputs_embeds: torch.FloatTensor | None = None,
530
+ use_cache: bool | None = None,
531
+ **kwargs: Unpack[TransformersKwargs],
532
+ ) -> MoeModelOutputWithPast:
533
+ if (input_ids is None) ^ (inputs_embeds is not None):
534
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
535
+
536
+ if inputs_embeds is None:
537
+ inputs_embeds = self.embed_tokens(input_ids)
538
+
539
+ if use_cache and past_key_values is None:
540
+ past_key_values = DynamicCache(config=self.config)
541
+
542
+ if position_ids is None:
543
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
544
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
545
+ position_ids = position_ids.unsqueeze(0)
546
+
547
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
548
+ mask_kwargs = {
549
+ "config": self.config,
550
+ "inputs_embeds": inputs_embeds,
551
+ "attention_mask": attention_mask,
552
+ "past_key_values": past_key_values,
553
+ "position_ids": position_ids,
554
+ }
555
+ mask_creation_functions = {
556
+ "full_attention": lambda: create_causal_mask(**mask_kwargs),
557
+ "sliding_attention": lambda: create_sliding_window_causal_mask(**mask_kwargs),
558
+ }
559
+ causal_mask_mapping = {}
560
+ for layer_type in set(self.config.layer_types):
561
+ causal_mask_mapping[layer_type] = mask_creation_functions[layer_type]()
562
+
563
+ hidden_states = inputs_embeds
564
+ position_embeddings = {}
565
+ for layer_type in set(self.config.layer_types):
566
+ position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)
567
+
568
+ for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
569
+ hidden_states = decoder_layer(
570
+ hidden_states,
571
+ attention_mask=causal_mask_mapping[self.config.layer_types[i]],
572
+ position_embeddings=position_embeddings[self.config.layer_types[i]],
573
+ position_ids=position_ids,
574
+ past_key_values=past_key_values,
575
+ **kwargs,
576
+ )
577
+
578
+ hidden_states = self.norm(hidden_states)
579
+
580
+ return MoeModelOutputWithPast(
581
+ last_hidden_state=hidden_states,
582
+ past_key_values=past_key_values if use_cache else None,
583
+ )
584
+
585
+
586
+ def load_balancing_loss_func(
587
+ gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
588
+ num_experts: int | None = None,
589
+ top_k=2,
590
+ attention_mask: torch.Tensor | None = None,
591
+ ) -> torch.Tensor | int:
592
+ r"""
593
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
594
+
595
+ See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
596
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
597
+ experts is too unbalanced.
598
+
599
+ Args:
600
+ gate_logits:
601
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
602
+ shape [batch_size X sequence_length, num_experts].
603
+ num_experts:
604
+ Number of experts
605
+ top_k:
606
+ The number of experts to route per-token, can be also interpreted as the `top-k` routing
607
+ parameter.
608
+ attention_mask (`torch.Tensor`, *optional*):
609
+ The attention_mask used in forward function
610
+ shape [batch_size X sequence_length] if not None.
611
+
612
+ Returns:
613
+ The auxiliary loss.
614
+ """
615
+ if gate_logits is None or not isinstance(gate_logits, tuple):
616
+ return 0
617
+
618
+ if isinstance(gate_logits, tuple):
619
+ compute_device = gate_logits[0].device
620
+ concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
621
+
622
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
623
+
624
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
625
+
626
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
627
+
628
+ if attention_mask is None:
629
+ # Compute the percentage of tokens routed to each experts
630
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
631
+
632
+ # Compute the average probability of routing to these experts
633
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
634
+ else:
635
+ batch_size, sequence_length = attention_mask.shape
636
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
637
+
638
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
639
+ expert_attention_mask = (
640
+ attention_mask[None, :, :, None, None]
641
+ .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
642
+ .reshape(-1, top_k, num_experts)
643
+ .to(compute_device)
644
+ )
645
+
646
+ # Compute the percentage of tokens routed to each experts
647
+ tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
648
+ expert_attention_mask, dim=0
649
+ )
650
+
651
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
652
+ router_per_expert_attention_mask = (
653
+ attention_mask[None, :, :, None]
654
+ .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
655
+ .reshape(-1, num_experts)
656
+ .to(compute_device)
657
+ )
658
+
659
+ # Compute the average probability of routing to these experts
660
+ router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
661
+ router_per_expert_attention_mask, dim=0
662
+ )
663
+
664
+ overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
665
+ return overall_loss * num_experts
666
+
667
+
668
+ @auto_docstring
669
+ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
670
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
671
+ _tp_plan = {"lm_head": "colwise_gather_output"}
672
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
673
+
674
+ def __init__(self, config):
675
+ super().__init__(config)
676
+ self.model = LagunaModel(config)
677
+ self.vocab_size = config.vocab_size
678
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
679
+ self.router_aux_loss_coef = config.router_aux_loss_coef
680
+ self.num_experts = config.num_experts
681
+ self.num_experts_per_tok = config.num_experts_per_tok
682
+
683
+ # Initialize weights and apply final processing
684
+ self.post_init()
685
+
686
+ @can_return_tuple
687
+ @auto_docstring
688
+ def forward(
689
+ self,
690
+ input_ids: torch.LongTensor | None = None,
691
+ attention_mask: torch.Tensor | None = None,
692
+ position_ids: torch.LongTensor | None = None,
693
+ past_key_values: Cache | None = None,
694
+ inputs_embeds: torch.FloatTensor | None = None,
695
+ labels: torch.LongTensor | None = None,
696
+ use_cache: bool | None = None,
697
+ output_router_logits: bool | None = None,
698
+ logits_to_keep: int | torch.Tensor = 0,
699
+ **kwargs: Unpack[TransformersKwargs],
700
+ ) -> MoeCausalLMOutputWithPast:
701
+ r"""
702
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
703
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
704
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
705
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
706
+ """
707
+
708
+ output_router_logits = (
709
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
710
+ )
711
+
712
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
713
+ outputs: MoeModelOutputWithPast = self.model(
714
+ input_ids=input_ids,
715
+ attention_mask=attention_mask,
716
+ position_ids=position_ids,
717
+ past_key_values=past_key_values,
718
+ inputs_embeds=inputs_embeds,
719
+ use_cache=use_cache,
720
+ output_router_logits=output_router_logits,
721
+ **kwargs,
722
+ )
723
+
724
+ hidden_states = outputs.last_hidden_state
725
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
726
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
727
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
728
+
729
+ loss = None
730
+ if labels is not None:
731
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
732
+
733
+ aux_loss = None
734
+ if output_router_logits:
735
+ aux_loss = load_balancing_loss_func(
736
+ outputs.router_logits,
737
+ self.num_experts,
738
+ self.num_experts_per_tok,
739
+ attention_mask,
740
+ )
741
+ if labels is not None:
742
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
743
+
744
+ return MoeCausalLMOutputWithPast(
745
+ loss=loss,
746
+ aux_loss=aux_loss,
747
+ logits=logits,
748
+ past_key_values=outputs.past_key_values,
749
+ hidden_states=outputs.hidden_states,
750
+ attentions=outputs.attentions,
751
+ router_logits=outputs.router_logits,
752
+ )
753
+
754
+
755
+ __all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaPreTrainedModel"]
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "〈|EOS|〉",
4
+ "clean_up_tokenization_spaces": false,
5
+ "cls_token": "〈|CLS|〉",
6
+ "eos_token": "〈|EOS|〉",
7
+ "is_local": true,
8
+ "mask_token": "〈|MASK|〉",
9
+ "model_max_length": 1000000000000000019884624838656,
10
+ "pad_token": "〈|PAD|〉",
11
+ "sep_token": "〈|SEP|〉",
12
+ "tokenizer_class": "TokenizersBackend",
13
+ "tool_parser_type": "laguna",
14
+ "unk_token": "〈|UNK|〉"
15
+ }