Upload optimized Gemma4 checkpoint

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,19 @@
+# gemma4-zero-compute
+
+Optimized Gemma4 text-only checkpoint exported from this repository.
+
+## Loading
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("haysonC/gemma4-zero-compute", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("haysonC/gemma4-zero-compute", trust_remote_code=True)
+```
+
+## Notes
+
+- Base model: `google/gemma-4-26B-A4B-it`
+- Architecture: `OptimizedGemma4ForCausalLM`
+- Scope: text-only causal language model export
+- Router checkpoint loaded: `True`

__init__.py

@@ -0,0 +1,2 @@
+from .configuration_gemma4 import Gemma4TextConfig
+from .gemma4_optimization import OptimizedGemma4ForCausalLM

chat_template.jinja

@@ -0,0 +1,347 @@
+{%- macro format_parameters(properties, required) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([])) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {{- messages[0]['content'] | trim -}}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+    {%- set prev_nt = namespace(role=None, found=false) -%}
+    {%- if loop.index0 > 0 -%}
+        {%- for j in range(loop.index0 - 1, -1, -1) -%}
+            {%- if not prev_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                    {%- set prev_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message['tool_calls'] -%}
+                {%- for tool_call in message['tool_calls'] -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is string -%}
+                        {{- function['arguments'] -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message['tool_responses'] -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+                        {%- for tc in message['tool_calls'] -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- if message['content'] is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message['content'] is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item['type'] == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item['type'] == 'image' -%}
+                        {{- '<|image|>' -}}
+                        {%- set ns.prev_message_type = 'image' -%}
+                    {%- elif item['type'] == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                        {%- set ns.prev_message_type = 'audio' -%}
+                    {%- elif item['type'] == 'video' -%}
+                        {{- '<|video|>' -}}
+                        {%- set ns.prev_message_type = 'video' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif not (ns_tr_out.flag and not message.get('content')) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,89 @@
+{
+  "add_zero_compute_expert": true,
+  "architectures": [
+    "OptimizedGemma4ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_k_eq_v": true,
+  "auto_map": {
+    "AutoConfig": "configuration_gemma4.Gemma4TextConfig",
+    "AutoModelForCausalLM": "gemma4_optimization.OptimizedGemma4ForCausalLM"
+  },
+  "bos_token_id": 2,
+  "dtype": "bfloat16",
+  "enable_moe_block": true,
+  "eos_token_id": 1,
+  "final_logit_softcapping": 30.0,
+  "global_head_dim": 512,
+  "head_dim": 256,
+  "hidden_activation": "gelu_pytorch_tanh",
+  "hidden_size": 2816,
+  "hidden_size_per_layer_input": 0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2112,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "model_type": "gemma4_text",
+  "moe_intermediate_size": 704,
+  "num_attention_heads": 16,
+  "num_experts": 128,
+  "num_global_key_value_heads": 2,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 8,
+  "num_kv_shared_layers": 0,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "full_attention": {
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 1000000.0,
+      "rope_type": "proportional"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sliding_window": 1024,
+  "tie_word_embeddings": true,
+  "top_k_experts": 8,
+  "transformers_version": "5.5.3",
+  "use_bidirectional_attention": "vision",
+  "use_cache": true,
+  "use_double_wide_mlp": false,
+  "use_zero_compute_optimization": true,
+  "vocab_size": 262144,
+  "vocab_size_per_layer_input": 262144
+}

configuration_gemma4.py

@@ -0,0 +1,383 @@
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Literal
+
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import auto_docstring, logging
+from transformers.utils.type_validators import interval
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="google/gemma-4-e2b-it")
+@strict
+class Gemma4AudioConfig(PreTrainedConfig):
+    r"""
+    subsampling_conv_channels (`list[int]`, defaults to `[128, 32]`):
+        Channel sizes for the convolutional layers in the Sub-sample Convolution Projection.
+    residual_weight (`float`, defaults to `0.5`):
+        Scaling applied to hidden_states prior to combining with the residual in the feedforward.
+    attention_chunk_size (`int`, defaults to `12`):
+        The sub-sequence size for attention processing.
+    attention_context_left (`int`, defaults to `13`):
+        The leftward context size for the attention chunk.
+    attention_context_right (`int`, defaults to `0`):
+        The rightward context size for the attention chunk.
+    attention_logit_cap (`float`, defaults to `50.0`):
+        Cap applied to attention weights.
+    attention_invalid_logits_value (`float`, defaults to `1e-9`):
+        Value to use for invalid logits in attention.
+    use_clipped_linears (`bool`, defaults to `True`):
+        If true, apply clipping to the Linear layers, drawing bounds from the model checkpoint.
+    gradient_clipping (`float`, defaults to `1e10`):
+        Clipping value used to stabilize extremely large gradient values.
+    output_proj_dims (`int`, defaults to `1536`):
+        Dimension of the final linear projection from `hidden_size` to the model's output.
+    """
+
+    model_type = "gemma4_audio"
+
+    hidden_size: int = 1024
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 8
+    hidden_act: str = "silu"
+
+    # subsampling parameters
+    subsampling_conv_channels: list[int] | tuple[int, int] = (128, 32)
+
+    # conformer parameters
+    conv_kernel_size: int = 5
+    residual_weight: float = 0.5
+    attention_chunk_size: int = 12
+    attention_context_left: int = 13
+    attention_context_right: int = 0
+    attention_logit_cap: float = 50.0
+    attention_invalid_logits_value: float = -1.0e9
+
+    use_clipped_linears: bool = True
+    rms_norm_eps: float = 1e-6
+    gradient_clipping: float = 1e10
+    output_proj_dims: int = 1536
+    initializer_range: float = interval(min=0.0, max=1.0)(default=0.02)
+
+    def __post_init__(self, **kwargs):
+        # JSON serialization converts tuples to lists, convert back
+        if isinstance(self.subsampling_conv_channels, tuple):
+            self.subsampling_conv_channels = list(self.subsampling_conv_channels)
+        super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="google/gemma-4-e2b-it")
+@strict
+class Gemma4TextConfig(PreTrainedConfig):
+    r"""
+    use_bidirectional_attention (`str`, *optional*):
+        Controls bidirectional attention behavior. When set to `"vision"`, vision tokens
+        attend bidirectionally while text tokens use causal attention. When set to `"all"`,
+        all tokens use bidirectional attention.
+    vocab_size_per_layer_input (`int`, defaults to 262144):
+        Vocabulary size for the per-layer input embeddings. Used by models with per-layer
+        residual streams where a smaller embedding is added at each decoder layer.
+    hidden_size_per_layer_input (`int`, defaults to 256):
+        Hidden dimension for the per-layer input embeddings. Controls the width of the
+        per-layer residual embedding vectors.
+    num_global_key_value_heads (`int`, *optional*):
+        Number of key-value heads for global (full) attention layers. If `None`, defaults
+        to `num_key_value_heads`.
+    global_head_dim (`int`, defaults to 512):
+        Dimension of each attention head in global (full) attention layers.
+    attention_k_eq_v (`bool`, defaults to `False`):
+        Whether keys and values share the same projection weights. When `True`, the key
+        projection output is reused as the value projection.
+    num_kv_shared_layers (`int`, defaults to 0):
+        Number of consecutive decoder layers that share the same key-value projections.
+        A value of 0 means no sharing (each layer has independent KV projections).
+    enable_moe_block (`bool`, defaults to `False`):
+        Whether to enable Mixture-of-Experts (MoE) blocks in the decoder layers. When
+        `True`, eligible layers will use a sparse MoE feed-forward network.
+    use_double_wide_mlp (`bool`, defaults to `False`):
+        Whether to use a double-width MLP with fused gate and up projections.
+    top_k_experts (`int`, *optional*):
+        Number of experts activated per token in MoE layers. Only used when
+        `enable_moe_block=True`.
+    moe_intermediate_size (`int`, *optional*):
+        Intermediate (hidden) size of each expert's feed-forward network in MoE layers.
+        Only used when `enable_moe_block=True`.
+    add_zero_compute_expert (`bool`, defaults to `False`):
+        Whether to append a router-only expert slot that performs no expert compute. This
+        keeps the original expert weights intact while allowing the router to learn to
+        send tokens to a zero-compute path.
+    use_zero_compute_optimization (`bool`, defaults to `False`):
+        Signals higher-level orchestration to build the optimized Gemma4 text stack instead
+        of the original one while keeping the base architecture definitions available.
+    """
+
+    model_type = "gemma4_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+        "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+        "layers.*.experts.gate_up_proj": "packed_colwise",
+        "layers.*.experts.down_proj": "rowwise",
+        "layers.*.experts": "moe_tp_experts",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    vocab_size: int = 262_144
+    hidden_size: int = 2304
+    intermediate_size: int = 9216
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    hidden_activation: str = "gelu_pytorch_tanh"
+    max_position_embeddings: int = 131_072
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int | None = 0
+    eos_token_id: int | list[int] | None = 1
+    bos_token_id: int | None = 2
+    tie_word_embeddings: bool = True
+    rope_parameters: dict | None = None
+    attention_bias: bool = False
+    attention_dropout: int | float | None = 0.0
+    sliding_window: int = 512
+    layer_types: list[str] | None = None
+    final_logit_softcapping: float | None = None
+    use_bidirectional_attention: Literal["all", "vision"] | None = None
+    vocab_size_per_layer_input: int = 262_144
+    hidden_size_per_layer_input: int = 256
+    num_global_key_value_heads: int | None = None
+    global_head_dim: int = 512
+    attention_k_eq_v: bool = False
+    num_kv_shared_layers: int = 0
+    enable_moe_block: bool = False
+    use_double_wide_mlp: bool = False
+    num_experts: int | None = None
+    top_k_experts: int | None = None
+    moe_intermediate_size: int | None = None
+    add_zero_compute_expert: bool = False
+    use_zero_compute_optimization: bool = False
+
+    def __post_init__(self, **kwargs):
+        if self.use_bidirectional_attention == "all":
+            self.sliding_window = (self.sliding_window // 2) + 1  # due to fa we set exclusive bounds
+
+        if self.layer_types is None:
+            sliding_window_pattern = 6  # by default 5:1
+            self.layer_types = [
+                "sliding_attention" if bool((i + 1) % sliding_window_pattern) else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+
+        if self.layer_types and (last_layer_type := self.layer_types[-1]) != "full_attention":
+            logger.warning(
+                f"Last layer must use `full_attention`, but got `{last_layer_type}`. Forcing last layer to `full_attention`."
+            )
+            self.layer_types[-1] = "full_attention"
+
+        default_rope_params: dict[Literal["full_attention", "sliding_attention"] : dict[str, Any]] = {
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
+            "full_attention": {"rope_type": "proportional", "partial_rotary_factor": 0.25, "rope_theta": 1_000_000.0},
+        }
+        active_layer_types = set(self.layer_types)
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                layer_type: dict(default_rope_params[layer_type]) for layer_type in active_layer_types
+            }
+        elif set(self.rope_parameters.keys()).issubset(default_rope_params):
+            self.rope_parameters = {
+                layer_type: dict(rope_params)
+                for layer_type, rope_params in self.rope_parameters.items()
+                if layer_type in active_layer_types
+            }
+
+        if self.num_experts is not None and self.top_k_experts is not None:
+            total_num_experts = self.num_experts + int(self.add_zero_compute_expert)
+            if self.top_k_experts > total_num_experts:
+                logger.warning(
+                    "top_k_experts=%s exceeds the available expert count %s. "
+                    "Clamping top_k_experts to %s.",
+                    self.top_k_experts,
+                    total_num_experts,
+                    total_num_experts,
+                )
+                self.top_k_experts = total_num_experts
+
+        if self.add_zero_compute_expert:
+            self.use_zero_compute_optimization = True
+
+        super().__post_init__(**kwargs)
+
+    def convert_rope_params_to_dict(self, **kwargs):
+        # No need to handle BC for new models, because they have no old-format `rope_scaling`
+        return kwargs
+
+
+@auto_docstring(checkpoint="google/gemma-4-e2b-it")
+@strict
+class Gemma4VisionConfig(PreTrainedConfig):
+    r"""
+    pooling_kernel_size (`int`, *optional*):
+        Spatial pooling kernel size applied after patchification.
+    position_embedding_size (`int`, defaults to 10240):
+        Maximum number of position embeddings for the vision encoder. Controls the size of
+        the learned 2D position embedding table used by the patch embedder.
+    use_clipped_linears (`bool`, defaults to `False`):
+        Whether to use weight-clipped linear layers. When enabled, linear layer weights are
+        clamped to a fixed range during the forward pass to improve numerical stability.
+    standardize (`bool`, defaults to `False`):
+        If true, applies a bias and scale to the soft tokens returned from the pooler.
+    """
+
+    model_type = "gemma4_vision"
+    base_model_tp_plan = {
+        "encoder.layers.*.self_attn.q_proj": "colwise",
+        "encoder.layers.*.self_attn.k_proj": "colwise",
+        "encoder.layers.*.self_attn.v_proj": "colwise",
+        "encoder.layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+        "encoder.layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+        "encoder.layers.*.self_attn.o_proj": "rowwise",
+        "encoder.layers.*.mlp.gate_proj": "colwise",
+        "encoder.layers.*.mlp.up_proj": "colwise",
+        "encoder.layers.*.mlp.down_proj": "rowwise",
+    }
+    default_theta = 100.0
+
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    num_hidden_layers: int = 16
+    num_attention_heads: int = 12
+    num_key_value_heads: int = 12
+    head_dim: int = 64
+    hidden_activation: str = "gelu_pytorch_tanh"
+    rms_norm_eps: float = 1e-6
+    max_position_embeddings: int = 131_072
+    attention_bias: bool | None = False
+    attention_dropout: float | None = 0.0
+    rope_parameters: dict | None = None
+    pooling_kernel_size: int = 3
+    patch_size: int = 16
+    position_embedding_size: int = 10 * 1024
+    use_clipped_linears: bool = False
+    standardize: bool = False
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        if self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "rope_theta": 100.0}
+
+        super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="google/gemma-4-e2b-it")
+@strict
+class Gemma4Config(PreTrainedConfig):
+    r"""
+    boi_token_id (`int`, *optional*, defaults to 255999):
+        The begin-of-image token index to wrap the image prompt.
+    eoi_token_id (`int`, *optional*, defaults to 258882):
+        The end-of-image token index to wrap the image prompt.
+    boa_token_id (`int`, *optional*, defaults to 256000):
+        The begin-of-audio token index to wrap the audio prompt.
+    eoa_token_index (`int`, *optional*, defaults to 258883):
+        The end-of-audio token index to wrap the audio prompt.
+
+    Example:
+
+    ```python
+    >>> from transformers import (
+    >>>     Gemma4AudioConfig,
+    >>>     Gemma4Config,
+    >>>     Gemma4ForConditionalGeneration,
+    >>>     Gemma4TextConfig,
+    >>>     Gemma4VisionConfig,
+    >>> )
+
+    >>> # Initializing a Gemma 4 Audio config.
+    >>> audio_config = Gemma4AudioConfig()
+
+    >>> # Initializing a Gemma 4 Text config.
+    >>> text_config = Gemma4TextConfig()
+
+    >>> # Initializing a Gemma 4 vision config.
+    >>> vision_config = Gemma4VisionConfig()
+
+    >>> # Initializing a Gemma 4 config similar to google/gemma-4-e2b-it
+    >>> configuration = Gemma4Config(text_config, vision_config, audio_config)
+
+    >>> # Initializing a model from the google/gemma-4-e2b-it configuration
+    >>> model = Gemma4ForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "gemma4"
+    sub_configs = {
+        "text_config": Gemma4TextConfig,
+        "vision_config": Gemma4VisionConfig,
+        "audio_config": Gemma4AudioConfig,
+    }
+
+    text_config: Gemma4TextConfig | dict[str, Any] | None = None
+    vision_config: Gemma4VisionConfig | dict[str, Any] | None = None
+    audio_config: Gemma4AudioConfig | dict[str, Any] | None = None
+    boi_token_id: int | None = 255_999
+    eoi_token_id: int | None = 258_882
+    image_token_id: int | None = 258_880
+    video_token_id: int | None = 258_884
+    boa_token_id: int | None = 256_000
+    eoa_token_index: int | None = 258_883
+    audio_token_id: int | None = 258_881
+    initializer_range: float | None = 0.02
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self, **kwargs):
+        if self.text_config is None:
+            self.text_config = Gemma4TextConfig()
+            logger.info("text_config is None. Using default Gemma4TextConfig.")
+        elif isinstance(self.text_config, dict):
+            self.text_config = Gemma4TextConfig(**self.text_config)
+
+        if self.vision_config is None:
+            logger.info("vision_config is None. Gemma4Model.vision_tower will not be initialized.")
+        if isinstance(self.vision_config, dict):
+            self.vision_config = Gemma4VisionConfig(**self.vision_config)
+
+        if self.audio_config is None:
+            logger.info("audio_config is None. Gemma4Model.audio_tower will not be initialized.")
+        if isinstance(self.audio_config, dict):
+            self.audio_config = Gemma4AudioConfig(**self.audio_config)
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["Gemma4AudioConfig", "Gemma4Config", "Gemma4TextConfig", "Gemma4VisionConfig"]

export_summary.json

@@ -0,0 +1,348 @@
+{
+  "base_model_class": "Gemma4TextCausalLMProxy",
+  "code_files": [
+    "/tmp/gemma4-hf-export-zxa_ushe/configuration_gemma4.py",
+    "/tmp/gemma4-hf-export-zxa_ushe/modeling_gemma4.py",
+    "/tmp/gemma4-hf-export-zxa_ushe/gemma4_optimization.py",
+    "/tmp/gemma4-hf-export-zxa_ushe/__init__.py"
+  ],
+  "max_shard_size": "5GB",
+  "model_class": "OptimizedGemma4ForCausalLM",
+  "output_dir": "/tmp/gemma4-hf-export-zxa_ushe",
+  "repo_id": "haysonC/gemma4-zero-compute",
+  "router_checkpoint": {
+    "config": {
+      "metadata": {
+        "resume_summary": {
+          "config": {
+            "metadata": {
+              "resume_summary": {
+                "config": {
+                  "metadata": {
+                    "resume_summary": {
+                      "loaded": false,
+                      "reason": "resume disabled"
+                    },
+                    "step": 500,
+                    "training_metrics": {
+                      "current_lambda_zero_compute": 2.0,
+                      "effective_batch_size": 16,
+                      "entropy_loss": 4.538632531960806,
+                      "entropy_term": 0.0,
+                      "expert_usage_sample": [
+                        1240.0,
+                        2139.0,
+                        1213.0,
+                        793.0,
+                        1992.0,
+                        1957.0,
+                        503.0,
+                        1528.0,
+                        1398.0,
+                        2252.0,
+                        1835.0,
+                        1146.0,
+                        1440.0,
+                        1527.0,
+                        1757.0,
+                        803.0
+                      ],
+                      "grad_norm": 70.5,
+                      "gradient_accumulation_steps": 8,
+                      "lambda_entropy": 0.0,
+                      "lambda_router": 1.0,
+                      "lambda_zero_compute": 2.0,
+                      "loss": 2.327640622854233,
+                      "micro_batch_size": 2,
+                      "output_kl": 1.1904816403985023,
+                      "output_kl_term": 1.1904816403985023,
+                      "probe_output_kl": 1.1182771921157837,
+                      "probe_router_entropy": 4.540449047088623,
+                      "probe_router_kl": 0.051853783428668976,
+                      "probe_same_expert_ratio": 0.8435872395833334,
+                      "probe_zero_compute_loss": 0.5133160352706909,
+                      "probe_zero_compute_margin_gap": 0.41453187317432216,
+                      "probe_zero_compute_mass": 0.013335910812020301,
+                      "probe_zero_compute_top1_ratio": 0.00654296875,
+                      "probe_zero_compute_topk_ratio": 0.29026692708333335,
+                      "router_entropy": 4.538632531960806,
+                      "router_kl": 0.05338437343016267,
+                      "router_kl_term": 0.05338437343016267,
+                      "same_expert_ratio": 0.8424682617187501,
+                      "step": 500,
+                      "tokens_per_optimizer_step": 8192,
+                      "zero_compute_loss": 0.5418873056769371,
+                      "zero_compute_margin_gap": 0.4672557485134652,
+                      "zero_compute_mass": 0.012517090452214084,
+                      "zero_compute_ramp_steps": 50,
+                      "zero_compute_term": 1.0837746113538742,
+                      "zero_compute_top1_ratio": 0.00439453125,
+                      "zero_compute_topk_margin": 0.0,
+                      "zero_compute_topk_ratio": 0.24339599609374998,
+                      "zero_compute_warmup_steps": 0,
+                      "zero_expert_usage": 1080.0
+                    }
+                  },
+                  "model_config": {
+                    "add_zero_compute_expert": true,
+                    "num_experts": 128,
+                    "top_k_experts": 8,
+                    "use_zero_compute_optimization": true
+                  },
+                  "num_router_keys": 90,
+                  "router_keys_sample": [
+                    "model.layers.0.router.per_expert_scale",
+                    "model.layers.0.router.proj.weight",
+                    "model.layers.0.router.scale",
+                    "model.layers.1.router.per_expert_scale",
+                    "model.layers.1.router.proj.weight",
+                    "model.layers.1.router.scale",
+                    "model.layers.10.router.per_expert_scale",
+                    "model.layers.10.router.proj.weight",
+                    "model.layers.10.router.scale",
+                    "model.layers.11.router.per_expert_scale",
+                    "model.layers.11.router.proj.weight",
+                    "model.layers.11.router.scale"
+                  ],
+                  "source_model_id": ""
+                },
+                "config_path": "/cache/router_artifacts/router_config.json",
+                "loaded": true,
+                "loaded_key_count": 90,
+                "loaded_keys_sample": [
+                  "model.layers.0.router.per_expert_scale",
+                  "model.layers.0.router.proj.weight",
+                  "model.layers.0.router.scale",
+                  "model.layers.1.router.per_expert_scale",
+                  "model.layers.1.router.proj.weight",
+                  "model.layers.1.router.scale",
+                  "model.layers.10.router.per_expert_scale",
+                  "model.layers.10.router.proj.weight",
+                  "model.layers.10.router.scale",
+                  "model.layers.11.router.per_expert_scale",
+                  "model.layers.11.router.proj.weight",
+                  "model.layers.11.router.scale"
+                ],
+                "path": "/cache/router_artifacts/router_state_dict.pt"
+              },
+              "step": 100,
+              "training_metrics": {
+                "current_lambda_zero_compute": 3.0,
+                "effective_batch_size": 16,
+                "entropy_loss": 4.542559911807379,
+                "entropy_term": 0.0,
+                "expert_usage_sample": [
+                  1239.0,
+                  2070.0,
+                  1500.0,
+                  1201.0,
+                  2009.0,
+                  1821.0,
+                  670.0,
+                  1778.0,
+                  1452.0,
+                  2154.0,
+                  2007.0,
+                  980.0,
+                  1320.0,
+                  1568.0,
+                  1522.0,
+                  700.0
+                ],
+                "grad_norm": 446.0,
+                "gradient_accumulation_steps": 8,
+                "lambda_entropy": 0.0,
+                "lambda_router": 1.0,
+                "lambda_zero_compute": 3.0,
+                "loss": 2.6639687418937683,
+                "micro_batch_size": 2,
+                "output_kl": 0.9797117039561272,
+                "output_kl_term": 0.9797117039561272,
+                "probe_output_kl": 1.0213916301727295,
+                "probe_router_entropy": 4.537958733240763,
+                "probe_router_kl": 0.05228007212281227,
+                "probe_same_expert_ratio": 0.8512044270833333,
+                "probe_zero_compute_loss": 0.5116603970527649,
+                "probe_zero_compute_margin_gap": 0.37900154244465134,
+                "probe_zero_compute_mass": 0.014502804105480513,
+                "probe_zero_compute_top1_ratio": 0.02294921875,
+                "probe_zero_compute_topk_ratio": 0.3021158854166667,
+                "router_entropy": 4.542559911807379,
+                "router_kl": 0.05281998496502638,
+                "router_kl_term": 0.05281998496502638,
+                "same_expert_ratio": 0.85335693359375,
+                "step": 100,
+                "tokens_per_optimizer_step": 8192,
+                "zero_compute_loss": 0.5438123419880867,
+                "zero_compute_margin_gap": 0.4289153911076331,
+                "zero_compute_mass": 0.013749805480862657,
+                "zero_compute_ramp_steps": 50,
+                "zero_compute_term": 1.6314370036125183,
+                "zero_compute_top1_ratio": 0.019120279947916666,
+                "zero_compute_topk_margin": 0.0,
+                "zero_compute_topk_ratio": 0.26689453125,
+                "zero_compute_warmup_steps": 0,
+                "zero_expert_usage": 4699.0
+              }
+            },
+            "model_config": {
+              "add_zero_compute_expert": true,
+              "num_experts": 128,
+              "top_k_experts": 8,
+              "use_zero_compute_optimization": true
+            },
+            "num_router_keys": 90,
+            "router_keys_sample": [
+              "model.layers.0.router.per_expert_scale",
+              "model.layers.0.router.proj.weight",
+              "model.layers.0.router.scale",
+              "model.layers.1.router.per_expert_scale",
+              "model.layers.1.router.proj.weight",
+              "model.layers.1.router.scale",
+              "model.layers.10.router.per_expert_scale",
+              "model.layers.10.router.proj.weight",
+              "model.layers.10.router.scale",
+              "model.layers.11.router.per_expert_scale",
+              "model.layers.11.router.proj.weight",
+              "model.layers.11.router.scale"
+            ],
+            "source_model_id": ""
+          },
+          "config_path": "/cache/router_artifacts/router_config.json",
+          "loaded": true,
+          "loaded_key_count": 90,
+          "loaded_keys_sample": [
+            "model.layers.0.router.per_expert_scale",
+            "model.layers.0.router.proj.weight",
+            "model.layers.0.router.scale",
+            "model.layers.1.router.per_expert_scale",
+            "model.layers.1.router.proj.weight",
+            "model.layers.1.router.scale",
+            "model.layers.10.router.per_expert_scale",
+            "model.layers.10.router.proj.weight",
+            "model.layers.10.router.scale",
+            "model.layers.11.router.per_expert_scale",
+            "model.layers.11.router.proj.weight",
+            "model.layers.11.router.scale"
+          ],
+          "path": "/cache/router_artifacts/router_state_dict.pt"
+        },
+        "step": 500,
+        "training_metrics": {
+          "current_lambda_zero_compute": 3.0,
+          "easy_token_ratio": 0.88720703125,
+          "effective_batch_size": 16,
+          "entropy_loss": 4.538989106814067,
+          "entropy_term": 0.009077978213628133,
+          "expert_usage_sample": [
+            1200.0,
+            2028.0,
+            1208.0,
+            783.0,
+            2040.0,
+            1847.0,
+            497.0,
+            1490.0,
+            1385.0,
+            2217.0,
+            1843.0,
+            1175.0,
+            1406.0,
+            1502.0,
+            1687.0,
+            771.0
+          ],
+          "grad_norm": 204.0,
+          "gradient_accumulation_steps": 8,
+          "lambda_entropy": 0.002,
+          "lambda_router": 1.0,
+          "lambda_zero_compute": 3.0,
+          "loss": 2.76311457157135,
+          "micro_batch_size": 2,
+          "output_kl": 1.1192611530423164,
+          "output_kl_term": 1.1192611530423164,
+          "probe_easy_token_ratio": 0.9248046875,
+          "probe_output_kl": 1.1542232036590576,
+          "probe_router_entropy": 4.538172864913941,
+          "probe_router_kl": 0.0544092059135437,
+          "probe_same_expert_ratio": 0.8352864583333334,
+          "probe_teacher_confidence_mean": 0.6827144622802734,
+          "probe_zero_compute_loss": 0.5005475282669067,
+          "probe_zero_compute_margin_gap": 0.3712589807061819,
+          "probe_zero_compute_mass": 0.01432527024565543,
+          "probe_zero_compute_token_weight_mean": 0.6357069611549377,
+          "probe_zero_compute_top1_hits_actual": 568.0,
+          "probe_zero_compute_top1_ratio": 0.017643229166666666,
+          "probe_zero_compute_top1_ratio_actual": 0.018489583333333334,
+          "probe_zero_compute_topk_hits_actual": 10002.0,
+          "probe_zero_compute_topk_ratio": 0.31997760956028976,
+          "probe_zero_compute_topk_ratio_actual": 0.0406982421875,
+          "router_entropy": 4.538989106814067,
+          "router_kl": 0.054629013407975435,
+          "router_kl_term": 0.054629013407975435,
+          "same_expert_ratio": 0.8434204101562499,
+          "step": 500,
+          "teacher_confidence_mean": 0.6225322559475899,
+          "tokens_per_optimizer_step": 8192,
+          "zero_compute_loss": 0.5267154797911644,
+          "zero_compute_margin_gap": 0.418523011850672,
+          "zero_compute_mass": 0.01356517664706988,
+          "zero_compute_ramp_steps": 50,
+          "zero_compute_term": 1.5801464468240738,
+          "zero_compute_token_weight_mean": 0.5501668378710747,
+          "zero_compute_top1_hits_actual": 459.0,
+          "zero_compute_top1_ratio": 0.014103190104166665,
+          "zero_compute_top1_ratio_actual": 0.014941406249999997,
+          "zero_compute_topk_hits_actual": 8916.125,
+          "zero_compute_topk_margin": 0.0,
+          "zero_compute_topk_ratio": 0.27892842232367093,
+          "zero_compute_topk_ratio_actual": 0.03627980550130208,
+          "zero_compute_warmup_steps": 0,
+          "zero_expert_usage": 3466.0
+        }
+      },
+      "model_config": {
+        "add_zero_compute_expert": true,
+        "num_experts": 128,
+        "top_k_experts": 8,
+        "use_zero_compute_optimization": true
+      },
+      "num_router_keys": 90,
+      "router_keys_sample": [
+        "model.layers.0.router.per_expert_scale",
+        "model.layers.0.router.proj.weight",
+        "model.layers.0.router.scale",
+        "model.layers.1.router.per_expert_scale",
+        "model.layers.1.router.proj.weight",
+        "model.layers.1.router.scale",
+        "model.layers.10.router.per_expert_scale",
+        "model.layers.10.router.proj.weight",
+        "model.layers.10.router.scale",
+        "model.layers.11.router.per_expert_scale",
+        "model.layers.11.router.proj.weight",
+        "model.layers.11.router.scale"
+      ],
+      "source_model_id": ""
+    },
+    "config_path": "/cache/router_artifacts/router_config.json",
+    "loaded": true,
+    "loaded_key_count": 90,
+    "loaded_keys_sample": [
+      "model.layers.0.router.per_expert_scale",
+      "model.layers.0.router.proj.weight",
+      "model.layers.0.router.scale",
+      "model.layers.1.router.per_expert_scale",
+      "model.layers.1.router.proj.weight",
+      "model.layers.1.router.scale",
+      "model.layers.10.router.per_expert_scale",
+      "model.layers.10.router.proj.weight",
+      "model.layers.10.router.scale",
+      "model.layers.11.router.per_expert_scale",
+      "model.layers.11.router.proj.weight",
+      "model.layers.11.router.scale"
+    ],
+    "path": "/cache/router_artifacts/router_state_dict.pt"
+  },
+  "source_model_id": "google/gemma-4-26B-A4B-it",
+  "torch_dtype": "torch.bfloat16"
+}

gemma4_optimization.py

@@ -0,0 +1,203 @@
+# gemma4_optimization.py
+
+"""
+Hayson Cheung, 2026, Oringinal Script written to optimize
+Gemma4 on Hugging Face's Transformers library. 
+
+LICENSED UNDER THE MIT LICENSE.
+
+This file contains optimized variants of Gemma4 text model components, including a mixin for remapping weights from original Gemma4 models to optimized versions. The optimizations include support for an additional zero-compute expert in the MoE router and experts, as well as adjustments to the router's projection and scaling parameters to accommodate the expanded expert set. The load_optimization_weights method enables loading weights from a base Gemma4 model while remapping tensors as needed for the optimized architecture.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from .modeling_gemma4 import (
+    Gemma4ForCausalLM,
+    Gemma4TextDecoderLayer,
+    Gemma4TextExperts,
+    Gemma4TextModel,
+    Gemma4TextRouter,
+)
+
+
+@dataclass(frozen=True)
+class Gemma4OptimizationLoadResult:
+    loaded_keys: tuple[str, ...]
+    skipped_keys: tuple[str, ...]
+
+    @property
+    def loaded_count(self) -> int:
+        return len(self.loaded_keys)
+
+    @property
+    def skipped_count(self) -> int:
+        return len(self.skipped_keys)
+
+
+class Gemma4OptimizationWeightsMixin:
+    """
+    Mixin for modules that need a custom remount step when loading weights
+    from an original Gemma4 model into an optimized variant.
+    """
+
+    def _remap_optimization_tensors(
+        self,
+        base_state_dict: dict[str, torch.Tensor],
+        target_state_dict: dict[str, torch.Tensor],
+    ) -> dict[str, torch.Tensor]:
+        return {}
+
+    def load_optimization_weights(self, base_model: nn.Module) -> Gemma4OptimizationLoadResult:
+        if not isinstance(self, nn.Module):
+            raise TypeError("Gemma4OptimizationWeightsMixin can only be used with nn.Module subclasses.")
+        if not isinstance(base_model, nn.Module):
+            raise TypeError("base_model must be an nn.Module.")
+
+        target_state_dict = self.state_dict()
+        loaded: dict[str, torch.Tensor] = {}
+
+        for module_name, module in self.named_modules():
+            if not isinstance(module, Gemma4OptimizationWeightsMixin):
+                continue
+
+            try:
+                base_module = base_model if module_name == "" else base_model.get_submodule(module_name)
+            except AttributeError:
+                continue
+
+            remapped_tensors = module._remap_optimization_tensors(base_module.state_dict(), module.state_dict())
+            for tensor_name, tensor_value in remapped_tensors.items():
+                full_name = f"{module_name}.{tensor_name}" if module_name else tensor_name
+                loaded[full_name] = tensor_value.to(
+                    device=target_state_dict[full_name].device,
+                    dtype=target_state_dict[full_name].dtype,
+                )
+
+        for tensor_name, tensor_value in base_model.state_dict().items():
+            if tensor_name in loaded:
+                continue
+            target_tensor = target_state_dict.get(tensor_name)
+            if target_tensor is None or target_tensor.shape != tensor_value.shape:
+                continue
+            loaded[tensor_name] = tensor_value.to(device=target_tensor.device, dtype=target_tensor.dtype)
+
+        self.load_state_dict(loaded, strict=False)
+
+        skipped = tuple(sorted(set(base_model.state_dict()) - set(loaded)))
+        return Gemma4OptimizationLoadResult(tuple(sorted(loaded)), skipped)
+
+    def _load_weights(self, base_model: nn.Module) -> Gemma4OptimizationLoadResult:
+        return self.load_optimization_weights(base_model)
+
+
+def get_total_optimized_experts(num_experts: int, add_zero_compute_expert: bool) -> int:
+    return num_experts + int(add_zero_compute_expert)
+
+
+class OptimizedGemma4TextExperts(Gemma4TextExperts):
+    def __init__(self, config):
+        super().__init__(config)
+        self.total_num_experts = get_total_optimized_experts(
+            self.num_experts, getattr(config, "add_zero_compute_expert", False)
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.total_num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx >= self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+
+class OptimizedGemma4TextRouter(Gemma4OptimizationWeightsMixin, Gemma4TextRouter):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_experts = config.num_experts
+        self.total_num_experts = get_total_optimized_experts(
+            self.num_experts, getattr(config, "add_zero_compute_expert", False)
+        )
+        self.proj = nn.Linear(config.hidden_size, self.total_num_experts, bias=False)
+        self.per_expert_scale = nn.Parameter(torch.ones(self.total_num_experts))
+
+    def _remap_optimization_tensors(
+        self,
+        base_state_dict: dict[str, torch.Tensor],
+        target_state_dict: dict[str, torch.Tensor],
+    ) -> dict[str, torch.Tensor]:
+        remapped: dict[str, torch.Tensor] = {}
+
+        base_proj = base_state_dict.get("proj.weight")
+        target_proj = target_state_dict.get("proj.weight")
+        if (
+            base_proj is not None
+            and target_proj is not None
+            and target_proj.shape[1] == base_proj.shape[1]
+            and target_proj.shape[0] == base_proj.shape[0] + 1
+        ):
+            expanded_proj = target_proj.clone()
+            expanded_proj.zero_()
+            expanded_proj[: base_proj.shape[0]].copy_(base_proj)
+            remapped["proj.weight"] = expanded_proj
+
+        base_per_expert_scale = base_state_dict.get("per_expert_scale")
+        target_per_expert_scale = target_state_dict.get("per_expert_scale")
+        if (
+            base_per_expert_scale is not None
+            and target_per_expert_scale is not None
+            and target_per_expert_scale.shape[0] == base_per_expert_scale.shape[0] + 1
+        ):
+            expanded_per_expert_scale = target_per_expert_scale.clone()
+            expanded_per_expert_scale.fill_(1.0)
+            expanded_per_expert_scale[: base_per_expert_scale.shape[0]].copy_(base_per_expert_scale)
+            remapped["per_expert_scale"] = expanded_per_expert_scale
+
+        return remapped
+
+
+class OptimizedGemma4TextDecoderLayer(Gemma4TextDecoderLayer):
+    router_class = OptimizedGemma4TextRouter
+    experts_class = OptimizedGemma4TextExperts
+
+
+class OptimizedGemma4TextModel(Gemma4OptimizationWeightsMixin, Gemma4TextModel):
+    decoder_layer_class = OptimizedGemma4TextDecoderLayer
+
+
+class OptimizedGemma4ForCausalLM(Gemma4OptimizationWeightsMixin, Gemma4ForCausalLM):
+    text_model_class = OptimizedGemma4TextModel
+
+
+__all__ = [
+    "Gemma4OptimizationLoadResult",
+    "Gemma4OptimizationWeightsMixin",
+    "OptimizedGemma4ForCausalLM",
+    "OptimizedGemma4TextDecoderLayer",
+    "OptimizedGemma4TextExperts",
+    "OptimizedGemma4TextModel",
+    "OptimizedGemma4TextRouter",
+    "get_total_optimized_experts",
+]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 2,
+  "eos_token_id": 1,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.5.3",
+  "use_cache": true
+}

model-00001-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3439313309f3d3a25d6f24c8d91c444f86f21b3c3def0b6c122cf722858ec4f
+size 4838439644

model-00002-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffad46c8227677da8b448b548da345e8151c6f320f4f79a378db0f012282f41f
+size 4913431124

model-00003-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb8becf37055a085fa3aa38beac7585ee3f3106451f3f7ee8944bd3e19468800
+size 4884594380

model-00004-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:806b09e6afa1fa84cb2e9e77f2627991ac8279785ad48da03620d9b2bf37e39d
+size 4913431164

model-00005-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f66e5b58eb65af1d13e440580fe6dc696988604ca60cf3d082f12bf5caa86c56
+size 4884594444

model-00006-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05ee10d6a8c3f74c1db8519a8089ba4871335f8fc0ffc196525c35e1a85d558d
+size 4913431196

model-00007-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d94e9293dd7099334cb7a72c2a529d82559d87bde5298e95dd36048c8fd9b85e
+size 4884594444

model-00008-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e139cbd5b61c7506513cd1fba878ae2a3738e84121438b1042865516c956608
+size 4913431196

model-00009-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c7e39645e9d4608d84b2327d9fb59d635d8d23201b62296572523016ddedb8a
+size 4884594444

model-00010-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52cc2118792e66eca31769eeffff2ccce8e9659f5f7c83bc9a09c6c3c4de727f
+size 4913431196

model-00011-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4574aeba49be3a8ebf26e718af8fd27ff0f1537ea47df763f03d74230fda7b1
+size 2998950720

model.safetensors.index.json

@@ -0,0 +1,666 @@
+{
+  "metadata": {
+    "total_parameters": 25971423774,
+    "total_size": 51942847608
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00011-of-00011.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.experts.down_proj": "model-00001-of-00011.safetensors",
+    "model.layers.0.experts.gate_up_proj": "model-00001-of-00011.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.layer_scalar": "model-00001-of-00011.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.post_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.post_feedforward_layernorm_1.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.post_feedforward_layernorm_2.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.pre_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.pre_feedforward_layernorm_2.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.router.per_expert_scale": "model-00001-of-00011.safetensors",
+    "model.layers.0.router.proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.router.scale": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.experts.down_proj": "model-00001-of-00011.safetensors",
+    "model.layers.1.experts.gate_up_proj": "model-00001-of-00011.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.layer_scalar": "model-00001-of-00011.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.post_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.post_feedforward_layernorm_1.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.post_feedforward_layernorm_2.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.pre_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.pre_feedforward_layernorm_2.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.router.per_expert_scale": "model-00001-of-00011.safetensors",
+    "model.layers.1.router.proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.router.scale": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.10.experts.down_proj": "model-00004-of-00011.safetensors",
+    "model.layers.10.experts.gate_up_proj": "model-00004-of-00011.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.layer_scalar": "model-00004-of-00011.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.post_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.post_feedforward_layernorm_1.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.post_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.pre_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.pre_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.router.per_expert_scale": "model-00004-of-00011.safetensors",
+    "model.layers.10.router.proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.router.scale": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.experts.down_proj": "model-00005-of-00011.safetensors",
+    "model.layers.11.experts.gate_up_proj": "model-00005-of-00011.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.layer_scalar": "model-00004-of-00011.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.post_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.post_feedforward_layernorm_1.weight": "model-00005-of-00011.safetensors",
+    "model.layers.11.post_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.11.pre_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.pre_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.11.router.per_expert_scale": "model-00004-of-00011.safetensors",
+    "model.layers.11.router.proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.router.scale": "model-00004-of-00011.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.12.experts.down_proj": "model-00005-of-00011.safetensors",
+    "model.layers.12.experts.gate_up_proj": "model-00005-of-00011.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.layer_scalar": "model-00005-of-00011.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.post_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.post_feedforward_layernorm_1.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.post_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.pre_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.pre_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.router.per_expert_scale": "model-00005-of-00011.safetensors",
+    "model.layers.12.router.proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.router.scale": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.experts.down_proj": "model-00005-of-00011.safetensors",
+    "model.layers.13.experts.gate_up_proj": "model-00005-of-00011.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.layer_scalar": "model-00005-of-00011.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.post_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.post_feedforward_layernorm_1.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.post_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.pre_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.pre_feedforward_layernorm_2.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.router.per_expert_scale": "model-00005-of-00011.safetensors",
+    "model.layers.13.router.proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.router.scale": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.experts.down_proj": "model-00006-of-00011.safetensors",
+    "model.layers.14.experts.gate_up_proj": "model-00006-of-00011.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.layer_scalar": "model-00005-of-00011.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.post_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.post_feedforward_layernorm_1.weight": "model-00006-of-00011.safetensors",
+    "model.layers.14.post_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.14.pre_feedforward_layernorm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.pre_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.14.router.per_expert_scale": "model-00005-of-00011.safetensors",
+    "model.layers.14.router.proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.router.scale": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00005-of-00011.safetensors",
+    "model.layers.15.experts.down_proj": "model-00006-of-00011.safetensors",
+    "model.layers.15.experts.gate_up_proj": "model-00006-of-00011.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.layer_scalar": "model-00006-of-00011.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.post_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.post_feedforward_layernorm_1.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.post_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.pre_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.pre_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.router.per_expert_scale": "model-00006-of-00011.safetensors",
+    "model.layers.15.router.proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.router.scale": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.experts.down_proj": "model-00006-of-00011.safetensors",
+    "model.layers.16.experts.gate_up_proj": "model-00006-of-00011.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.layer_scalar": "model-00006-of-00011.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.post_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.post_feedforward_layernorm_1.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.post_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.pre_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.pre_feedforward_layernorm_2.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.router.per_expert_scale": "model-00006-of-00011.safetensors",
+    "model.layers.16.router.proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.router.scale": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.experts.down_proj": "model-00007-of-00011.safetensors",
+    "model.layers.17.experts.gate_up_proj": "model-00007-of-00011.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.layer_scalar": "model-00006-of-00011.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.post_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.post_feedforward_layernorm_1.weight": "model-00007-of-00011.safetensors",
+    "model.layers.17.post_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.17.pre_feedforward_layernorm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.pre_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.17.router.per_expert_scale": "model-00006-of-00011.safetensors",
+    "model.layers.17.router.proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.router.scale": "model-00006-of-00011.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00006-of-00011.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00006-of-00011.safetensors",
+    "model.layers.18.experts.down_proj": "model-00007-of-00011.safetensors",
+    "model.layers.18.experts.gate_up_proj": "model-00007-of-00011.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.layer_scalar": "model-00007-of-00011.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.post_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.post_feedforward_layernorm_1.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.post_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.pre_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.pre_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.router.per_expert_scale": "model-00007-of-00011.safetensors",
+    "model.layers.18.router.proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.router.scale": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.experts.down_proj": "model-00007-of-00011.safetensors",
+    "model.layers.19.experts.gate_up_proj": "model-00007-of-00011.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.layer_scalar": "model-00007-of-00011.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.post_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.post_feedforward_layernorm_1.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.post_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.pre_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.pre_feedforward_layernorm_2.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.router.per_expert_scale": "model-00007-of-00011.safetensors",
+    "model.layers.19.router.proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.router.scale": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.2.experts.down_proj": "model-00002-of-00011.safetensors",
+    "model.layers.2.experts.gate_up_proj": "model-00002-of-00011.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.layer_scalar": "model-00001-of-00011.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.post_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.post_feedforward_layernorm_1.weight": "model-00002-of-00011.safetensors",
+    "model.layers.2.post_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.2.pre_feedforward_layernorm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.pre_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.2.router.per_expert_scale": "model-00001-of-00011.safetensors",
+    "model.layers.2.router.proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.router.scale": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00011.safetensors",
+    "model.layers.20.experts.down_proj": "model-00008-of-00011.safetensors",
+    "model.layers.20.experts.gate_up_proj": "model-00008-of-00011.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.layer_scalar": "model-00007-of-00011.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.post_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.post_feedforward_layernorm_1.weight": "model-00008-of-00011.safetensors",
+    "model.layers.20.post_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.20.pre_feedforward_layernorm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.pre_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.20.router.per_expert_scale": "model-00007-of-00011.safetensors",
+    "model.layers.20.router.proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.router.scale": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00007-of-00011.safetensors",
+    "model.layers.21.experts.down_proj": "model-00008-of-00011.safetensors",
+    "model.layers.21.experts.gate_up_proj": "model-00008-of-00011.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.layer_scalar": "model-00008-of-00011.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.post_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.post_feedforward_layernorm_1.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.post_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.pre_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.pre_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.router.per_expert_scale": "model-00008-of-00011.safetensors",
+    "model.layers.21.router.proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.router.scale": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.experts.down_proj": "model-00008-of-00011.safetensors",
+    "model.layers.22.experts.gate_up_proj": "model-00008-of-00011.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.layer_scalar": "model-00008-of-00011.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.post_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.post_feedforward_layernorm_1.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.post_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.pre_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.pre_feedforward_layernorm_2.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.router.per_expert_scale": "model-00008-of-00011.safetensors",
+    "model.layers.22.router.proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.router.scale": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.experts.down_proj": "model-00009-of-00011.safetensors",
+    "model.layers.23.experts.gate_up_proj": "model-00009-of-00011.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.layer_scalar": "model-00008-of-00011.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.post_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.post_feedforward_layernorm_1.weight": "model-00009-of-00011.safetensors",
+    "model.layers.23.post_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.23.pre_feedforward_layernorm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.pre_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.23.router.per_expert_scale": "model-00008-of-00011.safetensors",
+    "model.layers.23.router.proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.router.scale": "model-00008-of-00011.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00008-of-00011.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00008-of-00011.safetensors",
+    "model.layers.24.experts.down_proj": "model-00009-of-00011.safetensors",
+    "model.layers.24.experts.gate_up_proj": "model-00009-of-00011.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.layer_scalar": "model-00009-of-00011.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.post_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.post_feedforward_layernorm_1.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.post_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.pre_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.pre_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.router.per_expert_scale": "model-00009-of-00011.safetensors",
+    "model.layers.24.router.proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.router.scale": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.experts.down_proj": "model-00009-of-00011.safetensors",
+    "model.layers.25.experts.gate_up_proj": "model-00009-of-00011.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.layer_scalar": "model-00009-of-00011.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.post_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.post_feedforward_layernorm_1.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.post_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.pre_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.pre_feedforward_layernorm_2.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.router.per_expert_scale": "model-00009-of-00011.safetensors",
+    "model.layers.25.router.proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.router.scale": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.experts.down_proj": "model-00010-of-00011.safetensors",
+    "model.layers.26.experts.gate_up_proj": "model-00010-of-00011.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.layer_scalar": "model-00009-of-00011.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.post_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.post_feedforward_layernorm_1.weight": "model-00010-of-00011.safetensors",
+    "model.layers.26.post_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.26.pre_feedforward_layernorm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.pre_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.26.router.per_expert_scale": "model-00009-of-00011.safetensors",
+    "model.layers.26.router.proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.router.scale": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00009-of-00011.safetensors",
+    "model.layers.27.experts.down_proj": "model-00010-of-00011.safetensors",
+    "model.layers.27.experts.gate_up_proj": "model-00010-of-00011.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.layer_scalar": "model-00010-of-00011.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.post_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.post_feedforward_layernorm_1.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.post_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.pre_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.pre_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.router.per_expert_scale": "model-00010-of-00011.safetensors",
+    "model.layers.27.router.proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.router.scale": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.experts.down_proj": "model-00010-of-00011.safetensors",
+    "model.layers.28.experts.gate_up_proj": "model-00010-of-00011.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.layer_scalar": "model-00010-of-00011.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.post_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.post_feedforward_layernorm_1.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.post_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.pre_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.pre_feedforward_layernorm_2.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.router.per_expert_scale": "model-00010-of-00011.safetensors",
+    "model.layers.28.router.proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.router.scale": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.experts.down_proj": "model-00011-of-00011.safetensors",
+    "model.layers.29.experts.gate_up_proj": "model-00011-of-00011.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.layer_scalar": "model-00010-of-00011.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.post_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.post_feedforward_layernorm_1.weight": "model-00011-of-00011.safetensors",
+    "model.layers.29.post_feedforward_layernorm_2.weight": "model-00011-of-00011.safetensors",
+    "model.layers.29.pre_feedforward_layernorm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.pre_feedforward_layernorm_2.weight": "model-00011-of-00011.safetensors",
+    "model.layers.29.router.per_expert_scale": "model-00010-of-00011.safetensors",
+    "model.layers.29.router.proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.router.scale": "model-00010-of-00011.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00010-of-00011.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00010-of-00011.safetensors",
+    "model.layers.3.experts.down_proj": "model-00002-of-00011.safetensors",
+    "model.layers.3.experts.gate_up_proj": "model-00002-of-00011.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.layer_scalar": "model-00002-of-00011.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.post_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.post_feedforward_layernorm_1.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.post_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.pre_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.pre_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.router.per_expert_scale": "model-00002-of-00011.safetensors",
+    "model.layers.3.router.proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.router.scale": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.experts.down_proj": "model-00002-of-00011.safetensors",
+    "model.layers.4.experts.gate_up_proj": "model-00002-of-00011.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.layer_scalar": "model-00002-of-00011.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.post_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.post_feedforward_layernorm_1.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.post_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.pre_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.pre_feedforward_layernorm_2.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.router.per_expert_scale": "model-00002-of-00011.safetensors",
+    "model.layers.4.router.proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.router.scale": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.experts.down_proj": "model-00003-of-00011.safetensors",
+    "model.layers.5.experts.gate_up_proj": "model-00003-of-00011.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.layer_scalar": "model-00002-of-00011.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.post_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.post_feedforward_layernorm_1.weight": "model-00003-of-00011.safetensors",
+    "model.layers.5.post_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.5.pre_feedforward_layernorm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.pre_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.5.router.per_expert_scale": "model-00002-of-00011.safetensors",
+    "model.layers.5.router.proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.router.scale": "model-00002-of-00011.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00002-of-00011.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00011.safetensors",
+    "model.layers.6.experts.down_proj": "model-00003-of-00011.safetensors",
+    "model.layers.6.experts.gate_up_proj": "model-00003-of-00011.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.layer_scalar": "model-00003-of-00011.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.post_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.post_feedforward_layernorm_1.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.post_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.pre_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.pre_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.router.per_expert_scale": "model-00003-of-00011.safetensors",
+    "model.layers.6.router.proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.router.scale": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.experts.down_proj": "model-00003-of-00011.safetensors",
+    "model.layers.7.experts.gate_up_proj": "model-00003-of-00011.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.layer_scalar": "model-00003-of-00011.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.post_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.post_feedforward_layernorm_1.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.post_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.pre_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.pre_feedforward_layernorm_2.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.router.per_expert_scale": "model-00003-of-00011.safetensors",
+    "model.layers.7.router.proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.router.scale": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.experts.down_proj": "model-00004-of-00011.safetensors",
+    "model.layers.8.experts.gate_up_proj": "model-00004-of-00011.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.layer_scalar": "model-00003-of-00011.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.post_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.post_feedforward_layernorm_1.weight": "model-00004-of-00011.safetensors",
+    "model.layers.8.post_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.8.pre_feedforward_layernorm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.pre_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.8.router.per_expert_scale": "model-00003-of-00011.safetensors",
+    "model.layers.8.router.proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.router.scale": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00003-of-00011.safetensors",
+    "model.layers.9.experts.down_proj": "model-00004-of-00011.safetensors",
+    "model.layers.9.experts.gate_up_proj": "model-00004-of-00011.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.layer_scalar": "model-00004-of-00011.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.post_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.post_feedforward_layernorm_1.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.post_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.pre_feedforward_layernorm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.pre_feedforward_layernorm_2.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.router.per_expert_scale": "model-00004-of-00011.safetensors",
+    "model.layers.9.router.proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.router.scale": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00004-of-00011.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00004-of-00011.safetensors",
+    "model.norm.weight": "model-00011-of-00011.safetensors"
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f
+size 32169626

tokenizer_config.json

@@ -0,0 +1,95 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<content>(?:(?!\\<\\|tool_call\\>)(?!\\<turn\\|\\>).)+)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?:\\<turn\\|\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}