From 1c9597a0068c7255035c96554257644f317f48cc Mon Sep 17 00:00:00 2001 From: Rajan Sharma Date: Mon, 6 Jul 2026 12:35:25 -0400 Subject: [PATCH 1/4] fix: infer bits/group_size for quantized paths missing from the per-tensor map load_model's class_predicate assumed any module whose weights already carry `.scales` should be quantized at the top-level default group_size/bits. For tensors a model's sanitize() derives and re-quantizes from a source tensor with its own per-tensor override (e.g. the absorbed MLA projections embed_q/unembed_out, derived from kv_b_proj in deepseek_v3, deepseek_v32, glm4_moe_lite, kimi_linear, and longcat_flash), that path never appears in config["quantization"] because the converter never saw it - it's created inside sanitize(). Loading a mixed-bit checkpoint where kv_b_proj carries a per-tensor override then allocates the derived module at the global default bits while the saved weight is packed at the override bits, raising a shape mismatch on load_weights. Infer the actual group_size/bits from the packed weight and scales shapes instead, mirroring the inference each sanitize() already performs on the source tensor. This is a no-op for the common case (no override, inferred bits equal the default already used to produce the checkpoint). Fixes #1451. --- mlx_lm/utils.py | 20 ++++++++++- tests/test_utils.py | 83 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 7b285063e..0b03491f9 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -366,7 +366,25 @@ 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 + # This path has no entry in the per-tensor quantization map but + # the checkpoint already carries a quantized weight for it (e.g. + # a model's sanitize() can derive and re-quantize tensors, such + # as the absorbed MLA projections in deepseek_v3/deepseek_v32/ + # glm4_moe_lite/kimi_linear/longcat_flash, from a source tensor + # that has its own per-tensor bits/group_size override). Infer + # the actual packing from the weight/scales shapes instead of + # assuming the top-level default, so the module nn.quantize + # allocates matches the weight that will be loaded into it. + in_dims = m.weight.shape[-1] + group_size = in_dims // weights[f"{p}.scales"].shape[-1] + bits = (weights[f"{p}.weight"].shape[-1] * 32) // in_dims + return { + "group_size": group_size, + "bits": bits, + "mode": quantization.get("mode", "affine"), + } nn.quantize( model, diff --git a/tests/test_utils.py b/tests/test_utils.py index 21f60da46..5d97b3266 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,5 +1,6 @@ # Copyright © 2024 Apple Inc. +import dataclasses import json import os import tempfile @@ -183,6 +184,88 @@ 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_load_model_with_mixed_bit_derived_mla_projection(self): + # Regression test for a mismatch that only shows up with a + # non-uniform per-tensor quantization map: deepseek_v3's sanitize() + # derives embed_q/unembed_out from kv_b_proj and re-quantizes them + # at kv_b_proj's own bits, but those derived paths never appear in + # the per-tensor quantization map (the converter never saw them), + # so load_model must infer their true bits from the saved weight + # shapes rather than assume the global default. + from mlx_lm.models import deepseek_v3 + + args = deepseek_v3.ModelArgs( + model_type="deepseek_v3", + 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, + n_group=1, + topk_group=1, + num_experts_per_tok=2, + n_shared_experts=1, + kv_lora_rank=32, + q_lora_rank=32, + qk_rope_head_dim=32, + v_head_dim=32, + qk_nope_head_dim=32, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + ) + model = deepseek_v3.Model(args) + model, config = utils.quantize_model( + model, dataclasses.asdict(args), group_size=32, bits=4 + ) + + # Simulate a checkpoint where kv_b_proj carried a per-tensor bits + # override: re-pack the already-derived embed_q/unembed_out at a + # different bit width than the model's global default, without + # adding their paths to config["quantization"]. + attn = model.layers[0].self_attn + for proj in (attn.embed_q, attn.unembed_out): + w = mx.dequantize( + proj.weight, + proj.scales, + proj.biases, + group_size=proj.group_size, + bits=proj.bits, + mode=proj.mode, + ) + proj.weight, proj.scales, proj.biases = mx.quantize( + w, group_size=32, bits=8, mode=proj.mode + ) + proj.group_size, proj.bits = 32, 8 + + 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.assertFalse( + any( + "embed_q" in k or "unembed_out" in k for k in config["quantization"] + ) + ) + + loaded, _ = utils.load_model(Path(mlx_path)) + + loaded_attn = loaded.layers[0].self_attn + self.assertEqual(loaded_attn.embed_q.bits, 8) + self.assertEqual(loaded_attn.unembed_out.bits, 8) + + logits = loaded(mx.array([[1, 2, 3]], dtype=mx.int32)) + mx.eval(logits) + self.assertEqual(logits.shape, (1, 3, args.vocab_size)) + CUSTOM_MODEL_FILE = """\ from pathlib import Path From d49377bc46456af4f07741c6adda7605dc891bd1 Mon Sep 17 00:00:00 2001 From: Michal Klein <46717574+michalk8@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:37:34 +0200 Subject: [PATCH 2/4] Add `infer_quant_config` --- mlx_lm/utils.py | 45 ++++++++++++++++++++++++++++----------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index bf1f2d3e6..a00fed8ab 100644 --- a/mlx_lm/utils.py +++ b/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, @@ -465,23 +492,7 @@ def class_predicate(p, m): return False if f"{p}.scales" not in weights: return False - # This path has no entry in the per-tensor quantization map but - # the checkpoint already carries a quantized weight for it (e.g. - # a model's sanitize() can derive and re-quantize tensors, such - # as the absorbed MLA projections in deepseek_v3/deepseek_v32/ - # glm4_moe_lite/kimi_linear/longcat_flash, from a source tensor - # that has its own per-tensor bits/group_size override). Infer - # the actual packing from the weight/scales shapes instead of - # assuming the top-level default, so the module nn.quantize - # allocates matches the weight that will be loaded into it. - in_dims = m.weight.shape[-1] - group_size = in_dims // weights[f"{p}.scales"].shape[-1] - bits = (weights[f"{p}.weight"].shape[-1] * 32) // in_dims - return { - "group_size": group_size, - "bits": bits, - "mode": quantization.get("mode", "affine"), - } + return infer_quant_config(p, m, weights) nn.quantize( model, From 7a6ed5cfd1010ac784f4d10f6ddbc5846c69af9f Mon Sep 17 00:00:00 2001 From: Michal Klein <46717574+michalk8@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:46:44 +0200 Subject: [PATCH 3/4] Add tests --- tests/test_utils.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_utils.py b/tests/test_utils.py index f7f37ce03..738be67fd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -255,6 +255,42 @@ 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): # Regression test for a mismatch that only shows up with a # non-uniform per-tensor quantization map: deepseek_v3's sanitize() From eb49c3c375110591a3e87b580a5dde64db48175d Mon Sep 17 00:00:00 2001 From: Michal Klein <46717574+michalk8@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:28:51 +0200 Subject: [PATCH 4/4] Update test --- tests/test_utils.py | 64 +++++++++++++++++---------------------------- 1 file changed, 24 insertions(+), 40 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 738be67fd..3560f5161 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -292,17 +292,12 @@ def test_infer_quant_config_unknown_packing(self): utils.infer_quant_config("l", layer, weights) def test_load_model_with_mixed_bit_derived_mla_projection(self): - # Regression test for a mismatch that only shows up with a - # non-uniform per-tensor quantization map: deepseek_v3's sanitize() - # derives embed_q/unembed_out from kv_b_proj and re-quantizes them - # at kv_b_proj's own bits, but those derived paths never appear in - # the per-tensor quantization map (the converter never saw them), - # so load_model must infer their true bits from the saved weight - # shapes rather than assume the global default. + # 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( - model_type="deepseek_v3", vocab_size=64, hidden_size=32, intermediate_size=64, @@ -311,67 +306,56 @@ def test_load_model_with_mixed_bit_derived_mla_projection(self): num_attention_heads=2, num_key_value_heads=2, n_routed_experts=2, - n_group=1, - topk_group=1, - num_experts_per_tok=2, - n_shared_experts=1, kv_lora_rank=32, q_lora_rank=32, qk_rope_head_dim=32, v_head_dim=32, qk_nope_head_dim=32, - rope_scaling={ - "beta_fast": 32, - "beta_slow": 1, - "factor": 40, - "mscale": 1.0, - "mscale_all_dim": 1.0, - "original_max_position_embeddings": 4096, - "type": "yarn", - }, ) 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=32, bits=4 + model, + dataclasses.asdict(args), + group_size=group_size, + bits=bits, + mode=mode, ) - # Simulate a checkpoint where kv_b_proj carried a per-tensor bits - # override: re-pack the already-derived embed_q/unembed_out at a - # different bit width than the model's global default, without - # adding their paths to config["quantization"]. + # 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, - proj.biases, group_size=proj.group_size, bits=proj.bits, mode=proj.mode, ) - proj.weight, proj.scales, proj.biases = mx.quantize( - w, group_size=32, bits=8, mode=proj.mode + proj.weight, proj.scales = mx.quantize( + w, group_size=group_size, bits=derived_bits, mode=derived_mode ) - proj.group_size, proj.bits = 32, 8 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.assertFalse( - any( - "embed_q" in k or "unembed_out" in k for k in config["quantization"] - ) + 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 - self.assertEqual(loaded_attn.embed_q.bits, 8) - self.assertEqual(loaded_attn.unembed_out.bits, 8) - - logits = loaded(mx.array([[1, 2, 3]], dtype=mx.int32)) - mx.eval(logits) - self.assertEqual(logits.shape, (1, 3, args.vocab_size)) + 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 = """\