diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index 0e228bb000..e9d1a01cf3 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -900,7 +900,11 @@ def process_cache(input_name, name_suffix): return flat_cos, flat_sin def load_weights(self, input_path): - # Load the Hugging Face model + # For quantized models (e.g., Quark, AWQ, GPTQ) or GGUF, use base class logic + # which loads weights directly via QuantModel + if self.quant_type is not None or input_path.endswith(".gguf"): + return super().load_weights(input_path) + print("Loading Qwen3VLForConditionalGeneration model...") return Qwen3VLForConditionalGeneration.from_pretrained( self.model_name_or_path, diff --git a/src/python/py/models/quantized_model.py b/src/python/py/models/quantized_model.py index 1da0445309..b06b060961 100644 --- a/src/python/py/models/quantized_model.py +++ b/src/python/py/models/quantized_model.py @@ -20,6 +20,24 @@ from safetensors.torch import load_file +def normalize_vlm_weight_name(name): + """Normalize a checkpoint tensor key for VLM/Quark conventions. + + Returns None if the tensor should be skipped (vision-tower weights), or + the normalized key string otherwise. + """ + # Skip vision tower weights in VLM checkpoints + if name.startswith(("model.visual.", "model.vision.", "visual.")): + return None + # Normalize common VLM prefix so existing LLM regex + parsing keeps working + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + # Normalize Quark weight_quantizer.* naming to flat weight_* naming + name = name.replace(".weight_quantizer.scale", ".weight_scale") + name = name.replace(".weight_quantizer.zero_point", ".weight_zero_point") + return name + + class QuantizedTensorModule: def __init__(self): self.qweight = None @@ -217,7 +235,11 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme weights = load_file(os.path.join(input_path, weight_file)) # Map weights to modules - for name, tensor in weights.items(): + for raw_name, tensor in weights.items(): + name = normalize_vlm_weight_name(raw_name) + if name is None: + continue + # Per-layer quantization support local_bits = self.get_layer_bits(name) # codeql[py/init-calls-subclass] local_group_size = self.get_layer_group_size(name) # codeql[py/init-calls-subclass] diff --git a/test/python/test_quantized_model.py b/test/python/test_quantized_model.py index 52dac77e8c..9e6bafec32 100644 --- a/test/python/test_quantized_model.py +++ b/test/python/test_quantized_model.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License -"""Unit tests for quantized_model.py lm_head tensor loading. +"""Unit tests for quantized_model.py lm_head tensor loading and VLM key normalisation. These tests verify that lm_head tensors are assigned correctly regardless -of the iteration order returned by safetensors.torch.load_file(). +of the iteration order returned by safetensors.torch.load_file(), and that +the VLM/Quark checkpoint key normalisation introduced for Qwen3-VL-4B works +correctly so future refactors do not silently break quantised VLM loading. """ from __future__ import annotations @@ -14,6 +16,7 @@ QuantizedModel, QuantizedTensorModule, TensorModule, + normalize_vlm_weight_name, ) @@ -203,3 +206,61 @@ def test_lm_head_bias_assigned(): assert isinstance(model.lm_head, TensorModule) assert model.lm_head.bias is bias + + +# --------------------------------------------------------------------------- +# Regression tests for VLM / Quark checkpoint key normalisation (Qwen3-VL-4B) +# --------------------------------------------------------------------------- + + +def test_normalize_vlm_weight_name_skips_vision_keys(): + """Vision-tower tensors must be filtered out (return None).""" + assert normalize_vlm_weight_name("model.visual.patch_embed.weight") is None + assert normalize_vlm_weight_name("model.vision.encoder.layer.0.weight") is None + assert normalize_vlm_weight_name("visual.embed.weight") is None + + +def test_normalize_vlm_weight_name_keeps_non_vision_keys(): + """Non-vision keys that do not match any normalisation rule pass through unchanged.""" + assert normalize_vlm_weight_name("model.embed_tokens.weight") == "model.embed_tokens.weight" + assert normalize_vlm_weight_name("lm_head.weight") == "lm_head.weight" + assert normalize_vlm_weight_name("model.norm.weight") == "model.norm.weight" + + +def test_normalize_vlm_weight_name_strips_language_model_prefix(): + """'model.language_model.*' must be rewritten to 'model.*'.""" + assert ( + normalize_vlm_weight_name("model.language_model.embed_tokens.weight") + == "model.embed_tokens.weight" + ) + assert ( + normalize_vlm_weight_name("model.language_model.layers.0.self_attn.q_proj.weight") + == "model.layers.0.self_attn.q_proj.weight" + ) + assert ( + normalize_vlm_weight_name("model.language_model.norm.weight") + == "model.norm.weight" + ) + + +def test_normalize_vlm_weight_name_quark_scale_renamed(): + """Quark '.weight_quantizer.scale' must map to '.weight_scale'.""" + assert ( + normalize_vlm_weight_name("model.layers.0.self_attn.q_proj.weight_quantizer.scale") + == "model.layers.0.self_attn.q_proj.weight_scale" + ) + + +def test_normalize_vlm_weight_name_quark_zero_point_renamed(): + """Quark '.weight_quantizer.zero_point' must map to '.weight_zero_point'.""" + assert ( + normalize_vlm_weight_name("model.layers.0.mlp.gate_proj.weight_quantizer.zero_point") + == "model.layers.0.mlp.gate_proj.weight_zero_point" + ) + + +def test_normalize_vlm_weight_name_combined_vlm_prefix_and_quark(): + """VLM prefix stripping and Quark renaming must compose correctly.""" + raw = "model.language_model.layers.2.self_attn.v_proj.weight_quantizer.scale" + expected = "model.layers.2.self_attn.v_proj.weight_scale" + assert normalize_vlm_weight_name(raw) == expected