952
vllm_ascend/quantization/modelslim_config.py
Normal file
952
vllm_ascend/quantization/modelslim_config.py
Normal file
@@ -0,0 +1,952 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
"""ModelSlim quantization configuration and model mappings for Ascend.
|
||||
|
||||
This module provides the AscendModelSlimConfig class for parsing quantization
|
||||
configs generated by the ModelSlim tool, along with model-specific mappings.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Optional
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
from vllm.config import get_current_vllm_config, get_current_vllm_config_or_none
|
||||
from vllm.logger import logger
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
from vllm.model_executor.layers.quantization import register_quantization_config
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import UnquantizedEmbeddingMethod, VocabParallelEmbedding
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
from vllm_ascend.utils import (
|
||||
ASCEND_QUANTIZATION_METHOD,
|
||||
AscendDeviceType,
|
||||
calc_split_factor,
|
||||
get_ascend_device_type,
|
||||
vllm_version_is,
|
||||
)
|
||||
|
||||
if vllm_version_is("0.23.0"):
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoE
|
||||
else:
|
||||
from vllm.model_executor.layers.fused_moe import MoERunner, RoutedExperts
|
||||
|
||||
from .methods import get_scheme_class
|
||||
|
||||
|
||||
def _is_fused_moe_layer(layer: torch.nn.Module) -> bool:
|
||||
if vllm_version_is("0.23.0"):
|
||||
return isinstance(layer, FusedMoE)
|
||||
else:
|
||||
return isinstance(layer, (MoERunner, RoutedExperts))
|
||||
|
||||
|
||||
# The config filename that ModelSlim generates after quantizing a model.
|
||||
MODELSLIM_CONFIG_FILENAME = "quant_model_description.json"
|
||||
|
||||
# key: model_type
|
||||
# value: dict of fused module name -> list of original module names
|
||||
packed_modules_model_mapping: dict[str, dict[str, list[str]]] = {
|
||||
"qwen3_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"qwen3_5": {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
},
|
||||
"qwen3_5_moe": {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"deepseek_v2": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"deepseek_v3": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"deepseek_v4": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"pangu_ultra_moe": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"kimi_k2": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"deepseek_v32": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"glm_moe_dsa": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
# NOTE 1.The quantized MTP layer of deepseek on the NPU is not quantized;
|
||||
# NOTE 2.The description file generated by the current msmodelslim tool does not have
|
||||
# MTP layer info. Please manually add it and set the value to FLOAT.
|
||||
"deepseek_mtp": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"pangu_ultra_moe_mtp": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"qwen3_next": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"in_proj": ["in_proj_qkvz", "in_proj_ba"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"qwen2_5_vl": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
},
|
||||
"qwen3_vl_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"glm4_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"glm4_moe_lite": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"glm4v_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"glm4v_moe_text": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"longcat_flash": {
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
},
|
||||
"minimax_m2": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"experts": ["experts.0.w1", "experts.0.w2", "experts.0.w3"],
|
||||
},
|
||||
"qwen3_omni_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"attn_qkv_proj": [
|
||||
"attn_q_proj",
|
||||
"attn_k_proj",
|
||||
"attn_v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
"qwen2_5_omni": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"attn_qkv_proj": [
|
||||
"attn_q_proj",
|
||||
"attn_k_proj",
|
||||
"attn_v_proj",
|
||||
],
|
||||
"qkv": [
|
||||
"q",
|
||||
"k",
|
||||
"v",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
},
|
||||
"bailing_hybrid": {
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
||||
"o_proj": ["dense"],
|
||||
},
|
||||
"step3p5": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
# The step3.5 MTP draft (speculative.py sets model_type="step3p5_mtp")
|
||||
# reuses the same fused module layout as the verifier.
|
||||
"step3p5_mtp": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts": ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
QUANT_MODEL_PREFIX_MAPPINGS = {
|
||||
"deepseek_v4": {
|
||||
"layers.": "model.layers.",
|
||||
"embed.": "model.embed_tokens.",
|
||||
"head.": "lm_head.",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
QUANT_MODEL_SUBSTR_MAPPINGS = {
|
||||
"deepseek_v4": {
|
||||
".attn.": ".self_attn.",
|
||||
".w1.": ".gate_proj.",
|
||||
".w2.": ".down_proj.",
|
||||
".w3.": ".up_proj.",
|
||||
".ffn.": ".mlp.",
|
||||
".ffn_norm.": ".post_attention_layernorm.",
|
||||
".attn_norm.": ".input_layernorm.",
|
||||
},
|
||||
# The step3.5 MTP draft nests its decoder block under ".mtp_block.", but the
|
||||
# checkpoint's quant_model_description.json keys it without that infix
|
||||
# (e.g. "model.layers.45.self_attn.q_proj.weight"). Strip it so the quant
|
||||
# lookup matches the on-disk naming.
|
||||
"step3p5_mtp": {
|
||||
".mtp_block.": ".",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_packed_modules_mapping(model_type: str) -> dict[str, list[str]]:
|
||||
"""Get packed modules mapping for a model type.
|
||||
|
||||
Args:
|
||||
model_type: The model type string (e.g., "deepseek_v3").
|
||||
|
||||
Returns:
|
||||
Dictionary mapping fused module names to their component module names.
|
||||
Returns empty dict if model_type is not found.
|
||||
"""
|
||||
return packed_modules_model_mapping.get(model_type, {})
|
||||
|
||||
|
||||
def get_linear_quant_type(
|
||||
quant_description: dict[str, Any], prefix: str, packed_modules_mapping: dict[str, Any]
|
||||
) -> str | None:
|
||||
"""Determine the quantization type for a linear layer.
|
||||
|
||||
Args:
|
||||
quant_description: The quantization description dictionary.
|
||||
prefix: The layer prefix.
|
||||
packed_modules_mapping: Mapping for packed/fused modules.
|
||||
|
||||
Returns:
|
||||
The quantization type string (e.g., "W8A8_DYNAMIC").
|
||||
"""
|
||||
proj_name = prefix.split(".")[-1]
|
||||
if proj_name in packed_modules_mapping:
|
||||
quant_type = None
|
||||
shard_prefixes = [
|
||||
prefix.replace(proj_name, shard_proj_name) for shard_proj_name in packed_modules_mapping[proj_name]
|
||||
]
|
||||
for shard_prefix in shard_prefixes:
|
||||
shard_quant_type = quant_description[shard_prefix + ".weight"]
|
||||
|
||||
if quant_type is None:
|
||||
quant_type = shard_quant_type
|
||||
elif shard_quant_type != quant_type:
|
||||
err_msg = (
|
||||
f"Not all shards of {prefix} are quantized with same quant type. "
|
||||
f"Shard {proj_name} uses {shard_quant_type}, but another shard "
|
||||
f"uses {quant_type}. Please check quantization config."
|
||||
)
|
||||
logger.error(err_msg)
|
||||
raise ValueError(err_msg)
|
||||
else:
|
||||
quant_type = quant_description[prefix + ".weight"]
|
||||
return quant_type
|
||||
|
||||
|
||||
def get_quant_type_for_layer(
|
||||
quant_description: dict[str, Any],
|
||||
prefix: str,
|
||||
layer_type: str,
|
||||
packed_modules_mapping: dict[str, Any] | None = None,
|
||||
) -> str | None:
|
||||
"""Determine the quantization type for a layer.
|
||||
|
||||
Args:
|
||||
quant_description: The quantization description dictionary.
|
||||
prefix: The layer prefix.
|
||||
layer_type: The type of layer ("linear", "moe", "attention").
|
||||
packed_modules_mapping: Mapping for packed/fused modules.
|
||||
|
||||
Returns:
|
||||
The quantization type string (e.g., "W8A8_DYNAMIC").
|
||||
"""
|
||||
if packed_modules_mapping is None:
|
||||
packed_modules_mapping = dict()
|
||||
# Attention
|
||||
if layer_type == "attention":
|
||||
layer_indexer_quant_type = quant_description.get(f"{prefix}.indexer.quant_type")
|
||||
if layer_indexer_quant_type is not None:
|
||||
return layer_indexer_quant_type
|
||||
if "fa_quant_type" in quant_description:
|
||||
return quant_description["fa_quant_type"]
|
||||
if "indexer_quant_type" in quant_description:
|
||||
return quant_description["indexer_quant_type"]
|
||||
# Linear / MoE
|
||||
return get_linear_quant_type(quant_description, prefix, packed_modules_mapping)
|
||||
|
||||
|
||||
def create_scheme_for_layer(
|
||||
quant_description: dict[str, Any],
|
||||
prefix: str,
|
||||
layer_type: str,
|
||||
packed_modules_mapping: dict[str, Any] | None = None,
|
||||
):
|
||||
"""Create a quantization scheme instance for a layer.
|
||||
|
||||
Args:
|
||||
quant_description: The quantization description dictionary.
|
||||
prefix: The layer prefix.
|
||||
layer_type: The type of layer ("linear", "moe", "attention").
|
||||
packed_modules_mapping: Mapping for packed/fused modules.
|
||||
|
||||
Returns:
|
||||
An instance of the appropriate quantization scheme class.
|
||||
"""
|
||||
logger.info_once("Using the vLLM Ascend modelslim Quantization now!")
|
||||
quant_type = get_quant_type_for_layer(quant_description, prefix, layer_type, packed_modules_mapping)
|
||||
|
||||
if quant_type is None:
|
||||
err_msg = f"Could not determine quantization type for layer {prefix} (layer_type={layer_type})."
|
||||
logger.error(err_msg)
|
||||
raise ValueError(err_msg)
|
||||
|
||||
# Use registry to get scheme class
|
||||
scheme_cls = get_scheme_class(quant_type, layer_type)
|
||||
if scheme_cls is not None:
|
||||
return scheme_cls()
|
||||
|
||||
err_msg = (
|
||||
"Currently, vLLM Ascend doesn't support quant_type=%s for layer_type=%s. "
|
||||
"Please use a supported quantization format "
|
||||
"or load the model with its original float weights."
|
||||
)
|
||||
logger.error(err_msg, quant_type, layer_type)
|
||||
raise NotImplementedError(err_msg % (quant_type, layer_type))
|
||||
|
||||
|
||||
@register_quantization_config(ASCEND_QUANTIZATION_METHOD)
|
||||
class AscendModelSlimConfig(QuantizationConfig):
|
||||
"""Config class for Ascend ModelSlim quantization.
|
||||
|
||||
This class is a general class that parses quantization configs
|
||||
that are supported on Ascend hardware, specifically for models
|
||||
quantized using the ModelSlim tool.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: dict[str, Any] | None = None):
|
||||
super().__init__()
|
||||
self.quant_description = quant_config if quant_config is not None else {}
|
||||
self._apply_extra_quant_adaptations()
|
||||
self.model_type: str | None = None
|
||||
self.hf_to_vllm_mapper: WeightsMapper | None = None
|
||||
self._mapper_applied = False
|
||||
self._add_kvcache_quant_metadata()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "AscendModelSlimConfig:\n" + super().__repr__()
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return ASCEND_QUANTIZATION_METHOD
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.int8, torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
logger.error("Ascend hardware does not support 'get_min_capability' feature.")
|
||||
raise NotImplementedError('Ascend hardware dose not support "get_min_capability" feature.')
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
# Return empty list so that vllm's get_quant_config() skips the
|
||||
# file-based lookup (which raises an unfriendly "Cannot find the
|
||||
# config file for ascend" error when the model is not quantized).
|
||||
# Instead, the config file is loaded in maybe_update_config(),
|
||||
# which can provide a user-friendly error message.
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "AscendModelSlimConfig":
|
||||
return cls(config)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config: Any = None) -> str | None:
|
||||
if hf_quant_cfg is not None:
|
||||
quant_method = hf_quant_cfg.get("quant_method", None)
|
||||
if not quant_method and torch.npu.is_available():
|
||||
return ASCEND_QUANTIZATION_METHOD
|
||||
return None
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
||||
"""Apply the vLLM model-specific mapper to this quantization config.
|
||||
|
||||
This method is called by vLLM to apply the model-specific weight mapper
|
||||
to the quantization configuration. It directly uses the forward mapping
|
||||
(HF -> vLLM) to transform keys in quant_description from HF format to
|
||||
vLLM format.
|
||||
|
||||
Args:
|
||||
hf_to_vllm_mapper: The WeightsMapper instance provided by vLLM
|
||||
that contains model-specific prefix mappings (HF to vLLM).
|
||||
"""
|
||||
if self._mapper_applied and self.hf_to_vllm_mapper is hf_to_vllm_mapper:
|
||||
return
|
||||
vllm_config = get_current_vllm_config_or_none()
|
||||
model_type = None
|
||||
if vllm_config is not None:
|
||||
model_type = vllm_config.model_config.hf_config.model_type
|
||||
|
||||
if model_type == "qwen3_omni_moe":
|
||||
hf_to_vllm_mapper.orig_to_new_prefix = {
|
||||
**hf_to_vllm_mapper.orig_to_new_prefix,
|
||||
"model.": "language_model.model.",
|
||||
"lm_head.": "language_model.lm_head.",
|
||||
}
|
||||
|
||||
self.hf_to_vllm_mapper = hf_to_vllm_mapper
|
||||
self._mapper_applied = True
|
||||
|
||||
if self.quant_description:
|
||||
self.quant_description = hf_to_vllm_mapper.apply_dict(self.quant_description)
|
||||
self._add_kvcache_quant_metadata()
|
||||
logger.info("Applied hf_to_vllm_mapper to quant_description keys")
|
||||
|
||||
def get_cache_scale(self, name: str) -> str | None:
|
||||
"""Map checkpoint C8 KV scale/offset names to vLLM parameter names."""
|
||||
if self.quant_description.get("kv_cache_type") != "C8":
|
||||
return None
|
||||
_C8_SCALE_MAPPING = {
|
||||
"k_proj.kv_cache_scale": "attn.k_cache_scale",
|
||||
"k_proj.kv_cache_offset": "attn.k_cache_offset",
|
||||
"v_proj.kv_cache_scale": "attn.v_cache_scale",
|
||||
"v_proj.kv_cache_offset": "attn.v_cache_offset",
|
||||
}
|
||||
for src_suffix, dst_suffix in _C8_SCALE_MAPPING.items():
|
||||
if name.endswith(src_suffix):
|
||||
return name[: -len(src_suffix)] + dst_suffix
|
||||
return None
|
||||
|
||||
def _has_quant_weight(self, prefix: str, packed_modules_mapping: Mapping[str, list[str]]) -> bool:
|
||||
proj_name = prefix.split(".")[-1]
|
||||
if proj_name in packed_modules_mapping:
|
||||
return all(
|
||||
f"{prefix.replace(proj_name, shard_proj_name)}.weight" in self.quant_description
|
||||
for shard_proj_name in packed_modules_mapping[proj_name]
|
||||
)
|
||||
return f"{prefix}.weight" in self.quant_description
|
||||
|
||||
def quant_prefix_mapper(self, model_type: str, prefix: str) -> str:
|
||||
self.model_type = model_type
|
||||
# Some model paths, e.g. qwen3-vl and qwen3_5_moe MTP drafter,
|
||||
# initialize lm_head with prefix="lm_head", while the quant description
|
||||
# key is mapped to "language_model.lm_head.weight".
|
||||
if (
|
||||
prefix == "lm_head"
|
||||
and "lm_head.weight" not in self.quant_description
|
||||
and "language_model.lm_head.weight" in self.quant_description
|
||||
):
|
||||
prefix = "language_model.lm_head"
|
||||
prefix_mapping = QUANT_MODEL_PREFIX_MAPPINGS.get(model_type)
|
||||
substr_mapping = QUANT_MODEL_SUBSTR_MAPPINGS.get(model_type)
|
||||
if prefix_mapping or substr_mapping:
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix=prefix_mapping or {},
|
||||
orig_to_new_substr=substr_mapping or {},
|
||||
)
|
||||
prefix = hf_to_vllm_mapper._map_name(prefix)
|
||||
|
||||
if model_type == "step3p5_mtp" and prefix.startswith("model.layers."):
|
||||
# Step3P5 MTP and newly generated Step3P7 W8A8 MTP checkpoints use
|
||||
# ``model.layers.*``. The Step3P7 vLLM wrapper mapper rewrites
|
||||
# current ``model.layers.*`` quant descriptions to
|
||||
# ``language_model.model.layers.*``. The MTP draft module itself
|
||||
# is still Step3P5-shaped and queries ``model.layers.*``, so try
|
||||
# the Step3P7 wrapper alias only when the direct Step3P5/new-key
|
||||
# lookup misses.
|
||||
packed_modules_mapping = get_packed_modules_mapping(model_type)
|
||||
if not self._has_quant_weight(prefix, packed_modules_mapping):
|
||||
for candidate in (prefix.replace("model.layers.", "language_model.model.layers.", 1),):
|
||||
if self._has_quant_weight(candidate, packed_modules_mapping):
|
||||
return candidate
|
||||
return prefix
|
||||
|
||||
def get_quant_method(self, layer: torch.nn.Module, prefix: str, tid2eid=None) -> Optional["QuantizeMethodBase"]:
|
||||
from .method_adapters import (
|
||||
AscendEmbeddingMethod,
|
||||
AscendFusedMoEMethod,
|
||||
AscendKVCacheMethod,
|
||||
AscendLinearMethod,
|
||||
)
|
||||
|
||||
vllm_config = get_current_vllm_config()
|
||||
model_type = vllm_config.model_config.hf_config.model_type
|
||||
|
||||
if model_type in ["minimax", "minimax_m2"]:
|
||||
# Adapt to Minimax architecture: update layer names to MoE convention
|
||||
prefix = prefix.replace("mlp", "block_sparse_moe")
|
||||
# Normalize the prefix by stripping specific expert indices (e.g., 'experts.0' -> 'experts')
|
||||
parts = prefix.split(".")
|
||||
if "experts" in parts and len(parts) > 2:
|
||||
exp_idx = parts.index("experts")
|
||||
if exp_idx + 1 < len(parts) and parts[exp_idx + 1].isdigit():
|
||||
parts = parts[: exp_idx + 1]
|
||||
prefix = ".".join(parts)
|
||||
|
||||
if model_type in ["bailing_hybrid"]:
|
||||
# Adapt to bailing_hybrid architecture: update layer names to MoE convention
|
||||
prefix = prefix.replace("linear_attn", "attention")
|
||||
prefix = prefix.replace("self_attn", "attention")
|
||||
if model_type in packed_modules_model_mapping:
|
||||
self.packed_modules_mapping = packed_modules_model_mapping.get(model_type, {})
|
||||
prefix = self.quant_prefix_mapper(model_type, prefix)
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if self.is_layer_skipped_ascend(prefix, self.packed_modules_mapping):
|
||||
# Delayed import to avoid circular import
|
||||
from vllm_ascend.ops.linear import AscendUnquantizedLinearMethod
|
||||
|
||||
logger.debug("Select AscendUnquantizedLinearMethod for %s (layer=%s)", prefix, "LinearBase")
|
||||
return AscendUnquantizedLinearMethod()
|
||||
scheme = create_scheme_for_layer(self.quant_description, prefix, "linear", self.packed_modules_mapping)
|
||||
logger.debug("Select AscendLinearMethod for %s (layer=%s)", prefix, "LinearBase")
|
||||
return AscendLinearMethod(scheme)
|
||||
elif isinstance(layer, AttentionLayerBase) and (
|
||||
self.is_fa_quant_layer(prefix) or self.is_indexer_quant_layer(prefix)
|
||||
):
|
||||
scheme = create_scheme_for_layer(self.quant_description, prefix, "attention", self.packed_modules_mapping)
|
||||
logger.debug("Select AscendKVCacheMethod for %s (layer=%s)", prefix, "AttentionLayerBase[fa/indexer]")
|
||||
return AscendKVCacheMethod(scheme)
|
||||
elif isinstance(layer, AttentionLayerBase) and self.is_c8_quant_layer(prefix):
|
||||
from .methods.kv_c8 import AscendC8KVCacheAttentionMethod
|
||||
|
||||
logger.debug("Select AscendKVCacheMethod(C8) for %s (layer=%s)", prefix, "AttentionLayerBase[C8]")
|
||||
return AscendKVCacheMethod(AscendC8KVCacheAttentionMethod(self.quant_description, prefix))
|
||||
elif _is_fused_moe_layer(layer):
|
||||
if self.is_layer_skipped_ascend(prefix, self.packed_modules_mapping):
|
||||
# Delayed import to avoid circular import
|
||||
from vllm_ascend.ops.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod
|
||||
|
||||
logger.debug("Select AscendUnquantizedFusedMoEMethod for %s (layer=%s)", prefix, "FusedMoE")
|
||||
return AscendUnquantizedFusedMoEMethod(layer.moe_config)
|
||||
scheme = create_scheme_for_layer(self.quant_description, prefix, "moe", self.packed_modules_mapping)
|
||||
logger.debug("Select AscendFusedMoEMethod for %s (layer=%s)", prefix, "FusedMoE")
|
||||
return AscendFusedMoEMethod(scheme, layer.moe_config, tid2eid)
|
||||
elif isinstance(layer, VocabParallelEmbedding):
|
||||
if not self._has_quant_weight(prefix, self.packed_modules_mapping):
|
||||
logger.debug(
|
||||
"No ModelSlim quant entry for %s; select UnquantizedEmbeddingMethod",
|
||||
prefix,
|
||||
)
|
||||
return UnquantizedEmbeddingMethod()
|
||||
if self.is_layer_skipped_ascend(prefix, self.packed_modules_mapping):
|
||||
logger.debug("Select UnquantizedEmbeddingMethod for %s (layer=%s)", prefix, "VocabParallelEmbedding")
|
||||
return UnquantizedEmbeddingMethod()
|
||||
scheme = create_scheme_for_layer(self.quant_description, prefix, "linear", self.packed_modules_mapping)
|
||||
logger.debug("Select AscendEmbeddingMethod for %s (layer=%s)", prefix, "VocabParallelEmbedding")
|
||||
return AscendEmbeddingMethod(scheme)
|
||||
logger.debug("No quant method matched for %s, falling back to base", prefix)
|
||||
return None
|
||||
|
||||
def is_layer_skipped_ascend(self, prefix: str, fused_mapping: Mapping[str, list[str]] = MappingProxyType({})):
|
||||
# adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
|
||||
proj_name = prefix.split(".")[-1]
|
||||
if proj_name in fused_mapping:
|
||||
shard_prefixes = [
|
||||
prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]
|
||||
]
|
||||
|
||||
is_skipped = None
|
||||
for shard_prefix in shard_prefixes:
|
||||
is_shard_skipped = self.quant_description[shard_prefix + ".weight"] == "FLOAT"
|
||||
|
||||
if is_skipped is None:
|
||||
is_skipped = is_shard_skipped
|
||||
elif is_shard_skipped != is_skipped:
|
||||
raise ValueError(
|
||||
f"Detected some but not all shards of {prefix} "
|
||||
"are quantized. All shards of fused layers "
|
||||
"to have the same precision."
|
||||
)
|
||||
else:
|
||||
is_skipped = any(
|
||||
key.startswith(prefix) and key.endswith(".weight") and value == "FLOAT"
|
||||
for key, value in self.quant_description.items()
|
||||
)
|
||||
|
||||
assert is_skipped is not None
|
||||
return is_skipped
|
||||
|
||||
def is_fa_quant_layer(self, prefix):
|
||||
if self.enable_fa_quant:
|
||||
layer_id_str = "".join(re.findall(r"\.(\d+)\.", prefix))
|
||||
if layer_id_str.isdigit() and int(layer_id_str) in self.kvcache_quant_layers:
|
||||
return True
|
||||
return False
|
||||
|
||||
def enabling_fa_quant(self, vllm_config, layer_name) -> bool:
|
||||
is_decode_instance = (
|
||||
vllm_config.kv_transfer_config is not None
|
||||
and vllm_config.kv_transfer_config.is_kv_consumer
|
||||
and not vllm_config.kv_transfer_config.is_kv_producer
|
||||
)
|
||||
if get_ascend_device_type() == AscendDeviceType.A5:
|
||||
return self.is_fa_quant_layer(layer_name)
|
||||
else:
|
||||
return bool(is_decode_instance and self.is_fa_quant_layer(layer_name))
|
||||
|
||||
def is_indexer_quant_layer(self, prefix):
|
||||
if self.enable_indexer_quant:
|
||||
layer_id_str = "".join(re.findall(r"\.(\d+)\.", prefix))
|
||||
if layer_id_str.isdigit() and int(layer_id_str) in self.indexer_quant_layers:
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_c8_quant_layer(self, prefix):
|
||||
if self.enable_c8_quant:
|
||||
layer_id_str = "".join(re.findall(r"\.(\d+)\.", prefix))
|
||||
if layer_id_str.isdigit() and int(layer_id_str) in self.c8_quant_layers:
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_kv_quant_dtype(self, layer_name, cache_dtype, model_config):
|
||||
if self.enable_fa_quant and self.is_fa_quant_layer(layer_name):
|
||||
ori_dtype = model_config.dtype
|
||||
quant_dtype = torch.float8_e4m3fn if get_ascend_device_type() == AscendDeviceType.A5 else torch.int8
|
||||
# For MLA models like deepseek, we only quantify K cache to ensure accuracy
|
||||
if model_config.use_mla:
|
||||
return quant_dtype, ori_dtype
|
||||
else:
|
||||
return quant_dtype, quant_dtype
|
||||
return cache_dtype, cache_dtype
|
||||
|
||||
def get_kv_quant_split_factor(self, layer_name, kv_head_dim_list):
|
||||
if self.enable_fa_quant and self.is_fa_quant_layer(layer_name):
|
||||
k_quant_head_dim = kv_head_dim_list[0]
|
||||
v_quant_head_dim = kv_head_dim_list[1] * 2
|
||||
kv_head_dim_list = [k_quant_head_dim, v_quant_head_dim]
|
||||
return calc_split_factor(kv_head_dim_list)
|
||||
|
||||
def maybe_update_config(
|
||||
self,
|
||||
model_name: str,
|
||||
hf_config: PretrainedConfig | None = None,
|
||||
revision: str | None = None,
|
||||
) -> None:
|
||||
"""Load the ModelSlim quantization config from model directory.
|
||||
|
||||
This method is called by vllm after get_quant_config() returns
|
||||
successfully. Since we return an empty list from get_config_filenames()
|
||||
to bypass vllm's built-in file lookup, we do the actual config loading
|
||||
here and provide user-friendly error messages when the config is missing.
|
||||
|
||||
Works with both local directories (``/path/to/model``) and remote
|
||||
repository identifiers (``org/model-name``). For remote repos the
|
||||
lookup goes through the HuggingFace / ModelScope cache via
|
||||
``get_model_file`` to fetch the config if not already cached.
|
||||
|
||||
Args:
|
||||
model_name: Path to the model directory or HuggingFace /
|
||||
ModelScope repo id.
|
||||
hf_config: The Hugging Face config of the model
|
||||
revision: Optional revision (branch, tag, or commit hash) for
|
||||
remote repos.
|
||||
"""
|
||||
from vllm_ascend.quantization.utils import get_model_file
|
||||
|
||||
# If quant_description is already populated (e.g. from from_config()),
|
||||
# there is nothing to do.
|
||||
if self.quant_description:
|
||||
return
|
||||
|
||||
# Try to get the config file (local or remote)
|
||||
config_path = get_model_file(model_name, MODELSLIM_CONFIG_FILENAME, revision=revision)
|
||||
|
||||
if config_path is not None:
|
||||
with open(config_path) as f:
|
||||
self.quant_description = json.load(f)
|
||||
self._apply_extra_quant_adaptations()
|
||||
self._add_kvcache_quant_metadata()
|
||||
return
|
||||
|
||||
# Collect diagnostic info for the error message
|
||||
json_names: list[str] = []
|
||||
if os.path.isdir(model_name):
|
||||
json_files = glob.glob(os.path.join(model_name, "*.json"))
|
||||
json_names = [os.path.basename(f) for f in json_files]
|
||||
|
||||
# Config file not found - raise a friendly error message
|
||||
logger.error(
|
||||
"ModelSlim quantization config not found for model '%s'. Searched path: %s. Found JSON files: %s.",
|
||||
model_name,
|
||||
model_name,
|
||||
json_names if json_names else "N/A",
|
||||
)
|
||||
raise ValueError(
|
||||
"\n"
|
||||
+ "=" * 80
|
||||
+ "\n"
|
||||
+ "ERROR: ModelSlim Quantization Config Not Found\n"
|
||||
+ "=" * 80
|
||||
+ "\n"
|
||||
+ "\n"
|
||||
+ f"You have enabled '--quantization {ASCEND_QUANTIZATION_METHOD}' "
|
||||
+ "(ModelSlim quantization),\n"
|
||||
+ f"but the model '{model_name}' does not contain the required\n"
|
||||
+ f"quantization config file ('{MODELSLIM_CONFIG_FILENAME}').\n"
|
||||
+ "\n"
|
||||
+ "This usually means the model weights are NOT quantized by "
|
||||
+ "ModelSlim.\n"
|
||||
+ "\n"
|
||||
+ "Please choose one of the following solutions:\n"
|
||||
+ "\n"
|
||||
+ " Solution 1: Remove the quantization option "
|
||||
+ "(for float/unquantized models)\n"
|
||||
+ " "
|
||||
+ "-" * 58
|
||||
+ "\n"
|
||||
+ f" Remove '--quantization {ASCEND_QUANTIZATION_METHOD}' from "
|
||||
+ "your command if you want to\n"
|
||||
+ " run the model with the original (float) weights.\n"
|
||||
+ "\n"
|
||||
+ " Example:\n"
|
||||
+ f" vllm serve {model_name}\n"
|
||||
+ "\n"
|
||||
+ " Solution 2: Quantize your model weights with ModelSlim first\n"
|
||||
+ " "
|
||||
+ "-" * 58
|
||||
+ "\n"
|
||||
+ " Use the ModelSlim tool to quantize your model weights "
|
||||
+ "before deployment.\n"
|
||||
+ " After quantization, the model directory should contain "
|
||||
+ f"'{MODELSLIM_CONFIG_FILENAME}'.\n"
|
||||
+ " For more information, please refer to:\n"
|
||||
+ " https://gitee.com/ascend/msit/tree/master/msmodelslim\n"
|
||||
+ "\n"
|
||||
+ (f" (Found JSON files in model directory: {json_names})\n" if json_names else "")
|
||||
+ "=" * 80
|
||||
)
|
||||
|
||||
def _apply_extra_quant_adaptations(self) -> None:
|
||||
"""Apply extra adaptations to the quant_description dict.
|
||||
|
||||
This handles known key transformations such as shared_head and
|
||||
weight_packed mappings.
|
||||
"""
|
||||
if "hc_head_fn" in self.quant_description:
|
||||
# TODO
|
||||
extra_quant_dict = {}
|
||||
for name in self.quant_description:
|
||||
new_name = name
|
||||
if not name.startswith("model"):
|
||||
new_name = f"model.{name}"
|
||||
extra_quant_dict[new_name] = self.quant_description[name]
|
||||
self.quant_description.update(extra_quant_dict)
|
||||
|
||||
extra_quant_dict = {}
|
||||
for name in self.quant_description:
|
||||
new_name = name
|
||||
if "attn" in name and "self_attn" not in name:
|
||||
new_name = name.replace(".attn.", ".self_attn.")
|
||||
extra_quant_dict[new_name] = self.quant_description[name]
|
||||
self.quant_description.update(extra_quant_dict)
|
||||
|
||||
extra_quant_dict = {}
|
||||
for name in self.quant_description:
|
||||
new_name = name
|
||||
if "ffn" in name:
|
||||
new_name = name.replace("ffn", "mlp")
|
||||
extra_quant_dict[new_name] = self.quant_description[name]
|
||||
self.quant_description.update(extra_quant_dict)
|
||||
|
||||
extra_quant_dict = {}
|
||||
for name in self.quant_description:
|
||||
new_name = name
|
||||
if "w1" in name:
|
||||
new_name = name.replace(".w1.", ".gate_proj.")
|
||||
if "w2" in name:
|
||||
new_name = name.replace(".w2.", ".down_proj.")
|
||||
if "w3" in name:
|
||||
new_name = name.replace(".w3.", ".up_proj.")
|
||||
|
||||
if "head" in name and "lm_head" not in name:
|
||||
new_name = name.replace("head", "lm_head")
|
||||
if "embed" in name and "embed_tokens" not in name:
|
||||
new_name = name.replace("embed", "embed_tokens")
|
||||
extra_quant_dict[new_name] = self.quant_description[name]
|
||||
self.quant_description.update(extra_quant_dict)
|
||||
|
||||
extra_quant_dict = {}
|
||||
for k in self.quant_description:
|
||||
if "shared_head" in k:
|
||||
new_k = k.replace(".shared_head.", ".")
|
||||
extra_quant_dict[new_k] = self.quant_description[k]
|
||||
if "transformer.shared_head.output." in k:
|
||||
# Step3.5 MTP checkpoints describe per-layer draft logits heads
|
||||
# as ``transformer.shared_head.output``. The vLLM model module
|
||||
# exposes the same parameter as ``shared_head.head``.
|
||||
new_k = k.replace(
|
||||
"transformer.shared_head.output.",
|
||||
"shared_head.head.",
|
||||
)
|
||||
extra_quant_dict[new_k] = self.quant_description[k]
|
||||
if "transformer.shared_head.norm." in k:
|
||||
new_k = k.replace(
|
||||
"transformer.shared_head.norm.",
|
||||
"shared_head.norm.",
|
||||
)
|
||||
extra_quant_dict[new_k] = self.quant_description[k]
|
||||
if "weight_packed" in k:
|
||||
new_k = k.replace("weight_packed", "weight")
|
||||
extra_quant_dict[new_k] = self.quant_description[k]
|
||||
self.quant_description.update(extra_quant_dict)
|
||||
|
||||
def _add_kvcache_quant_metadata(self):
|
||||
fa_quant_type = self.quant_description.get("fa_quant_type", "")
|
||||
self.enable_fa_quant = fa_quant_type != ""
|
||||
self.kvcache_quant_layers = []
|
||||
indexer_quant_type = self.quant_description.get("indexer_quant_type", "")
|
||||
self.enable_indexer_quant = indexer_quant_type != ""
|
||||
self.indexer_quant_layers = []
|
||||
kv_quant_type = self.quant_description.get("kv_cache_type", "")
|
||||
self.enable_c8_quant = kv_quant_type == "C8"
|
||||
self.c8_quant_layers = []
|
||||
if self.enable_fa_quant or self.enable_indexer_quant or self.enable_c8_quant:
|
||||
for key in self.quant_description:
|
||||
_id = "".join(re.findall(r"\.(\d+)\.", key))
|
||||
if "fa_k.scale" in key:
|
||||
self.kvcache_quant_layers.append(int(_id))
|
||||
if "indexer.quant_type" in key:
|
||||
self.indexer_quant_layers.append(int(_id))
|
||||
if "k_proj.kv_cache_scale" in key:
|
||||
self.c8_quant_layers.append(int(_id))
|
||||
Reference in New Issue
Block a user