allow per model quant config (#349)

mlx_lm/convert.py

@@ -49,7 +49,6 @@ def mixed_quant_predicate_builder(
     def mixed_quant_predicate(
         path: str,
         module: nn.Module,
-        config: dict,
     ) -> Union[bool, dict]:
         """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
         Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265

mlx_lm/models/qwen3_moe.py

@@ -235,6 +235,15 @@ class Model(nn.Module):
                     weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
         return weights
 
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/utils.py

@@ -476,6 +476,8 @@ def quantize_model(
     quantized_config = copy.deepcopy(config)
     quantized_config["quantization"] = {"group_size": q_group_size, "bits": q_bits}
 
+    quant_predicate = quant_predicate or getattr(model, "quant_predicate", None)
+
     def base_predicate(path, module):
         if not hasattr(module, "to_quantized"):
             return False
@@ -487,8 +489,9 @@ def quantize_model(
     def wrapped_predicate(p, m):
         bool_or_params = base_predicate(p, m)
         if bool_or_params:
-            bool_or_params = quant_predicate(p, m, config)
-        quantized_config["quantization"][p] = bool_or_params
+            bool_or_params = quant_predicate(p, m)
+        if isinstance(bool_or_params, dict):
+            quantized_config["quantization"][p] = bool_or_params
         return bool_or_params
 
     nn.quantize(