| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib |
| 4 | import json |
| 5 | import types |
| 6 | from importlib.util import find_spec |
| 7 | from typing import Any |
| 8 | |
| 9 | import regex as re |
| 10 | import torch |
| 11 | import torch.nn.functional as F |
| 12 | from packaging import version |
| 13 | from torch.nn.parameter import Parameter |
| 14 | |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.model_executor.layers.linear import ( |
| 17 | LinearBase, |
| 18 | LinearMethodBase, |
| 19 | UnquantizedLinearMethod, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 22 | from vllm.model_executor.layers.quantization.base_config import ( |
| 23 | QuantizationConfig, |
| 24 | QuantizeMethodBase, |
| 25 | ) |
| 26 | from vllm.model_executor.utils import set_weight_attrs |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | def _bond_method_to_cls(func, obj): |
| 32 | if hasattr(func, "__self__") or not callable(func): |
| 33 | # If the function is already bound to an instance, return it as is |
| 34 | return func |
| 35 | else: |
| 36 | return types.MethodType(func, obj) |
| 37 | |
| 38 | |
| 39 | def _get_weight_attrs(param): |
| 40 | # record attributes attached to the weight, so we can |
| 41 | # recover later |
| 42 | recorded_weight_attr = {} |
| 43 | for key in param.__dict__: |
| 44 | if hasattr(param, key): |
| 45 | attr = getattr(param, key) |
| 46 | if not callable(attr): |
| 47 | recorded_weight_attr[key] = attr |
| 48 | elif hasattr(attr, "__self__") and param is attr.__self__: |
| 49 | # if attr is a bonded method for an instance, and |
| 50 | # attr.__self__ points to the instance (param) |
| 51 | # we'll record the underlying function object |
| 52 | recorded_weight_attr[key] = attr.__func__ |
| 53 | else: |
| 54 | recorded_weight_attr[key] = attr |
| 55 | return recorded_weight_attr |
| 56 | |
| 57 | |
| 58 | def _restore_weight_attrs(param, recorded_weight_attr): |
| 59 | for attr_name, attr in recorded_weight_attr.items(): |
| 60 | if not hasattr(param, attr_name): |
| 61 | setattr(param, attr_name, _bond_method_to_cls(attr, param)) |
| 62 | |
| 63 | |
| 64 | def torchao_version_at_least(torchao_version: str) -> bool: |
| 65 | if find_spec("torchao"): |
| 66 | try: |
| 67 | if version.parse(importlib.metadata.version("torchao")) >= version.parse( |
| 68 | torchao_version |
| 69 | ): |
| 70 | return True |
| 71 | except (ImportError, version.InvalidVersion): |
| 72 | return False |
| 73 | return False |
| 74 | |
| 75 | |
| 76 | def should_skip(prefix: str, skip_modules: list[str]) -> bool: |
| 77 | """Robust skipping logic: |
| 78 | should_skip("model.model.layers.1.q_proj", |
| 79 | ["model.model.layers.1.q_proj"]) # True |
| 80 | should_skip("model.model.layers.10.o_proj", ["o_proj"]) -> True |
| 81 | should_skip("visual.model.layers.1.q_proj", ["visual"]) -> True |
| 82 | should_skip("model.model.layers.1.q_proj", ["layers.1"]) -> True |
| 83 | should_skip("model.model.layers.11.q_proj", ["layers.1"]) -> False |
| 84 | """ |
| 85 | for s in skip_modules: |
| 86 | if prefix == s: |
| 87 | return True |
| 88 | if f".{s}." in f".{prefix}.": |
| 89 | return True |
| 90 | return False |
| 91 | |
| 92 | |
| 93 | if torchao_version_at_least("0.15.0"): |
| 94 | from torchao.prototype.tensor_conversion.api import ( |
| 95 | convert_to_packed_tensor_based_on_current_hardware, |
| 96 | ) |
| 97 | else: |
| 98 | convert_to_packed_tensor_based_on_current_hardware = lambda t: t |
| 99 | |
| 100 | |
| 101 | def _check_torchao_fp8_activation_capability(torchao_config) -> None: |
| 102 | """Check if the current GPU supports FP8 activation quantization. |
| 103 | |
| 104 | FP8 activation configs (e.g., Float8DynamicActivationFloat8WeightConfig) |
| 105 | require GPU compute capability >= 8.9 (Ada Lovelace / Hopper) on NVIDIA, |
| 106 | or MI300+ on AMD. This check provides a clear error message before |
| 107 | torchao's internal assertion fires with a confusing message. |
| 108 | """ |
| 109 | config_name = type(torchao_config).__name__ |
| 110 | if "Float8" not in config_name or "Activation" not in config_name: |
| 111 | return |
| 112 | |
| 113 | from vllm.platforms import current_platform |
| 114 | |
| 115 | if current_platform.supports_fp8(): |
| 116 | return |
| 117 | |
| 118 | capability = current_platform.get_device_capability() |
| 119 | capability_str = ( |
| 120 | f" (current GPU compute capability: {capability.major}.{capability.minor})" |
| 121 | if capability is not None |
| 122 | else "" |
| 123 | ) |
| 124 | raise ValueError( |
| 125 | f"torchao FP8 activation quantization config '{config_name}' " |
| 126 | f"requires GPU compute capability >= 8.9 (e.g., NVIDIA Ada Lovelace " |
| 127 | f"/ Hopper or AMD MI300+){capability_str}. " |
| 128 | f"For older GPUs, consider using a non-FP8 config such as " |
| 129 | f"Int8WeightOnlyConfig or Int4WeightOnlyConfig." |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | class TorchAOConfig(QuantizationConfig): |
| 134 | """Config class for torchao.""" |
| 135 | |
| 136 | def __init__( |
| 137 | self, |
| 138 | torchao_config, |
| 139 | skip_modules: list[str] | None = None, |
| 140 | is_checkpoint_torchao_serialized: bool = False, |
| 141 | ) -> None: |
| 142 | super().__init__() |
| 143 | self.torchao_config = torchao_config |
| 144 | self.skip_modules = skip_modules or [] |
| 145 | self.is_checkpoint_torchao_serialized = is_checkpoint_torchao_serialized |
| 146 | |
| 147 | def __repr__(self) -> str: |
| 148 | return ( |
| 149 | f"TorchAOConfig({self.torchao_config=}, {self.skip_modules=}, " |
| 150 | f"{self.is_checkpoint_torchao_serialized=})" |
| 151 | ) |
| 152 | |
| 153 | def get_name(self) -> QuantizationMethods: |
| 154 | return "torchao" |
| 155 | |
| 156 | def get_supported_act_dtypes(self) -> list[torch.dtype]: |
| 157 | return [torch.float32, torch.float16, torch.bfloat16] |
| 158 | |
| 159 | @classmethod |
| 160 | def get_min_capability(cls) -> int: |
| 161 | return 75 |
| 162 | |
| 163 | @staticmethod |
| 164 | def get_config_filenames() -> list[str]: |
| 165 | """Torchao doesn't require additional config files, we use |
| 166 | `config.json` from huggingface: `model_config.hf_config` |
| 167 | """ |
| 168 | return [] |
| 169 | |
| 170 | @classmethod |
| 171 | def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig": |
| 172 | """Create the quant config from an hf model config.""" |
| 173 | try: |
| 174 | from torchao.core.config import config_from_dict |
| 175 | except ImportError as err: |
| 176 | raise ImportError( |
| 177 | "Please install torchao>=0.10.0 via " |
| 178 | "`pip install torchao>=0.10.0` to use torchao quantization." |
| 179 | ) from err |
| 180 | |
| 181 | quant_method = cls.get_from_keys_or(config, ["quant_method"], None) |
| 182 | is_checkpoint_torchao_serialized = ( |
| 183 | quant_method is not None and "torchao" in quant_method |
| 184 | ) |
| 185 | |
| 186 | hf_config = cls.get_from_keys_or(config, ["quant_type"], None) |
| 187 | assert hf_config is not None, "quant_type must be specified" |
| 188 | assert len(hf_config) == 1 and "default" in hf_config, ( |
| 189 | "Expected only one key 'default' in quant_type dictionary" |
| 190 | ) |
| 191 | quant_type = hf_config["default"] |
| 192 | ao_config = config_from_dict(quant_type) |
| 193 | |
| 194 | # Adds skipped modules defined in "modules_to_not_convert" |
| 195 | skip_modules = config.get("modules_to_not_convert", []) or [] |
| 196 | |
| 197 | # Adds skipped modules defined in "module_fqn_to_config" |
| 198 | _data = quant_type.get("_data", {}) |
| 199 | if not isinstance(_data, dict): |
| 200 | _data = {} |
| 201 | |
| 202 | module_fqn = _data.get("module_fqn_to_config", {}) |
| 203 | if not isinstance(module_fqn, dict): |
| 204 | module_fqn = {} |
| 205 | |
| 206 | for layer, layer_cfg in module_fqn.items(): |
| 207 | if layer_cfg is None: |
| 208 | skip_modules.append(layer) |
| 209 | |
| 210 | return cls(ao_config, skip_modules, is_checkpoint_torchao_serialized) |
| 211 | |
| 212 | @classmethod |
| 213 | def from_config_file(cls, config_file: str) -> "TorchAOConfig": |
| 214 | """Initialize class from a config file. Example: |
| 215 | ``` |
| 216 | config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()) |
| 217 | fn = "torchao_config.json" |
| 218 | |
| 219 | with open(fn, "w") as f: |
| 220 | f.write(json.dumps(config_to_dict(config))) |
| 221 | ``` |
| 222 | """ |
| 223 | with open(config_file) as f: |
| 224 | f.seek(0) |
| 225 | f_read = f.read() |
| 226 | config_dict = json.loads(f_read) |
| 227 | |
| 228 | hf_config = {"quant_type": {"default": config_dict}} |
| 229 | return cls.from_config(hf_config) |
| 230 | |
| 231 | @classmethod |
| 232 | def from_config_dict_json(cls, config_dict_json: str) -> "TorchAOConfig": |
| 233 | """Initialize class from a config_dict json string, got from |
| 234 | torchao_config_object = some AOBaseConfig object |
| 235 | json.dumps(config_to_dict(torchao_config_object)) |
| 236 | """ |
| 237 | config_dict = json.loads(config_dict_json) |
| 238 | hf_config = {"quant_type": {"default": config_dict}} |
| 239 | return cls.from_config(hf_config) |
| 240 | |
| 241 | def get_quant_method( |
| 242 | self, layer: torch.nn.Module, prefix: str |
| 243 | ) -> "QuantizeMethodBase | None": |
| 244 | if not isinstance(layer, LinearBase): |
| 245 | return None |
| 246 | |
| 247 | from torchao.quantization import ModuleFqnToConfig |
| 248 | |
| 249 | if should_skip(prefix, self.skip_modules): |
| 250 | return UnquantizedLinearMethod() |
| 251 | |
| 252 | module_fqn = prefix |
| 253 | if isinstance(self.torchao_config, ModuleFqnToConfig): |
| 254 | module_fqn_to_config = self.torchao_config.module_fqn_to_config |
| 255 | c = None |
| 256 | if module_fqn in module_fqn_to_config: |
| 257 | assert not module_fqn.startswith("re:"), ( |
| 258 | "module fqn should not start with" |
| 259 | "`re:`, which is used for specifying regex" |
| 260 | ) |
| 261 | c = module_fqn_to_config[module_fqn] |
| 262 | else: |
| 263 | for maybe_module_fqn_pattern in module_fqn_to_config: |
| 264 | if not maybe_module_fqn_pattern.startswith("re:"): |
| 265 | continue |
| 266 | elif re.fullmatch(maybe_module_fqn_pattern[3:], module_fqn): |
| 267 | # we'll apply the config for first fully matched pattern |
| 268 | c = module_fqn_to_config[maybe_module_fqn_pattern] |
| 269 | break |
| 270 | else: |
| 271 | # fallback to use default if no module specific |
| 272 | # config is provided |
| 273 | c = module_fqn_to_config.get("_default", None) |
| 274 | |
| 275 | if c is not None: |
| 276 | current_torchao_config = TorchAOConfig( |
| 277 | c, self.skip_modules, self.is_checkpoint_torchao_serialized |
| 278 | ) |
| 279 | return TorchAOLinearMethod(current_torchao_config) |
| 280 | else: |
| 281 | return UnquantizedLinearMethod() |
| 282 | |
| 283 | return TorchAOLinearMethod(self) |
| 284 | |
| 285 | |
| 286 | def torchao_quantize_param_data( |
| 287 | param: torch.Tensor, torchao_config: Any |
| 288 | ) -> torch.nn.Parameter: |
| 289 | """Quantize a Tensor with torchao quantization specified by torchao_config. |
| 290 | |
| 291 | Args: |
| 292 | param: weight parameter of the linear module |
| 293 | torchao_config: type of quantization and their arguments we want to |
| 294 | use to quantize the Tensor |
| 295 | |
| 296 | """ |
| 297 | from torchao.core.config import AOBaseConfig |
| 298 | from torchao.quantization import quantize_ |
| 299 | |
| 300 | assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}" |
| 301 | _check_torchao_fp8_activation_capability(torchao_config) |
| 302 | """ |
| 303 | Avoid real weight allocation for faster load, since we will |
| 304 | end up setting it to param. |
| 305 | """ |
| 306 | with torch.device("meta"): |
| 307 | # linear can't be top level module since quantize_ is inplace |
| 308 | # while some of our configs need to do module swap, and only non-top |
| 309 | # level modules support module swap |
| 310 | dummy_linear = torch.nn.Sequential( |
| 311 | torch.nn.Linear(param.shape[1], param.shape[0], bias=False) |
| 312 | ) |
| 313 | |
| 314 | dummy_linear[0].weight = param |
| 315 | quantize_(dummy_linear, torchao_config) |
| 316 | return dummy_linear[0].weight |
| 317 | |
| 318 | |
| 319 | class TorchAOLinearMethod(LinearMethodBase): |
| 320 | """Linear method for torchao. |
| 321 | |
| 322 | Args: |
| 323 | quant_config: The torchao quantization config, a string that encodes |
| 324 | the type of quantization and all relevant arguments. |
| 325 | |
| 326 | """ |
| 327 | |
| 328 | def __init__(self, quant_config: TorchAOConfig): |
| 329 | self.quant_config = quant_config |
| 330 | |
| 331 | def create_weights( |
| 332 | self, |
| 333 | layer: torch.nn.Module, |
| 334 | input_size_per_partition: int, |
| 335 | output_partition_sizes: list[int], |
| 336 | input_size: int, |
| 337 | output_size: int, |
| 338 | params_dtype: torch.dtype, |
| 339 | **extra_weight_attrs, |
| 340 | ): |
| 341 | weight = Parameter( |
| 342 | torch.empty( |
| 343 | sum(output_partition_sizes), |
| 344 | input_size_per_partition, |
| 345 | dtype=params_dtype, |
| 346 | ), |
| 347 | requires_grad=False, |
| 348 | ) |
| 349 | if self.quant_config.is_checkpoint_torchao_serialized: |
| 350 | weight = torchao_quantize_param_data( |
| 351 | weight, self.quant_config.torchao_config |
| 352 | ) |
| 353 | |
| 354 | set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) |
| 355 | |
| 356 | layer.register_parameter("weight", weight) |
| 357 | set_weight_attrs(weight, extra_weight_attrs) |
| 358 | |
| 359 | def apply( |
| 360 | self, |
| 361 | layer: torch.nn.Module, |
| 362 | x: torch.Tensor, |
| 363 | bias: torch.Tensor | None = None, |
| 364 | ) -> torch.Tensor: |
| 365 | return F.linear(x, layer.weight, bias) |
| 366 | |
| 367 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 368 | if self.quant_config.is_checkpoint_torchao_serialized: |
| 369 | if not hasattr(layer, "weight"): |
| 370 | return |
| 371 | |
| 372 | # record attributes attached to the weight, so we can |
| 373 | # recover later |
| 374 | recorded_weight_attr = _get_weight_attrs(layer.weight) |
| 375 | |
| 376 | layer.weight = Parameter( |
| 377 | convert_to_packed_tensor_based_on_current_hardware(layer.weight), |
| 378 | requires_grad=layer.weight.requires_grad, |
| 379 | ) |
| 380 | |
| 381 | _restore_weight_attrs(layer.weight, recorded_weight_attr) |
| 382 | return |
| 383 | |
| 384 | # online quantize the weight if the checkpoint is not already |
| 385 | # quantized by torchao |
| 386 | recorded_weight_attr = _get_weight_attrs(layer.weight) |
| 387 | |
| 388 | weight = torchao_quantize_param_data( |
| 389 | layer.weight, self.quant_config.torchao_config |
| 390 | ) |
| 391 | weight = torch.nn.Parameter( |
| 392 | convert_to_packed_tensor_based_on_current_hardware(weight), |
| 393 | weight.requires_grad, |
| 394 | ) |
| 395 | |
| 396 | _restore_weight_attrs(weight, recorded_weight_attr) |
| 397 | layer.register_parameter("weight", weight) |