Fix MLX model_file: exclude probe from quantization, sanitize KV-shared layers

gemma_4_e2b_it_hybrid.py

@@ -186,6 +186,48 @@ class Model(gemma4_text.Model):
             per_layer_inputs=per_layer_inputs,
         )
 
+    # --- weight handling ----------------------------------------------------
+
+    def sanitize(self, weights):
+        """Stock sanitize + drop KV tensors the module tree does not build.
+
+        The HF checkpoint ships k/v projections and norms for the KV-shared
+        layers; the mlx model creates no such modules there (``has_kv`` is
+        False), and ``load_weights`` is strict about extra keys.
+        """
+        weights = super().sanitize(weights)
+        shared_start = self.args.num_hidden_layers - self.args.num_kv_shared_layers
+        if self.args.num_kv_shared_layers <= 0:
+            return weights
+
+        def is_unused_shared_kv(key: str) -> bool:
+            parts = key.split(".")
+            return (
+                len(parts) >= 5
+                and parts[0] == "model"
+                and parts[1] == "layers"
+                and parts[2].isdigit()
+                and int(parts[2]) >= shared_start
+                and parts[3] == "self_attn"
+                and parts[4] in ("k_proj", "v_proj", "k_norm", "v_norm")
+            )
+
+        return {k: v for k, v in weights.items() if not is_unused_shared_kv(k)}
+
+    @property
+    def quant_predicate(self):
+        """Never quantize the probe: its math is float32 by contract, and its
+        tiny Linears would otherwise be 4-bit-quantized by ``mlx_lm.convert``
+        (their input dims are multiples of the group size)."""
+        base = gemma4_text.Model.quant_predicate.fget(self)
+
+        def predicate(path, module):
+            if path.startswith("handoff_probe"):
+                return False
+            return base(path, module)
+
+        return predicate
+
     # --- scoring -----------------------------------------------------------
 
     def reset_probe(self) -> None: