chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix qwenvl compressed tensors quantization weight loader (#11914)

Browse Source 
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: Peng Zhang <aniz1905@gmail.com>
This commit is contained in:
LHXuuu
2025-12-17 18:01:43 +08:00
committed by GitHub
parent 8c34e18140
commit 712f44ee2b
7 changed files with 166 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/layers/quantization/base_config.py
View File

@@ -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:
"""

27
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
View File

@@ -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,

4
python/sglang/srt/model_loader/loader.py
View File

@@ -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 = {

20
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -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,

17
python/sglang/srt/models/qwen2_vl.py
View File

@@ -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,

21
python/sglang/srt/models/qwen3_vl.py
View File

@@ -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,

71
python/sglang/srt/models/utils.py
View File

@@ -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."""