Fix mixed-bit quantized load for sanitize()-derived MLA projections

mlx_lm/utils.py

@@ -377,6 +377,33 @@ def load_config(model_path: Path) -> dict:
     return config
 
 
+def infer_quant_config(path: str, module: nn.Module, weights: dict) -> dict:
+    """Recover the group_size, bits and mode a saved weight was packed with.
+
+    Use this for paths the per-tensor quantization map does not name, where the
+    top-level default can be wrong. ``module`` must still be unquantized.
+    """
+    scales = weights[f"{path}.scales"]
+    in_dims = module.weight.shape[-1]
+    group_size = in_dims // scales.shape[-1]
+    bits = (weights[f"{path}.weight"].shape[-1] * 32) // in_dims
+    # Only affine keeps the scales in the weight dtype. Each of the other modes
+    # allows exactly one (bits, group_size) pair.
+    if scales.dtype != mx.uint8:
+        return {"group_size": group_size, "bits": bits, "mode": "affine"}
+    if (bits, group_size) == (4, 16):
+        return {"group_size": group_size, "bits": bits, "mode": "nvfp4"}
+    if (bits, group_size) == (4, 32):
+        return {"group_size": group_size, "bits": bits, "mode": "mxfp4"}
+    if (bits, group_size) == (8, 32):
+        return {"group_size": group_size, "bits": bits, "mode": "mxfp8"}
+
+    raise ValueError(
+        f"Cannot infer the quantization mode of {path}: "
+        f"{bits} bits with group size {group_size}."
+    )
+
+
 def load_model(
     model_path: Path,
     lazy: bool = False,
@@ -463,7 +490,9 @@ def class_predicate(p, m):
                 return config["quantization"][p]
             if not hasattr(m, "to_quantized"):
                 return False
-            return f"{p}.scales" in weights
+            if f"{p}.scales" not in weights:
+                return False
+            return infer_quant_config(p, m, weights)
 
         nn.quantize(
             model,

tests/test_utils.py

@@ -1,5 +1,6 @@
 # Copyright © 2024 Apple Inc.
 
+import dataclasses
 import json
 import os
 import tempfile
@@ -254,6 +255,108 @@ def test_load_model_gemma4_with_per_layer_projection_quantization(self):
             mx.eval(logits)
             self.assertEqual(logits.shape, (1, 3, args.vocab_size))
 
+    def test_infer_quant_config(self):
+        from mlx_lm.models.mla import MultiLinear
+
+        for mode, bits, group_size in [
+            ("affine", 3, 64),
+            ("affine", 4, 32),
+            ("affine", 8, 128),
+            ("mxfp4", 4, 32),
+            ("mxfp8", 8, 32),
+            ("nvfp4", 4, 16),
+        ]:
+            for name, layer in [
+                ("linear", nn.Linear(256, 128, bias=False)),
+                ("multi_linear", MultiLinear(256, 128, 4)),
+            ]:
+                with self.subTest(
+                    mode=mode, bits=bits, group_size=group_size, layer=name
+                ):
+                    q = layer.to_quantized(group_size=group_size, bits=bits, mode=mode)
+                    weights = {"l.weight": q.weight, "l.scales": q.scales}
+                    self.assertEqual(
+                        utils.infer_quant_config("l", layer, weights),
+                        {"group_size": group_size, "bits": bits, "mode": mode},
+                    )
+
+    def test_infer_quant_config_unknown_packing(self):
+        # uint8 scales say the weight is not affine, but no mode packs 6 bits
+        # into groups of 64.
+        layer = nn.Linear(256, 128, bias=False)
+        weights = {
+            "l.weight": mx.zeros((128, 48), mx.uint32),
+            "l.scales": mx.zeros((128, 4), mx.uint8),
+        }
+        with self.assertRaises(ValueError):
+            utils.infer_quant_config("l", layer, weights)
+
+    def test_load_model_with_mixed_bit_derived_mla_projection(self):
+        # deepseek_v3's sanitize() derives embed_q/unembed_out from kv_b_proj at
+        # its bits, but those paths never reach config["quantization"], so
+        # load_model has to infer them instead of using the global default.
+        from mlx_lm.models import deepseek_v3
+
+        args = deepseek_v3.ModelArgs(
+            vocab_size=64,
+            hidden_size=32,
+            intermediate_size=64,
+            moe_intermediate_size=64,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            n_routed_experts=2,
+            kv_lora_rank=32,
+            q_lora_rank=32,
+            qk_rope_head_dim=32,
+            v_head_dim=32,
+            qk_nope_head_dim=32,
+        )
+        model = deepseek_v3.Model(args)
+        # The derived projections use a different mode than the rest of the
+        # model, so load_model has to infer the mode too, not just the bits.
+        group_size, bits, mode = 32, 4, "mxfp4"
+        derived_bits, derived_mode = 8, "mxfp8"
+        model, config = utils.quantize_model(
+            model,
+            dataclasses.asdict(args),
+            group_size=group_size,
+            bits=bits,
+            mode=mode,
+        )
+
+        # Re-pack the derived projections, as a per-tensor override on kv_b_proj
+        # would, leaving config["quantization"] untouched.
+        attn = model.layers[0].self_attn
+        for proj in (attn.embed_q, attn.unembed_out):
+            w = mx.dequantize(
+                proj.weight,
+                proj.scales,
+                group_size=proj.group_size,
+                bits=proj.bits,
+                mode=proj.mode,
+            )
+            proj.weight, proj.scales = mx.quantize(
+                w, group_size=group_size, bits=derived_bits, mode=derived_mode
+            )
+
+        with tempfile.TemporaryDirectory(dir=self.test_dir) as mlx_path:
+            utils.save_model(mlx_path, model)
+            utils.save_config(config, os.path.join(mlx_path, "config.json"))
+            self.assertEqual(
+                config["quantization"],
+                {"group_size": group_size, "bits": bits, "mode": mode},
+            )
+
+            loaded, _ = utils.load_model(Path(mlx_path))
+
+            loaded_attn = loaded.layers[0].self_attn
+            for proj in (loaded_attn.embed_q, loaded_attn.unembed_out):
+                self.assertEqual(
+                    (proj.bits, proj.group_size, proj.mode),
+                    (derived_bits, group_size, derived_mode),
+                )
+
 
 CUSTOM_MODEL_FILE = """\
 from pathlib import Path