| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import inspect |
| 5 | from abc import ABC, abstractmethod |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import regex as re |
| 9 | import torch |
| 10 | from torch import nn |
| 11 | from transformers import PreTrainedConfig |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 15 | from vllm.model_executor.layers.quantization.online.base import ( |
| 16 | OnlineQuantizationConfig, |
| 17 | ) |
| 18 | from vllm.model_executor.models.utils import WeightsMapper |
| 19 | else: |
| 20 | QuantizationMethods = str |
| 21 | |
| 22 | |
| 23 | class QuantizeMethodBase(ABC): |
| 24 | """Base class for different quantized methods.""" |
| 25 | |
| 26 | requires_device_loading: bool = True |
| 27 | """Whether post-load processing requires parameters on the target device.""" |
| 28 | |
| 29 | uses_meta_device: bool = False |
| 30 | """ |
| 31 | Whether this method creates weights on meta device for online quantization. |
| 32 | When True, weights are created on meta device and quantized layer-wise |
| 33 | in process_weights_after_loading, reducing peak memory during loading. |
| 34 | """ |
| 35 | |
| 36 | supports_pre_processed_weights: bool = False |
| 37 | """ |
| 38 | Whether ``process_weights_after_loading`` supports running under |
| 39 | ``weights_already_processed``. Methods must skip tensor transforms in |
| 40 | that mode; the loader driver rejects methods that do not declare support. |
| 41 | """ |
| 42 | |
| 43 | @abstractmethod |
| 44 | def create_weights( |
| 45 | self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs |
| 46 | ): |
| 47 | """Create weights for a layer. |
| 48 | |
| 49 | The weights will be set as attributes of the layer.""" |
| 50 | raise NotImplementedError |
| 51 | |
| 52 | @abstractmethod |
| 53 | def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: |
| 54 | """Apply the weights in layer to the input tensor. |
| 55 | |
| 56 | Expects create_weights to have been called before on the layer.""" |
| 57 | raise NotImplementedError |
| 58 | |
| 59 | # Not required functions |
| 60 | def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: |
| 61 | """Gather embeddings in the layer based on indices in the input tensor. |
| 62 | |
| 63 | Expects create_weights to have been called before on the layer.""" |
| 64 | raise NotImplementedError |
| 65 | |
| 66 | # Not required functions |
| 67 | def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module): |
| 68 | """Tie ``layer``'s weight to ``embed_tokens``' weight. |
| 69 | |
| 70 | The default shares the weight tensor, which is the standard behavior for |
| 71 | tied word embeddings and matches what ``ParallelLMHead.tie_weights`` did |
| 72 | directly before quantization methods became responsible for it. |
| 73 | Quantization methods that need special weight handling (e.g. repacked |
| 74 | weights) override this. |
| 75 | |
| 76 | Expects create_weights to have been called before on the layer.""" |
| 77 | layer.weight = embed_tokens.weight |
| 78 | return layer |
| 79 | |
| 80 | def process_weights_after_loading(self, layer: nn.Module) -> None: |
| 81 | """Process the weight after loading. |
| 82 | |
| 83 | This can be used for example, to transpose weights for computation. |
| 84 | """ |
| 85 | return |
| 86 | |
| 87 | |
| 88 | def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: |
| 89 | """Not all quant methods have embedding implemented, so we need to check that |
| 90 | it exists for our given method. We check this by making sure the function |
| 91 | has been changed from the base implementation. |
| 92 | """ |
| 93 | base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None) |
| 94 | class_embedding = inspect.getattr_static(method_class, "embedding", None) |
| 95 | |
| 96 | return class_embedding is not None and class_embedding is not base_embedding |
| 97 | |
| 98 | |
| 99 | class QuantizationConfig(ABC): |
| 100 | """Base class for quantization configs.""" |
| 101 | |
| 102 | _ignore_unexpected_suffixes = ( |
| 103 | ".q_scale", |
| 104 | ".k_scale", |
| 105 | ".v_scale", |
| 106 | ".q_zero_point", |
| 107 | ".k_zero_point", |
| 108 | ".v_zero_point", |
| 109 | ) |
| 110 | """Suffixes of quantization parameters that may be present in the checkpoint but |
| 111 | not in the model, and should be ignored if unexpected during loading. These are used |
| 112 | after remapping, so should be in vLLM format (e.g. .q_scale, not .q.scale).""" |
| 113 | online_quantization_config: "OnlineQuantizationConfig | None" = None |
| 114 | |
| 115 | def __init__(self): |
| 116 | super().__init__() |
| 117 | # mapping is updated by models as they initialize |
| 118 | self.packed_modules_mapping: dict[str, list[str]] = dict() |
| 119 | |
| 120 | @abstractmethod |
| 121 | def get_name(self) -> QuantizationMethods: |
| 122 | """Name of the quantization method.""" |
| 123 | raise NotImplementedError |
| 124 | |
| 125 | @abstractmethod |
| 126 | def get_supported_act_dtypes(self) -> list[torch.dtype]: |
| 127 | """List of supported activation dtypes.""" |
| 128 | raise NotImplementedError |
| 129 | |
| 130 | @classmethod |
| 131 | @abstractmethod |
| 132 | def get_min_capability(cls) -> int: |
| 133 | """Minimum GPU capability to support the quantization method. |
| 134 | |
| 135 | E.g., 70 for Volta, 75 for Turing, 80 for Ampere. |
| 136 | This requirement is due to the custom CUDA kernels used by the |
| 137 | quantization method. |
| 138 | """ |
| 139 | raise NotImplementedError |
| 140 | |
| 141 | @staticmethod |
| 142 | @abstractmethod |
| 143 | def get_config_filenames() -> list[str]: |
| 144 | """List of filenames to search for in the model directory.""" |
| 145 | raise NotImplementedError |
| 146 | |
| 147 | @classmethod |
| 148 | @abstractmethod |
| 149 | def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig": |
| 150 | """Create a config class from the model's quantization config.""" |
| 151 | raise NotImplementedError |
| 152 | |
| 153 | @classmethod |
| 154 | def override_quantization_method( |
| 155 | cls, |
| 156 | hf_quant_cfg: dict[str, Any], |
| 157 | user_quant: str | None, |
| 158 | hf_config: Any = None, |
| 159 | ) -> QuantizationMethods | None: |
| 160 | """Detects if this quantization method can support a given checkpoint |
| 161 | format by overriding the user specified quantization method -- |
| 162 | this method should only be overwritten by subclasses in exceptional |
| 163 | circumstances. |
| 164 | |
| 165 | Args: |
| 166 | hf_quant_cfg: The checkpoint's quantization config dict. |
| 167 | user_quant: The user-specified quantization method string. |
| 168 | hf_config: The HuggingFace model config object (e.g. for |
| 169 | model_type checks). May be None if not available. |
| 170 | |
| 171 | """ |
| 172 | return None |
| 173 | |
| 174 | @staticmethod |
| 175 | def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any: |
| 176 | """Get a value from the model's quantization config.""" |
| 177 | for key in keys: |
| 178 | if key in config: |
| 179 | return config[key] |
| 180 | raise ValueError( |
| 181 | f"Cannot find any of {keys} in the model's quantization config." |
| 182 | ) |
| 183 | |
| 184 | @staticmethod |
| 185 | def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any: |
| 186 | """Get an optional value from the model's quantization config.""" |
| 187 | try: |
| 188 | return QuantizationConfig.get_from_keys(config, keys) |
| 189 | except ValueError: |
| 190 | return default |
| 191 | |
| 192 | @abstractmethod |
| 193 | def get_quant_method( |
| 194 | self, layer: torch.nn.Module, prefix: str |
| 195 | ) -> QuantizeMethodBase | None: |
| 196 | """Get the quantize method to use for the quantized layer, from the |
| 197 | pre-quantized checkpoint quant_method. |
| 198 | |
| 199 | Args: |
| 200 | layer: The layer for the quant method. |
| 201 | prefix: The full name of the layer in the state dict |
| 202 | Returns: |
| 203 | The quantize method. None if the given layer doesn't support quant |
| 204 | method. |
| 205 | |
| 206 | """ |
| 207 | raise NotImplementedError |
| 208 | |
| 209 | @staticmethod |
| 210 | def get_cache_scale_mapper() -> "WeightsMapper": |
| 211 | """Mapping from checkpoint KV-cache scale names to vLLM scale names. |
| 212 | |
| 213 | Returning a mapper here causes `AutoWeightsLoader` to apply it to the |
| 214 | weight stream automatically; individual model `load_weights` methods |
| 215 | do not need to know about KV-cache scales. |
| 216 | """ |
| 217 | from vllm.model_executor.models.utils import WeightsMapper |
| 218 | |
| 219 | orig_to_new_regex = { |
| 220 | # Deprecated fused kv_scale -> attn.k_scale |
| 221 | re.compile(r"\.kv_scale$"): r".attn.k_scale", |
| 222 | # ModelOpt: .self_attn.{k,v}_proj.{k,v}_scale -> .self_attn.attn.* |
| 223 | re.compile(r"\.self_attn\.[kv]_proj\.([kv])_scale$"): ( |
| 224 | r".self_attn.attn.\1_scale" |
| 225 | ), |
| 226 | # Fused QKV / qkqkv proj: .self_attn.qk(qk)v_proj.{k,v}_scale -> attn |
| 227 | re.compile(r"\.self_attn\.qk(?:qk)?v_proj\.([kv])_scale$"): ( |
| 228 | r".self_attn.attn.\1_scale" |
| 229 | ), |
| 230 | # NemotronH: .mixer.{k,v}_proj.{k,v}_scale -> .mixer.attn.* |
| 231 | re.compile(r"\.mixer\.[kv]_proj\.([kv])_scale$"): r".mixer.attn.\1_scale", |
| 232 | # HYV3: .self_attn.q.scale -> .self_attn.attn.q_scale |
| 233 | re.compile(r"\.self_attn\.q\.scale$"): r".self_attn.attn.q_scale", |
| 234 | # HYV3: .self_attn.{k,v}_cache.scale -> .self_attn.attn.{k,v}_scale |
| 235 | re.compile(r"\.self_attn\.([kv])_cache\.scale$"): ( |
| 236 | r".self_attn.attn.\1_scale" |
| 237 | ), |
| 238 | # Default: .{q,k,v}_scale -> .attn.{q,k,v}_scale (unless already .attn) |
| 239 | re.compile(r"(?<!\.attn)\.([qkv])_scale$"): r".attn.\1_scale", |
| 240 | re.compile(r"(?<!\.attn)\.([qkv])_zero_point$"): r".attn.\1_zero_point", |
| 241 | } |
| 242 | return WeightsMapper(orig_to_new_regex=orig_to_new_regex) |
| 243 | |
| 244 | @staticmethod |
| 245 | def get_checkpoint_weight_mapper() -> "WeightsMapper": |
| 246 | """Discard activation-order metadata unused by supported kernels.""" |
| 247 | from vllm.model_executor.models.utils import WeightsMapper |
| 248 | |
| 249 | return WeightsMapper(orig_to_new_suffix={".g_idx": None}) |
| 250 | |
| 251 | def apply_vllm_mapper( # noqa: B027 |
| 252 | self, hf_to_vllm_mapper: "WeightsMapper" |
| 253 | ): |
| 254 | """Interface for models to update module names referenced in |
| 255 | quantization configs in order to reflect the vllm model structure |
| 256 | |
| 257 | Args: |
| 258 | hf_to_vllm_mapper: maps from hf model structure (the assumed |
| 259 | structure of the qconfig) to vllm model structure |
| 260 | |
| 261 | """ |
| 262 | # TODO (@kylesayrs): add implementations for all subclasses |
| 263 | pass |
| 264 | |
| 265 | def maybe_update_config( # noqa: B027 |
| 266 | self, |
| 267 | model_name: str, |
| 268 | hf_config: PreTrainedConfig | None = None, |
| 269 | revision: str | None = None, |
| 270 | ): |
| 271 | """Interface to update values after config initialization. |
| 272 | |
| 273 | Args: |
| 274 | model_name: The name of the model |
| 275 | hf_config: The Hugging Face config of the model |
| 276 | revision: The revision of the model |
| 277 | Returns: |
| 278 | |
| 279 | """ |
| 280 | # TODO: revision is never passed currently in vllm.py, |
| 281 | # but is used in subclasses, should we remove this parameter? |
| 282 | pass |
| 283 | |
| 284 | |
| 285 | def resolve_quant_method( |
| 286 | quant_config: QuantizationConfig, layer: torch.nn.Module, prefix: str |
| 287 | ) -> QuantizeMethodBase | None: |
| 288 | """Return the checkpoint method with configured online quantization.""" |
| 289 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 290 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 291 | UnquantizedFusedMoEMethod, |
| 292 | ) |
| 293 | from vllm.model_executor.layers.linear import ( |
| 294 | LinearBase, |
| 295 | UnquantizedLinearMethod, |
| 296 | ) |
| 297 | |
| 298 | base_quant_method = quant_config.get_quant_method(layer, prefix) |
| 299 | if quant_config.online_quantization_config is None: |
| 300 | return base_quant_method |
| 301 | # Online quantization currently supports only LinearBase and RoutedExperts. |
| 302 | # Embeddings and ParallelLMHead retain their checkpoint quantization method. |
| 303 | if not isinstance(layer, (LinearBase, RoutedExperts)): |
| 304 | return base_quant_method |
| 305 | |
| 306 | quant_config.online_quantization_config.packed_modules_mapping = ( |
| 307 | quant_config.packed_modules_mapping |
| 308 | ) |
| 309 | checkpoint_is_quantized = base_quant_method is not None and not isinstance( |
| 310 | base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) |
| 311 | ) |
| 312 | online_target = quant_config.online_quantization_config.resolve_quant_method_cls( |
| 313 | layer, prefix |
| 314 | ) |
| 315 | if checkpoint_is_quantized: |
| 316 | if online_target is not None: |
| 317 | raise ValueError( |
| 318 | f"Cannot apply requested online quantization {online_target[3]} to " |
| 319 | f"pre-quantized layer {prefix}: {base_quant_method} was already " |
| 320 | "selected by the checkpoint quantization config." |
| 321 | ) |
| 322 | return base_quant_method |
| 323 | if online_target is None: |
| 324 | return base_quant_method |
| 325 | return quant_config.online_quantization_config.get_quant_method(layer, prefix) |