fix qwenvl compressed tensors quantization weight loader (#11914)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: Peng Zhang <aniz1905@gmail.com>

python/sglang/srt/layers/quantization/base_config.py

@@ -11,6 +11,7 @@ from torch import nn
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
     from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
+    from sglang.srt.models.utils import WeightsMapper
 
 
 class QuantizeMethodBase(ABC):
@@ -229,6 +230,15 @@ class QuantizationConfig(ABC):
         """
         raise NotImplementedError()
 
+    def apply_sglang_mapper(self, hf_to_sglang_mapper: "WeightsMapper"):  # noqa: B027
+        """
+        Interface for models to update module names referenced in
+        quantization configs in order to reflect the sglang model structure
+        :param hf_to_sglang_mapper: maps from hf model structure (the assumed
+            structure of the qconfig) to sglang model structure
+        """
+        pass
+
 
 def method_has_implemented_embedding(method_class: Type[QuantizeMethodBase]) -> bool:
     """

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -4,7 +4,17 @@ from __future__ import annotations
 
 import logging
 from contextlib import suppress
-from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    NamedTuple,
+    Optional,
+    Tuple,
+    cast,
+)
 
 import torch
 from compressed_tensors.config import (
@@ -44,6 +54,9 @@ from sglang.srt.layers.quantization.compressed_tensors.utils import (
 from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 
+if TYPE_CHECKING:
+    from sglang.srt.models.utils import WeightsMapper
+
 logger = logging.getLogger(__name__)
 
 __all__ = ["CompressedTensorsLinearMethod"]
@@ -110,6 +123,18 @@ class CompressedTensorsConfig(QuantizationConfig):
     def get_scaled_act_names(self) -> List[str]:
         return []
 
+    def apply_sglang_mapper(self, hf_to_sglang_mapper: "WeightsMapper"):
+        self.target_scheme_map = hf_to_sglang_mapper.apply_dict(self.target_scheme_map)
+        self.ignore = hf_to_sglang_mapper.apply_list(self.ignore)
+        self.sparsity_scheme_map = hf_to_sglang_mapper.apply_dict(
+            self.sparsity_scheme_map
+        )
+        self.sparsity_ignore_list = hf_to_sglang_mapper.apply_list(
+            self.sparsity_ignore_list
+        )
+        if self.kv_cache_scheme is not None:
+            self.kv_cache_scheme = hf_to_sglang_mapper.apply_dict(self.kv_cache_scheme)
+
     def get_quant_method(
         self,
         layer: torch.nn.Module,

python/sglang/srt/model_loader/loader.py

@@ -257,6 +257,10 @@ def _initialize_model(
     quant_config = _get_quantization_config(
         model_config, load_config, packed_modules_mapping, remap_prefix
     )
+    hf_to_sglang_mapper = getattr(model_class, "hf_to_sglang_mapper", None)
+    # pass mappings by reference to quant_config
+    if hf_to_sglang_mapper is not None and quant_config is not None:
+        quant_config.apply_sglang_mapper(hf_to_sglang_mapper)
 
     # Build kwargs conditionally
     kwargs = {

python/sglang/srt/models/qwen2_5_vl.py

@@ -70,7 +70,7 @@ from sglang.srt.managers.schedule_batch import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen2 import Qwen2Model
-from sglang.srt.models.utils import RotaryPosMixin, permute_inv
+from sglang.srt.models.utils import RotaryPosMixin, WeightsMapper, permute_inv
 from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
 from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
 from sglang.srt.server_args import get_global_server_args
@@ -547,6 +547,24 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
         "up_proj": ("gate_up_proj", 1),
     }
 
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+    # To ensure correct weight loading and mapping.
+    hf_to_sglang_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "attn.qkv": "attn.qkv_proj",
+        },
+        orig_to_new_prefix={
+            # mapping for new names in checkpoint saved after transformers v4.52
+            "model.language_model.": "language_model.model.",
+            "model.visual.": "visual.",
+            # mapping for original checkpoint
+            "lm_head.": "language_model.lm_head.",
+            "model.": "language_model.model.",
+        },
+    )
+
     def __init__(
         self,
         config: Qwen2_5_VLConfig,

python/sglang/srt/models/qwen2_vl.py

@@ -47,7 +47,7 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInp
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen2 import Qwen2Model
-from sglang.srt.models.utils import compute_cu_seqlens_from_grid_numpy
+from sglang.srt.models.utils import WeightsMapper, compute_cu_seqlens_from_grid_numpy
 from sglang.srt.utils import add_prefix
 from sglang.srt.utils.hf_transformers_utils import get_processor
 
@@ -426,6 +426,21 @@ class Qwen2VLForConditionalGeneration(nn.Module):
         "up_proj": ("gate_up_proj", 1),
     }
 
+    # To ensure correct weight loading and mapping.
+    hf_to_sglang_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "attn.qkv": "attn.qkv_proj",
+        },
+        orig_to_new_prefix={
+            # mapping for new names in checkpoint saved after transformers v4.52
+            "model.language_model.": "language_model.model.",
+            "model.visual.": "visual.",
+            # mapping for original checkpoint
+            "lm_head.": "language_model.lm_head.",
+            "model.": "language_model.model.",
+        },
+    )
+
     def __init__(
         self,
         config: Qwen2VLConfig,

python/sglang/srt/models/qwen3_vl.py

@@ -53,7 +53,11 @@ from sglang.srt.managers.schedule_batch import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen3 import Qwen3Model
-from sglang.srt.models.utils import RotaryPosMixin, compute_cu_seqlens_from_grid_numpy
+from sglang.srt.models.utils import (
+    RotaryPosMixin,
+    WeightsMapper,
+    compute_cu_seqlens_from_grid_numpy,
+)
 from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, get_int_env_var
@@ -594,6 +598,21 @@ class Qwen3LLMModel(Qwen3Model):
 
 
 class Qwen3VLForConditionalGeneration(nn.Module):
+    # To ensure correct weight loading and mapping.
+    hf_to_sglang_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "attn.qkv": "attn.qkv_proj",
+        },
+        orig_to_new_prefix={
+            # mapping for new names in checkpoint saved after transformers v4.52
+            "model.language_model.": "language_model.model.",
+            "model.visual.": "visual.",
+            # mapping for original checkpoint
+            "lm_head.": "language_model.lm_head.",
+            "model.": "language_model.model.",
+        },
+    )
+
     def __init__(
         self,
         config: Qwen3VLConfig,

python/sglang/srt/models/utils.py

@@ -12,7 +12,10 @@
 # limitations under the License.
 # ==============================================================================
 
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass, field
 from functools import lru_cache
+from typing import Any, Optional
 
 import numpy as np
 import torch
@@ -27,6 +30,74 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import FusedSetKVBufferArg
 
+WeightsMapping = Mapping[str, Optional[str]]
+"""If a key maps to a value of `None`, the corresponding weight is ignored."""
+
+
+@dataclass
+class WeightsMapper:
+    """Maps the name of each weight if they match the following patterns."""
+
+    orig_to_new_substr: WeightsMapping = field(default_factory=dict)
+    orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
+    orig_to_new_suffix: WeightsMapping = field(default_factory=dict)
+
+    def _map_name(self, key: str) -> Optional[str]:
+        for substr, new_key in sorted(
+            self.orig_to_new_substr.items(), key=lambda i: len(i[0]), reverse=True
+        ):
+            if substr in key:
+                if new_key is None:
+                    return None
+
+                key = key.replace(substr, new_key, 1)
+                break
+
+        for prefix, new_key in sorted(
+            self.orig_to_new_prefix.items(), key=lambda i: len(i[0]), reverse=True
+        ):
+            if key.startswith(prefix):
+                if new_key is None:
+                    return None
+
+                key = key.replace(prefix, new_key, 1)
+                break
+
+        for suffix, new_key in sorted(
+            self.orig_to_new_suffix.items(), key=lambda i: len(i[0]), reverse=True
+        ):
+            if key.endswith(suffix):
+                if new_key is None:
+                    return None
+
+                key = new_key.join(key.rsplit(suffix, 1))
+                break
+
+        return key
+
+    def apply(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        return (
+            (out_name, data)
+            for name, data in weights
+            if (out_name := self._map_name(name)) is not None
+        )
+
+    def apply_list(self, values: list[str]) -> list[str]:
+        return [
+            out_name
+            for name in values
+            if (out_name := self._map_name(name)) is not None
+        ]
+
+    def apply_dict(self, values: dict[str, Any]) -> dict[str, Any]:
+        return {
+            out_name: value
+            for name, value in values.items()
+            if (out_name := self._map_name(name)) is not None
+        }
+
 
 def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
     """Enable fused set_kv_buffer only on CUDA with bfloat16 KV cache."""