Add files using upload-large-folder tool

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "LagunaForCausalLM"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_laguna.LagunaConfig",
+    "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
+  },
   "model_type": "laguna",
   "vocab_size": 100352,
   "hidden_size": 4096,
@@ -38,7 +42,8 @@
       "original_max_position_embeddings": 4096,
       "beta_slow": 1.0,
       "beta_fast": 64.0,
-      "attention_factor": 1.0
+      "attention_factor": 1.0,
+      "partial_rotary_factor": 1.0
     }
   },
   "moe_apply_router_weight_on_input": false,
@@ -510,18 +515,5 @@
     "sparsity_config": {},
     "transform_config": {},
     "version": "0.11.0"
-  },
-  "auto_map": {
-    "AutoConfig": "configuration_laguna.LagunaConfig",
-    "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
-  },
-  "rope_theta": 500000.0,
-  "rope_scaling": {
-    "rope_type": "yarn",
-    "factor": 32.0,
-    "original_max_position_embeddings": 4096,
-    "beta_slow": 1.0,
-    "beta_fast": 64.0,
-    "attention_factor": 1.0
   }
-}
+}

configuration_laguna.py

@@ -12,17 +12,11 @@
 # 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.
-"""
-Laguna configuration for transformers 4.56-4.x (used by vLLM).
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.modeling_rope_utils import RopeParameters
 
-This uses rope_theta + rope_scaling (legacy format) instead of
-rope_parameters (v5 format).
-"""
 
-from transformers.configuration_utils import PretrainedConfig
-
-
-class LagunaConfig(PretrainedConfig):
+class LagunaConfig(PreTrainedConfig):
     r"""
     Configuration class for Laguna model.
 
@@ -43,6 +37,20 @@ class LagunaConfig(PretrainedConfig):
         gating (`bool`, *optional*, defaults to `True`):
             Whether to use softplus output gating on attention. When True, a g_proj linear
             layer is added and attn_output = attn_output * softplus(g_proj(x)).
+        sliding_window (`int`, *optional*):
+            Sliding window attention size. Used by layers whose type in ``layer_types``
+            is ``"sliding_attention"``. When ``None``, all layers use full attention.
+        layer_types (`list[str]`, *optional*):
+            Per-layer attention type. Each element should be ``"sliding_attention"`` or
+            ``"global_attention"``. Length must equal ``num_hidden_layers``. When ``None``,
+            all layers default to global attention.
+        swa_attention_sink_enabled (`bool`, *optional*, defaults to `False`):
+            Whether to enable learnable attention sinks on sliding-window attention layers.
+            When enabled, a per-head bias parameter is added that allows the model to attend
+            to position 0 even when it falls outside the sliding window.
+        swa_rope_parameters (`RopeParameters`, *optional*):
+            Separate RoPE configuration for sliding-window attention layers. When ``None``,
+            SWA layers use the same RoPE as global attention layers.
         vocab_size (`int`, *optional*, defaults to 100352):
             Vocabulary size of the Laguna model.
         hidden_size (`int`, *optional*, defaults to 2048):
@@ -59,19 +67,6 @@ class LagunaConfig(PretrainedConfig):
             Maximum sequence length.
         rms_norm_eps (`float`, *optional*, defaults to 1e-6):
             Epsilon for RMSNorm layers.
-        rope_theta (`float`, *optional*, defaults to 500000.0):
-            Base frequency for RoPE embeddings.
-        rope_scaling (`dict`, *optional*):
-            RoPE scaling configuration (e.g. YaRN, linear).
-        sliding_window (`int`, *optional*):
-            Sliding window attention size. Used by layers whose type in ``layer_types``
-            is ``"sliding_attention"``. When ``None``, all layers use full attention.
-        layer_types (`list[str]`, *optional*):
-            Per-layer attention type. Each element should be ``"sliding_attention"`` or
-            ``"global_attention"``. Length must equal ``num_hidden_layers``. When ``None``,
-            all layers default to global attention.
-        swa_attention_sink_enabled (`bool`, *optional*, defaults to `False`):
-            Whether to enable learnable attention sinks on sliding-window attention layers.
         num_experts (`int`, *optional*, defaults to 256):
             Number of routed experts.
         num_experts_per_tok (`int`, *optional*, defaults to 16):
@@ -88,6 +83,8 @@ class LagunaConfig(PretrainedConfig):
             Layer indices that use dense MLP instead of MoE.
         router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
             Auxiliary loss coefficient for load balancing.
+        rope_parameters (`RopeParameters`, *optional*):
+            RoPE configuration. Defaults to rope_theta=500000.0.
     """
 
     model_type = "laguna"
@@ -126,12 +123,12 @@ class LagunaConfig(PretrainedConfig):
         rms_norm_eps: float = 1e-6,
         use_cache: bool = True,
         tie_word_embeddings: bool = False,
-        rope_theta: float = 500000.0,
-        rope_scaling: dict | None = None,
+        rope_parameters: RopeParameters | dict[str, RopeParameters] | None = None,
         attention_dropout: float = 0.0,
         sliding_window: int | None = None,
         layer_types: list[str] | None = None,
         swa_attention_sink_enabled: bool = False,
+        swa_rope_parameters: RopeParameters | None = None,
         num_experts: int = 256,
         num_experts_per_tok: int = 16,
         moe_intermediate_size: int = 1024,
@@ -147,6 +144,10 @@ class LagunaConfig(PretrainedConfig):
         if mlp_only_layers is None:
             mlp_only_layers = [0]
 
+        # Default rope_parameters with Laguna's theta
+        if rope_parameters is None:
+            rope_parameters = {"rope_type": "default", "rope_theta": 500000.0}
+
         self.vocab_size = vocab_size
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
@@ -162,13 +163,13 @@ class LagunaConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.rms_norm_eps = rms_norm_eps
         self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
+        self.rope_parameters = rope_parameters
         self.attention_dropout = attention_dropout
         # Sliding window attention arguments
         self.sliding_window = sliding_window
         self.layer_types = layer_types
         self.swa_attention_sink_enabled = swa_attention_sink_enabled
+        self.swa_rope_parameters = swa_rope_parameters
         # MoE arguments
         self.num_experts = num_experts
         self.num_experts_per_tok = num_experts_per_tok

model-00001-of-00045.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:393a5863004d5c402c065129a1f93da263bc00501466fb34fa578cdb6903d627
-size 5118019912
+oid sha256:b290b33be8f24274158893c4581c92652ddab16ca649163964fb4015e90d2944
+size 5118025988

model-00002-of-00045.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4873cb659f6efa610cffb4ff42ff11528ce05c795ff669b7163c86512aea6599
-size 5076112752
+oid sha256:51497babd4107a6ce13ac33bd22397221b5c40da8f2c58dc6a940597c45a4c48
+size 5076119676

model-00003-of-00045.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:968265e5db97a1a82f110587e1fd90eef53e844dfed80bf8e2bd9eff1dc5494a
-size 5119812952
+oid sha256:3a7f85878afb3156b64de58640d08672aa6ffcc40294a878dd2e9d821fb9b297
+size 5119814648

model.safetensors.index.json

@@ -4,7 +4,9 @@
   },
   "weight_map": {
     "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.0.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.0.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.0.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.0.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -23,7 +25,9 @@
     "model.layers.0.mlp.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00045.safetensors",
     "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.1.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.1.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.1.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.1.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -42,7 +46,9 @@
     "model.layers.1.mlp.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00045.safetensors",
     "model.layers.10.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.10.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.10.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.10.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.10.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.10.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -63,7 +69,9 @@
     "model.layers.10.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.10.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.11.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.11.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.11.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.11.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.11.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.11.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -84,7 +92,9 @@
     "model.layers.11.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.11.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.12.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.12.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.12.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.12.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.12.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.12.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -105,7 +115,9 @@
     "model.layers.12.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.12.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.13.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.13.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.13.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.13.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.13.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.13.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -126,7 +138,9 @@
     "model.layers.13.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.13.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.14.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.14.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.14.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.14.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.14.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.14.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -147,7 +161,9 @@
     "model.layers.14.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.14.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.15.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.15.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.15.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.15.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.15.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.15.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -168,7 +184,9 @@
     "model.layers.15.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.15.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.16.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.16.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.16.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.16.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.16.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.16.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -189,7 +207,9 @@
     "model.layers.16.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.16.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.17.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.17.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.17.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.17.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.17.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.17.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -210,7 +230,9 @@
     "model.layers.17.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.17.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.18.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.18.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.18.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.18.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.18.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.18.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.18.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -231,7 +253,9 @@
     "model.layers.18.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.18.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.19.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.19.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.19.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.19.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.19.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.19.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.19.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -252,7 +276,9 @@
     "model.layers.19.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.19.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.2.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.2.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.2.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.2.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -271,7 +297,9 @@
     "model.layers.2.mlp.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00045.safetensors",
     "model.layers.20.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.20.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.20.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.20.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.20.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.20.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -292,7 +320,9 @@
     "model.layers.20.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.20.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.21.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.21.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.21.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.21.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.21.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.21.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.21.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -313,7 +343,9 @@
     "model.layers.21.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.21.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.22.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.22.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.22.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.22.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.22.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.22.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.22.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -334,7 +366,9 @@
     "model.layers.22.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.22.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.23.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.23.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.23.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.23.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.23.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.23.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.23.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -355,7 +389,9 @@
     "model.layers.23.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.23.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.24.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.24.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.24.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.24.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.24.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.24.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.24.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -376,7 +412,9 @@
     "model.layers.24.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.24.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.25.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.25.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.25.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.25.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.25.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.25.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.25.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -397,7 +435,9 @@
     "model.layers.25.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.25.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.26.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.26.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.26.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.26.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.26.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.26.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.26.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -418,7 +458,9 @@
     "model.layers.26.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.26.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.27.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.27.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.27.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.27.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.27.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.27.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.27.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -439,7 +481,9 @@
     "model.layers.27.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.27.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.28.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.28.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.28.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.28.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.28.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.28.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.28.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -460,7 +504,9 @@
     "model.layers.28.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.28.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.29.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.29.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.29.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.29.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.29.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.29.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.29.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -481,7 +527,9 @@
     "model.layers.29.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.29.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.3.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.3.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.3.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.3.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -502,7 +550,9 @@
     "model.layers.3.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.3.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.30.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.30.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.30.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.30.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.30.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.30.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.30.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -523,7 +573,9 @@
     "model.layers.30.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.30.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.31.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.31.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.31.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.31.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.31.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.31.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.31.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -544,7 +596,9 @@
     "model.layers.31.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.31.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.32.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.32.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.32.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.32.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.32.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.32.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.32.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -565,7 +619,9 @@
     "model.layers.32.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.32.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.33.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.33.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.33.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.33.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.33.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.33.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.33.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -586,7 +642,9 @@
     "model.layers.33.mlp.shared_expert.down_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.33.mlp.shared_expert.down_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.34.self_attn.k_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.34.self_attn.k_scale": "model-00001-of-00045.safetensors",
     "model.layers.34.self_attn.q_norm.weight": "model-00001-of-00045.safetensors",
+    "model.layers.34.self_attn.v_scale": "model-00001-of-00045.safetensors",
     "model.layers.34.self_attn.g_proj.weight": "model-00001-of-00045.safetensors",
     "model.layers.34.self_attn.g_proj.weight_scale": "model-00001-of-00045.safetensors",
     "model.layers.34.self_attn.k_proj.weight": "model-00001-of-00045.safetensors",
@@ -607,7 +665,9 @@
     "model.layers.34.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.34.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.35.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.35.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.35.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.35.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.35.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.35.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.35.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -628,7 +688,9 @@
     "model.layers.35.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.35.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.36.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.36.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.36.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.36.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.36.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.36.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.36.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -649,7 +711,9 @@
     "model.layers.36.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.36.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.37.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.37.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.37.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.37.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.37.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.37.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.37.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -670,7 +734,9 @@
     "model.layers.37.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.37.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.38.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.38.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.38.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.38.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.38.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.38.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.38.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -691,7 +757,9 @@
     "model.layers.38.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.38.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.39.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.39.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.39.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.39.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.39.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.39.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.39.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -712,7 +780,9 @@
     "model.layers.39.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.39.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.4.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.4.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.4.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.4.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.4.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.4.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -733,7 +803,9 @@
     "model.layers.4.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.4.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.40.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.40.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.40.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.40.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.40.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.40.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.40.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -754,7 +826,9 @@
     "model.layers.40.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.40.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.41.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.41.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.41.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.41.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.41.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.41.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.41.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -775,7 +849,9 @@
     "model.layers.41.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.41.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.42.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.42.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.42.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.42.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.42.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.42.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.42.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -796,7 +872,9 @@
     "model.layers.42.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.42.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.43.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.43.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.43.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.43.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.43.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.43.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.43.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -817,7 +895,9 @@
     "model.layers.43.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.43.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.44.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.44.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.44.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.44.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.44.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.44.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.44.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -838,7 +918,9 @@
     "model.layers.44.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.44.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.45.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.45.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.45.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.45.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.45.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.45.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.45.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -859,7 +941,9 @@
     "model.layers.45.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.45.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.46.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.46.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.46.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.46.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.46.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.46.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.46.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -880,7 +964,9 @@
     "model.layers.46.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.46.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.47.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.47.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.47.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.47.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.47.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.47.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.47.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -901,7 +987,9 @@
     "model.layers.47.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.47.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.48.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.48.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.48.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.48.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.48.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.48.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.48.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -922,7 +1010,9 @@
     "model.layers.48.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.48.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.49.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.49.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.49.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.49.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.49.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.49.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.49.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -943,7 +1033,9 @@
     "model.layers.49.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.49.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.5.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.5.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.5.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.5.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.5.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.5.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -964,7 +1056,9 @@
     "model.layers.5.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.5.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.50.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.50.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.50.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.50.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.50.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.50.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.50.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -985,7 +1079,9 @@
     "model.layers.50.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.50.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.51.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.51.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.51.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.51.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.51.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.51.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.51.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1006,7 +1102,9 @@
     "model.layers.51.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.51.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.52.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.52.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.52.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.52.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.52.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.52.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.52.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1027,7 +1125,9 @@
     "model.layers.52.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.52.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.53.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.53.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.53.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.53.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.53.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.53.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.53.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1048,7 +1148,9 @@
     "model.layers.53.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.53.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.54.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.54.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.54.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.54.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.54.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.54.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.54.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1069,7 +1171,9 @@
     "model.layers.54.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.54.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.55.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.55.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.55.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.55.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.55.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.55.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.55.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1090,7 +1194,9 @@
     "model.layers.55.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.55.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.56.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.56.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.56.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.56.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.56.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.56.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.56.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1111,7 +1217,9 @@
     "model.layers.56.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.56.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.57.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.57.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.57.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.57.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.57.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.57.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.57.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1132,7 +1240,9 @@
     "model.layers.57.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.57.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.58.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.58.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.58.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.58.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.58.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.58.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.58.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1153,7 +1263,9 @@
     "model.layers.58.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.58.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.59.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.59.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.59.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.59.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.59.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.59.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.59.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1174,7 +1286,9 @@
     "model.layers.59.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.59.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.6.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.6.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.6.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.6.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.6.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.6.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1195,7 +1309,9 @@
     "model.layers.6.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.6.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.60.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.60.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.60.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.60.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.60.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.60.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.60.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1216,7 +1332,9 @@
     "model.layers.60.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.60.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.61.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.61.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.61.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.61.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.61.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.61.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.61.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1237,7 +1355,9 @@
     "model.layers.61.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.61.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.62.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.62.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.62.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.62.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.62.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.62.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.62.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1258,7 +1378,9 @@
     "model.layers.62.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.62.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.63.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.63.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.63.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.63.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.63.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.63.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.63.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1279,7 +1401,9 @@
     "model.layers.63.mlp.shared_expert.down_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.63.mlp.shared_expert.down_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.64.self_attn.k_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.64.self_attn.k_scale": "model-00002-of-00045.safetensors",
     "model.layers.64.self_attn.q_norm.weight": "model-00002-of-00045.safetensors",
+    "model.layers.64.self_attn.v_scale": "model-00002-of-00045.safetensors",
     "model.layers.64.self_attn.g_proj.weight": "model-00002-of-00045.safetensors",
     "model.layers.64.self_attn.g_proj.weight_scale": "model-00002-of-00045.safetensors",
     "model.layers.64.self_attn.k_proj.weight": "model-00002-of-00045.safetensors",
@@ -1300,7 +1424,9 @@
     "model.layers.64.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.64.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.65.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.65.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.65.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.65.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.65.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.65.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.65.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1321,7 +1447,9 @@
     "model.layers.65.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.65.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.66.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.66.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.66.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.66.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.66.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.66.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.66.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1342,7 +1470,9 @@
     "model.layers.66.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.66.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.67.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.67.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.67.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.67.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.67.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.67.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.67.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1363,7 +1493,9 @@
     "model.layers.67.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.67.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.68.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.68.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.68.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.68.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.68.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.68.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.68.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1384,7 +1516,9 @@
     "model.layers.68.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.68.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.69.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.69.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.69.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.69.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.69.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.69.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.69.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1405,7 +1539,9 @@
     "model.layers.69.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.69.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.7.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.7.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.7.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.7.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.7.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.7.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.7.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1426,7 +1562,9 @@
     "model.layers.7.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.7.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.8.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.8.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.8.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.8.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.8.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.8.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.8.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -1447,7 +1585,9 @@
     "model.layers.8.mlp.shared_expert.down_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.8.mlp.shared_expert.down_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.9.self_attn.k_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.9.self_attn.k_scale": "model-00003-of-00045.safetensors",
     "model.layers.9.self_attn.q_norm.weight": "model-00003-of-00045.safetensors",
+    "model.layers.9.self_attn.v_scale": "model-00003-of-00045.safetensors",
     "model.layers.9.self_attn.g_proj.weight": "model-00003-of-00045.safetensors",
     "model.layers.9.self_attn.g_proj.weight_scale": "model-00003-of-00045.safetensors",
     "model.layers.9.self_attn.k_proj.weight": "model-00003-of-00045.safetensors",
@@ -104381,146 +104521,6 @@
     "model.layers.69.mlp.experts.255.down_proj.weight": "model-00045-of-00045.safetensors",
     "model.layers.69.mlp.experts.255.gate_proj.weight_scale": "model-00045-of-00045.safetensors",
     "model.layers.69.mlp.experts.255.up_proj.weight_scale": "model-00045-of-00045.safetensors",
-    "model.layers.69.mlp.experts.255.down_proj.weight_scale": "model-00045-of-00045.safetensors",
-    "model.layers.0.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.0.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.1.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.1.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.10.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.10.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.11.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.11.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.12.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.12.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.13.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.13.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.14.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.14.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.15.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.15.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.16.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.16.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.17.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.17.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.18.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.18.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.19.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.19.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.2.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.2.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.20.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.20.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.21.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.21.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.22.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.22.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.23.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.23.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.24.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.24.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.25.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.25.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.26.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.26.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.27.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.27.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.28.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.28.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.29.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.29.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.3.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.3.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.30.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.30.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.31.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.31.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.32.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.32.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.33.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.33.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.34.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.34.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.35.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.35.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.36.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.36.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.37.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.37.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.38.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.38.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.39.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.39.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.4.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.4.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.40.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.40.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.41.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.41.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.42.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.42.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.43.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.43.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.44.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.44.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.45.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.45.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.46.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.46.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.47.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.47.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.48.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.48.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.49.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.49.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.5.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.5.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.50.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.50.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.51.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.51.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.52.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.52.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.53.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.53.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.54.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.54.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.55.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.55.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.56.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.56.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.57.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.57.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.58.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.58.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.59.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.59.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.6.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.6.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.60.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.60.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.61.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.61.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.62.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.62.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.63.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.63.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.64.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.64.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.65.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.65.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.66.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.66.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.67.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.67.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.68.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.68.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.69.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.69.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.7.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.7.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.8.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.8.self_attn.v_scale": "model-kv_scales.safetensors",
-    "model.layers.9.self_attn.k_scale": "model-kv_scales.safetensors",
-    "model.layers.9.self_attn.v_scale": "model-kv_scales.safetensors"
+    "model.layers.69.mlp.experts.255.down_proj.weight_scale": "model-00045-of-00045.safetensors"
   }
-}
+}

modeling_laguna.py

@@ -12,30 +12,36 @@
 # 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.
-"""
-Laguna model implementation for transformers 4.56-4.x (used by vLLM).
-
-This avoids v5-only APIs (use_kernel_forward_from_hub, create_causal_mask,
-dynamic_rope_update, auto_docstring, can_return_tuple, etc.) while keeping
-the architecture identical to the v5 version.
-"""
 
 from typing import Optional
+from collections.abc import Callable
 
 import torch
 import torch.nn.functional as F
 from torch import nn
+from transformers import initialization as init
+from transformers.utils import auto_docstring, can_return_tuple, is_grouped_mm_available
 from transformers.generation import GenerationMixin
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
-from transformers.utils.generic import OutputRecorder, check_model_inputs
+from transformers.integrations import (
+    use_kernelized_func,
+    use_kernel_func_from_hub,
+    use_kernel_forward_from_hub,
+)
+from transformers.masking_utils import create_causal_mask
+from transformers.utils.generic import OutputRecorder, TransformersKwargs, maybe_autocast, check_model_inputs
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import MoeModelOutputWithPast, MoeCausalLMOutputWithPast
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.processing_utils import Unpack
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 
 from .configuration_laguna import LagunaConfig
 
 
+@use_kernel_forward_from_hub("RMSNorm")
 class LagunaRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         """
@@ -66,41 +72,55 @@ class LagunaRotaryEmbedding(nn.Module):
 
         self.config = config
 
-        # v4 uses rope_theta + rope_scaling (top-level config fields)
-        rope_type = "default"
-        if config.rope_scaling is not None:
-            rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default"))
-
-        self.rope_type = rope_type
-        if self.rope_type == "default":
-            inv_freq, self.attention_scaling = self._compute_default_rope_parameters(config, device)
-        else:
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
             rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-            inv_freq, self.attention_scaling = rope_init_fn(config, device)
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
 
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
 
     @staticmethod
-    def _compute_default_rope_parameters(
-        config: LagunaConfig,
+    def compute_default_rope_parameters(
+        config: LagunaConfig | None = None,
         device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
     ) -> tuple["torch.Tensor", float]:
-        base = config.rope_theta
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+
+        Returns
+        -------
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
         dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
-        attention_factor = 1.0
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
         inv_freq = 1.0 / (
             base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
         )
         return inv_freq, attention_factor
 
     @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
     def forward(self, x, position_ids):
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
         position_ids_expanded = position_ids[:, None, :].float()
 
         device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             emb = torch.cat((freqs, freqs), dim=-1)
             cos = emb.cos() * self.attention_scaling
@@ -200,8 +220,27 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
+@use_kernel_func_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors."""
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+
+    Returns
+    -------
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
     cos = cos.unsqueeze(unsqueeze_dim)
     sin = sin.unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)
@@ -229,7 +268,7 @@ def eager_attention_forward(
     attention_mask: torch.Tensor | None,
     scaling: float,
     dropout: float = 0.0,
-    **kwargs,
+    **kwargs: Unpack[TransformersKwargs],
 ):
     key_states = repeat_kv(key, module.num_key_value_groups)
     value_states = repeat_kv(value, module.num_key_value_groups)
@@ -247,9 +286,13 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
+# Laguna attention is identical to Qwen2MoE attention except:
+# - No QKV bias
+# - Explicit head_dim from config
+# - Output gating: attn_output = attn_output * softplus(g_proj(hidden_states))
+# - No sliding window (full attention only)
+@use_kernelized_func(apply_rotary_pos_emb)
 class LagunaAttention(nn.Module):
-    """Laguna attention with QK normalization and output gating."""
-
     def __init__(self, config: LagunaConfig, layer_idx: int):
         super().__init__()
         self.config = config
@@ -278,7 +321,7 @@ class LagunaAttention(nn.Module):
         attention_mask: torch.Tensor | None,
         past_key_values: Cache | None = None,
         cache_position: torch.LongTensor | None = None,
-        **kwargs,
+        **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
@@ -299,10 +342,11 @@ class LagunaAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
             cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
             key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
-        attention_interface = eager_attention_forward
+        attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
@@ -320,6 +364,7 @@ class LagunaAttention(nn.Module):
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
 
         # Laguna-specific: apply gating BEFORE o_proj
+        # gate values are computed from original hidden_states, applied in attention dimension
         gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
         attn_output = attn_output * gate
 
@@ -328,7 +373,7 @@ class LagunaAttention(nn.Module):
         return attn_output, attn_weights
 
 
-class LagunaDecoderLayer(nn.Module):
+class LagunaDecoderLayer(GradientCheckpointingLayer):
     """Laguna decoder layer with gated attention and sigmoid-routed MoE."""
 
     def __init__(self, config: LagunaConfig, layer_idx: int):
@@ -354,7 +399,7 @@ class LagunaDecoderLayer(nn.Module):
         use_cache: bool | None = False,
         cache_position: torch.LongTensor | None = None,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
-        **kwargs,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -379,63 +424,32 @@ class LagunaDecoderLayer(nn.Module):
         return hidden_states
 
 
+@auto_docstring
 class LagunaPreTrainedModel(PreTrainedModel):
-    config_class = LagunaConfig
+    config: LagunaConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["LagunaDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
-    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = True
-    _supports_cache_class = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = (
+        is_grouped_mm_available()
+    )  # https://huggingface.co/docs/transformers/experts_interface#torchcompile
+    _supports_attention_backend = True
     _can_record_outputs = {
         "router_logits": OutputRecorder(LagunaTopKRouter, index=0),
         "hidden_states": LagunaDecoderLayer,
         "attentions": LagunaAttention,
     }
 
+    @torch.no_grad()
     def _init_weights(self, module):
+        super()._init_weights(module)
         std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, LagunaTopKRouter):
-            module.weight.data.normal_(mean=0.0, std=std)
-
-
-def _prepare_4d_causal_attention_mask_with_cache_position(
-    attention_mask: torch.Tensor,
-    sequence_length: int,
-    target_length: int,
-    dtype: torch.dtype,
-    device: torch.device,
-    cache_position: torch.Tensor,
-    batch_size: int,
-):
-    """Create 4D causal mask from 2D attention mask, compatible with transformers 4.x."""
-    if attention_mask is not None and attention_mask.dim() == 4:
-        # Already a 4D mask
-        causal_mask = attention_mask
-    else:
-        min_dtype = torch.finfo(dtype).min
-        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
-        if sequence_length != 1:
-            causal_mask = torch.triu(causal_mask, diagonal=1)
-        causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
-        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
-        if attention_mask is not None:
-            causal_mask = causal_mask.clone()
-            mask_length = attention_mask.shape[-1]
-            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
-            padding_mask = padding_mask == 0
-            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
-
-    return causal_mask
+        if isinstance(module, LagunaTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
 
 
 class LagunaModel(LagunaPreTrainedModel):
@@ -465,13 +479,13 @@ class LagunaModel(LagunaPreTrainedModel):
         inputs_embeds: torch.FloatTensor | None = None,
         use_cache: bool | None = None,
         cache_position: torch.LongTensor | None = None,
-        **kwargs,
+        **kwargs: Unpack[TransformersKwargs],
     ):
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
+            past_key_values = DynamicCache(config=self.config)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
@@ -485,42 +499,30 @@ class LagunaModel(LagunaPreTrainedModel):
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(
-            attention_mask,
-            sequence_length=inputs_embeds.shape[1],
-            target_length=cache_position[-1].item() + 1 if cache_position is not None else inputs_embeds.shape[1],
-            dtype=inputs_embeds.dtype,
-            device=inputs_embeds.device,
+        # Laguna uses full attention only (no sliding window)
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
             cache_position=cache_position,
-            batch_size=inputs_embeds.shape[0],
+            past_key_values=past_key_values,
+            position_ids=position_ids,
         )
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            if self.gradient_checkpointing and self.training:
-                hidden_states = self._gradient_checkpointing_func(
-                    decoder_layer.__call__,
-                    hidden_states,
-                    causal_mask,
-                    position_ids,
-                    past_key_values,
-                    use_cache,
-                    cache_position,
-                    position_embeddings,
-                )
-            else:
-                hidden_states = decoder_layer(
-                    hidden_states,
-                    attention_mask=causal_mask,
-                    position_ids=position_ids,
-                    past_key_values=past_key_values,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                    position_embeddings=position_embeddings,
-                    **kwargs,
-                )
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
 
         hidden_states = self.norm(hidden_states)
 
@@ -537,9 +539,28 @@ def load_balancing_loss_func(
     attention_mask: torch.Tensor | None = None,
 ) -> torch.Tensor | int:
     r"""
-    Computes auxiliary load balancing loss as in Switch Transformer.
-
-    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details.
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns
+    -------
+        The auxiliary loss.
     """
     if gate_logits is None or not isinstance(gate_logits, tuple):
         return 0
@@ -555,12 +576,16 @@ def load_balancing_loss_func(
     expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
 
     if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
         tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
         router_prob_per_expert = torch.mean(routing_weights, dim=0)
     else:
         batch_size, sequence_length = attention_mask.shape
         num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
 
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
         expert_attention_mask = (
             attention_mask[None, :, :, None, None]
             .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
@@ -568,10 +593,12 @@ def load_balancing_loss_func(
             .to(compute_device)
         )
 
+        # Compute the percentage of tokens routed to each experts
         tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
             expert_attention_mask, dim=0
         )
 
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
         router_per_expert_attention_mask = (
             attention_mask[None, :, :, None]
             .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
@@ -579,6 +606,7 @@ def load_balancing_loss_func(
             .to(compute_device)
         )
 
+        # Compute the average probability of routing to these experts
         router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
             router_per_expert_attention_mask, dim=0
         )
@@ -587,9 +615,11 @@ def load_balancing_loss_func(
     return overall_loss * num_experts
 
 
+@auto_docstring
 class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
 
     def __init__(self, config):
         super().__init__(config)
@@ -603,6 +633,8 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
         # Initialize weights and apply final processing
         self.post_init()
 
+    @can_return_tuple
+    @auto_docstring
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
@@ -615,17 +647,21 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
         output_router_logits: bool | None = None,
         cache_position: torch.LongTensor | None = None,
         logits_to_keep: int | torch.Tensor = 0,
-        **kwargs,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> MoeCausalLMOutputWithPast:
         r"""
-        Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-        config.vocab_size]` or -100. Tokens with indices set to `-100` are ignored (masked), the loss is
-        only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        Labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
         """
+        # TODO (Joe) add example here after we got rid of the stale mistral example
+
         output_router_logits = (
             output_router_logits if output_router_logits is not None else self.config.output_router_logits
         )
 
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs: MoeModelOutputWithPast = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -639,6 +675,7 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
         )
 
         hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.lm_head(hidden_states[:, slice_indices, :])
 

tokenizer.json

@@ -649,7 +649,7 @@
       {
         "type": "Split",
         "pattern": {
-          "Regex": "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
+          "Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
         },
         "behavior": "Isolated",
         "invert": false
@@ -501192,4 +501192,4 @@
       ]
     ]
   }
-}
+}