diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 9c3104791..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, @@ -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, diff --git a/tests/test_utils.py b/tests/test_utils.py index 55a434c37..3560f5161 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 @@ -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