diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index cdc7bfaba..8297124cc 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/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: """ diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 14f4e5f86..b4bdc41b3 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/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, diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index b3fb7d898..a40e47ab4 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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 = { diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index 9336dbaf8..f2320ec76 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/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, diff --git a/python/sglang/srt/models/qwen2_vl.py b/python/sglang/srt/models/qwen2_vl.py index fa4055c36..24ec4c430 100644 --- a/python/sglang/srt/models/qwen2_vl.py +++ b/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, diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 3953b85cc..922d47632 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/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, diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 20854cc18..a536a55a8 100644 --- a/python/sglang/srt/models/utils.py +++ b/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."""