| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from fractions import Fraction |
| 5 | from typing import TYPE_CHECKING, Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.model_executor.layers.fused_moe import ( |
| 11 | RoutedExperts, |
| 12 | UnquantizedFusedMoEMethod, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.linear import ( |
| 15 | LinearBase, |
| 16 | UnquantizedLinearMethod, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.quantization import ( |
| 19 | QuantizationConfig, |
| 20 | QuantizationMethods, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 23 | |
| 24 | from .config_parser import INCConfigParser |
| 25 | |
| 26 | if TYPE_CHECKING: |
| 27 | from vllm.model_executor.models.utils import WeightsMapper |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | |
| 32 | class INCConfig(QuantizationConfig): |
| 33 | """Config class for Intel Neural Compressor (INC). |
| 34 | Repo: https://github.com/intel/neural-compressor |
| 35 | """ |
| 36 | |
| 37 | SUPPORTED_BITS = {2, 3, 4, 5, 6, 7, 8} |
| 38 | DEFAULT_INT_PACKING_FORMAT = "auto_round:auto_gptq" |
| 39 | SUPPORTED_DTYPES = {"int", "mx_fp", "fp"} |
| 40 | SUPPORTED_FORMATS = { |
| 41 | "auto_round:auto_gptq", |
| 42 | "auto_round:auto_awq", |
| 43 | "auto_round:fp8", |
| 44 | "auto_round:llm_compressor", |
| 45 | } |
| 46 | SUPPORTED_BACKENDS = { |
| 47 | "auto", |
| 48 | "gptq", |
| 49 | "gptq:marlin", |
| 50 | "awq", |
| 51 | "awq:marlin", |
| 52 | "marlin", |
| 53 | } |
| 54 | |
| 55 | FP8_BLOCK_BITS = 8 |
| 56 | FP8_BLOCK_DATA_TYPE = "fp" |
| 57 | FP8_BLOCK_PACKING_FORMAT = "auto_round:fp8" |
| 58 | FP8_BLOCK_FMT = "e4m3" |
| 59 | FP8_BLOCK_ACTIVATION_SCHEME = "dynamic" |
| 60 | |
| 61 | MXFP8_BITS = 8 |
| 62 | MXFP8_GROUP_SIZE = 32 |
| 63 | MXFP8_DATA_TYPE = "mx_fp" |
| 64 | MXFP8_PACKING_FORMAT = "auto_round:llm_compressor" |
| 65 | |
| 66 | @staticmethod |
| 67 | def _normalize_data_type(data_type: str) -> str: |
| 68 | # Keep a small alias for branches that already experimented with "fp8". |
| 69 | return "fp" if data_type == "fp8" else data_type |
| 70 | |
| 71 | def __init__( |
| 72 | self, |
| 73 | weight_bits: int, |
| 74 | group_size: int | tuple[int, int], |
| 75 | sym: bool = True, |
| 76 | packing_format: str = DEFAULT_INT_PACKING_FORMAT, |
| 77 | block_name_to_quantize: str | list[str] | None = None, |
| 78 | extra_config: dict[str, Any] | None = None, |
| 79 | data_type: str = "int", |
| 80 | backend: str = "auto", |
| 81 | ) -> None: |
| 82 | super().__init__() |
| 83 | |
| 84 | data_type = self._normalize_data_type(data_type) |
| 85 | group_size = INCConfigParser._normalize_group_size(group_size) |
| 86 | |
| 87 | if weight_bits not in self.SUPPORTED_BITS: |
| 88 | raise ValueError( |
| 89 | f"Unsupported weight_bits: {weight_bits}, " |
| 90 | f"currently only support {self.SUPPORTED_BITS}." |
| 91 | ) |
| 92 | # auto-round mxfp data_type is e.g. "mx_fp4" / "mx_fp4e2m1"; match the |
| 93 | # "mx_fp" family by substring like auto_round.compressors.is_mx_fp. |
| 94 | is_mxfp = "mx_fp" in data_type |
| 95 | if data_type not in self.SUPPORTED_DTYPES and not is_mxfp: |
| 96 | raise ValueError( |
| 97 | f"Unsupported data_type: {data_type}, " |
| 98 | f"currently only support {self.SUPPORTED_DTYPES}." |
| 99 | ) |
| 100 | if packing_format not in self.SUPPORTED_FORMATS: |
| 101 | raise ValueError( |
| 102 | f"Unsupported packing_format: {packing_format}, " |
| 103 | f"currently only support {self.SUPPORTED_FORMATS}." |
| 104 | ) |
| 105 | if backend not in self.SUPPORTED_BACKENDS: |
| 106 | raise ValueError( |
| 107 | f"Unsupported backend: {backend}, " |
| 108 | f"currently only support {self.SUPPORTED_BACKENDS}." |
| 109 | ) |
| 110 | |
| 111 | self.weight_bits = weight_bits |
| 112 | self.group_size = group_size |
| 113 | self.sym = sym |
| 114 | self.packing_format = packing_format |
| 115 | self.block_name_to_quantize = ( |
| 116 | block_name_to_quantize.split(",") |
| 117 | if isinstance(block_name_to_quantize, str) |
| 118 | else block_name_to_quantize |
| 119 | ) |
| 120 | self.extra_config = extra_config |
| 121 | self.data_type = data_type |
| 122 | self.backend = backend |
| 123 | self.pack_factor = Fraction(32, weight_bits) |
| 124 | self.config_parser = INCConfigParser(self) |
| 125 | |
| 126 | self._validate_supported_quantization() |
| 127 | |
| 128 | def __repr__(self) -> str: |
| 129 | return ( |
| 130 | f"INCConfig(weight_bits={self.weight_bits}, " |
| 131 | f"group_size={self.group_size}, sym={self.sym})" |
| 132 | ) |
| 133 | |
| 134 | @property |
| 135 | def is_mxfp(self) -> bool: |
| 136 | # auto-round mxfp data_type is the "mx_fp" family (e.g. "mx_fp", |
| 137 | # "mx_fp4e2m1"); match by substring like auto_round.compressors.is_mx_fp. |
| 138 | return "mx_fp" in self.data_type |
| 139 | |
| 140 | @property |
| 141 | def is_mxfp8(self) -> bool: |
| 142 | # MXFP4 and MXFP8 share data_type "mx_fp" and differ only by bit width. |
| 143 | return self.is_mxfp and self.weight_bits == self.MXFP8_BITS |
| 144 | |
| 145 | @classmethod |
| 146 | def _check_expected( |
| 147 | cls, |
| 148 | values: dict[str, Any], |
| 149 | expected_fields: dict[str, Any], |
| 150 | quantization_name: str, |
| 151 | ) -> None: |
| 152 | for field_name, expected_value in expected_fields.items(): |
| 153 | actual_value = cls.get_from_keys_or(values, [field_name], expected_value) |
| 154 | assert actual_value == expected_value, ( |
| 155 | f"{quantization_name} only supports " |
| 156 | f"{field_name}={expected_value!r}, " |
| 157 | f"but found {field_name}={actual_value!r}." |
| 158 | ) |
| 159 | |
| 160 | def _validate_int_quantization(self) -> None: |
| 161 | int_packing_formats = { |
| 162 | self.DEFAULT_INT_PACKING_FORMAT, |
| 163 | "auto_round:auto_awq", |
| 164 | } |
| 165 | assert isinstance(self.group_size, int), ( |
| 166 | "INC int quantization requires scalar group_size, " |
| 167 | f"but found group_size={self.group_size!r}." |
| 168 | ) |
| 169 | assert self.packing_format in int_packing_formats, ( |
| 170 | "INC int quantization only supports " |
| 171 | f"packing_format in {sorted(int_packing_formats)!r}, " |
| 172 | f"but found packing_format={self.packing_format!r}." |
| 173 | ) |
| 174 | |
| 175 | def _validate_fp8_block_quantization(self) -> None: |
| 176 | assert self.weight_bits == self.FP8_BLOCK_BITS, ( |
| 177 | "INC block-wise FP8 only supports bits=8, " |
| 178 | f"but found bits={self.weight_bits}." |
| 179 | ) |
| 180 | assert isinstance(self.group_size, tuple), ( |
| 181 | "INC block-wise FP8 requires a 2-D group_size, " |
| 182 | f"but found group_size={self.group_size!r}." |
| 183 | ) |
| 184 | assert self.sym, "INC block-wise FP8 only supports symmetric weights." |
| 185 | assert self.packing_format == self.FP8_BLOCK_PACKING_FORMAT, ( |
| 186 | "INC block-wise FP8 only supports " |
| 187 | f"packing_format={self.FP8_BLOCK_PACKING_FORMAT!r}, " |
| 188 | f"but found {self.packing_format!r}." |
| 189 | ) |
| 190 | assert self.backend == "auto", ( |
| 191 | "INC block-wise FP8 only supports backend='auto', " |
| 192 | f"but found backend={self.backend!r}." |
| 193 | ) |
| 194 | |
| 195 | def _validate_mxfp_quantization(self) -> None: |
| 196 | assert self.weight_bits in {4, self.MXFP8_BITS}, ( |
| 197 | "INC MXFP only supports bits=4 or bits=8, " |
| 198 | f"but found bits={self.weight_bits}." |
| 199 | ) |
| 200 | assert self.group_size == self.MXFP8_GROUP_SIZE, ( |
| 201 | "INC MXFP only supports group_size=32, " |
| 202 | f"but found group_size={self.group_size}." |
| 203 | ) |
| 204 | assert self.sym, "INC MXFP only supports symmetric weights." |
| 205 | assert self.packing_format == self.MXFP8_PACKING_FORMAT, ( |
| 206 | "INC MXFP only supports " |
| 207 | f"packing_format={self.MXFP8_PACKING_FORMAT!r}, " |
| 208 | f"but found {self.packing_format!r}." |
| 209 | ) |
| 210 | assert self.backend == "auto", ( |
| 211 | "INC MXFP only supports backend='auto', " |
| 212 | f"but found backend={self.backend!r}." |
| 213 | ) |
| 214 | |
| 215 | def _validate_supported_quantization(self) -> None: |
| 216 | if self.data_type == "int": |
| 217 | self._validate_int_quantization() |
| 218 | elif self.data_type == self.FP8_BLOCK_DATA_TYPE: |
| 219 | self._validate_fp8_block_quantization() |
| 220 | elif self.data_type == self.MXFP8_DATA_TYPE: |
| 221 | self._validate_mxfp_quantization() |
| 222 | |
| 223 | def _validate_raw_config(self, config: dict[str, Any]) -> None: |
| 224 | if self.data_type == self.FP8_BLOCK_DATA_TYPE: |
| 225 | assert isinstance(self.group_size, tuple), ( |
| 226 | "INC block-wise FP8 group_size must be a tuple." |
| 227 | ) |
| 228 | |
| 229 | expected_fields = { |
| 230 | "act_bits": self.FP8_BLOCK_BITS, |
| 231 | "act_group_size": self.group_size[0], |
| 232 | "act_data_type": self.FP8_BLOCK_DATA_TYPE, |
| 233 | "act_sym": True, |
| 234 | "act_dynamic": True, |
| 235 | "activation_scheme": self.FP8_BLOCK_ACTIVATION_SCHEME, |
| 236 | "fmt": self.FP8_BLOCK_FMT, |
| 237 | } |
| 238 | self._check_expected( |
| 239 | config, |
| 240 | expected_fields, |
| 241 | "INC block-wise FP8", |
| 242 | ) |
| 243 | elif self.is_mxfp8: |
| 244 | expected_fields = { |
| 245 | "act_bits": self.MXFP8_BITS, |
| 246 | "act_data_type": self.MXFP8_DATA_TYPE, |
| 247 | "act_group_size": self.MXFP8_GROUP_SIZE, |
| 248 | "act_sym": True, |
| 249 | "act_dynamic": True, |
| 250 | } |
| 251 | self._check_expected( |
| 252 | config, |
| 253 | expected_fields, |
| 254 | "INC MXFP8", |
| 255 | ) |
| 256 | |
| 257 | @classmethod |
| 258 | def get_name(cls) -> QuantizationMethods: |
| 259 | return "inc" |
| 260 | |
| 261 | @classmethod |
| 262 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 263 | return [torch.half, torch.bfloat16] |
| 264 | |
| 265 | @classmethod |
| 266 | def get_min_capability(cls) -> int: |
| 267 | return 60 |
| 268 | |
| 269 | @classmethod |
| 270 | def get_config_filenames(cls) -> list[str]: |
| 271 | return ["quantization_config.json"] |
| 272 | |
| 273 | @classmethod |
| 274 | def from_config(cls, config: dict[str, Any]) -> "INCConfig": |
| 275 | data_type = cls._normalize_data_type( |
| 276 | cls.get_from_keys_or(config, ["data_type"], "int") |
| 277 | ) |
| 278 | weight_bits = cls.get_from_keys(config, ["bits"]) |
| 279 | group_size = INCConfigParser._normalize_group_size( |
| 280 | cls.get_from_keys(config, ["group_size"]) |
| 281 | ) |
| 282 | packing_format = cls.get_from_keys_or(config, ["packing_format"], None) |
| 283 | |
| 284 | if packing_format is None: |
| 285 | if ( |
| 286 | data_type == cls.FP8_BLOCK_DATA_TYPE |
| 287 | and weight_bits == cls.FP8_BLOCK_BITS |
| 288 | and isinstance(group_size, tuple) |
| 289 | ): |
| 290 | packing_format = cls.FP8_BLOCK_PACKING_FORMAT |
| 291 | elif "mx_fp" in data_type: |
| 292 | packing_format = cls.MXFP8_PACKING_FORMAT |
| 293 | else: |
| 294 | packing_format = cls.DEFAULT_INT_PACKING_FORMAT |
| 295 | |
| 296 | quant_config = cls( |
| 297 | weight_bits=weight_bits, |
| 298 | group_size=group_size, |
| 299 | sym=cls.get_from_keys_or(config, ["sym"], True), |
| 300 | packing_format=packing_format, |
| 301 | block_name_to_quantize=cls.get_from_keys_or( |
| 302 | config, ["block_name_to_quantize", "to_quant_block_names"], None |
| 303 | ), |
| 304 | extra_config=cls.get_from_keys_or(config, ["extra_config"], None), |
| 305 | data_type=data_type, |
| 306 | backend=cls.get_from_keys_or(config, ["backend", "vllm_backend"], "auto"), |
| 307 | ) |
| 308 | quant_config._validate_raw_config(config) |
| 309 | return quant_config |
| 310 | |
| 311 | def get_layer_config(self, layer, layer_name: str): |
| 312 | return self.config_parser.get_layer_config(layer, layer_name) |
| 313 | |
| 314 | def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): |
| 315 | if self.block_name_to_quantize is not None: |
| 316 | self.block_name_to_quantize = hf_to_vllm_mapper.apply_list( |
| 317 | self.block_name_to_quantize |
| 318 | ) |
| 319 | if self.extra_config is not None: |
| 320 | self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config) |
| 321 | |
| 322 | def get_quant_method(self, layer: torch.nn.Module, prefix: str): |
| 323 | from .schemes.factory import resolve_scheme |
| 324 | |
| 325 | if prefix and self.extra_config: |
| 326 | for layer_name in self.extra_config: |
| 327 | if ( |
| 328 | layer_name == prefix or layer_name == f"model.{prefix}" |
| 329 | ) and self.extra_config[layer_name].get("bits", 16) >= 16: |
| 330 | if isinstance(layer, (LinearBase, ParallelLMHead)): |
| 331 | return UnquantizedLinearMethod() |
| 332 | if isinstance(layer, RoutedExperts): |
| 333 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 334 | return None |
| 335 | |
| 336 | layer_config = self.config_parser.resolve(layer, prefix) |
| 337 | if not layer_config.quantized: |
| 338 | if isinstance(layer, (LinearBase, ParallelLMHead)): |
| 339 | return UnquantizedLinearMethod() |
| 340 | if isinstance(layer, RoutedExperts): |
| 341 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 342 | return None |
| 343 | |
| 344 | logger.debug( |
| 345 | "[%s] Type: %s, Bits: %s, Group Size: %s, Sym: %s", |
| 346 | prefix, |
| 347 | layer.__class__.__name__, |
| 348 | layer_config.bits, |
| 349 | layer_config.group_size, |
| 350 | layer_config.sym, |
| 351 | ) |
| 352 | |
| 353 | scheme = resolve_scheme(layer_config) |
| 354 | if isinstance(layer, (LinearBase, ParallelLMHead)): |
| 355 | return scheme.get_linear_method(self, layer, prefix, layer_config) |
| 356 | if isinstance(layer, RoutedExperts): |
| 357 | return scheme.get_moe_method(self, layer, prefix, layer_config) |
| 358 | return None |
| 359 | |
| 360 | @classmethod |
| 361 | def override_quantization_method( |
| 362 | cls, hf_quant_cfg, user_quant, hf_config=None |
| 363 | ) -> "QuantizationMethods | None": |
| 364 | """Override the `auto-round` method to `inc`.""" |
| 365 | is_auto_round_format = hf_quant_cfg.get("quant_method", None) == "auto-round" |
| 366 | if is_auto_round_format: |
| 367 | return cls.get_name() |
| 368 | return None |