| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Mapping |
| 4 | from copy import deepcopy |
| 5 | from types import MappingProxyType |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import regex as re |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod |
| 12 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 13 | ParallelLMHead, |
| 14 | UnquantizedEmbeddingMethod, |
| 15 | ) |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from ..auto_gptq import AutoGPTQConfig |
| 19 | else: |
| 20 | AutoGPTQConfig = object |
| 21 | |
| 22 | |
| 23 | def normalize_and_validate_gptq_desc_act( |
| 24 | desc_act: bool, |
| 25 | group_size: int | None, |
| 26 | dynamic: Mapping[str, Any], |
| 27 | ) -> bool: |
| 28 | base_desc_act = desc_act |
| 29 | if desc_act and group_size == -1: |
| 30 | # With one group per output channel, activation ordering is a no-op. |
| 31 | desc_act = False |
| 32 | if desc_act: |
| 33 | raise ValueError( |
| 34 | "GPTQ group activation ordering (desc_act=True) is no longer " |
| 35 | "supported. Use a checkpoint with static activation ordering " |
| 36 | "or desc_act=False." |
| 37 | ) |
| 38 | |
| 39 | for pattern, overrides in dynamic.items(): |
| 40 | if pattern.startswith("-:") or not isinstance(overrides, Mapping): |
| 41 | continue |
| 42 | dynamic_desc_act = overrides.get("desc_act", base_desc_act) |
| 43 | dynamic_group_size = overrides.get("group_size", group_size) |
| 44 | if dynamic_desc_act and dynamic_group_size != -1: |
| 45 | raise ValueError( |
| 46 | "GPTQ group activation ordering (desc_act=True) is no longer " |
| 47 | f"supported, but is enabled by dynamic rule {pattern!r}. Use " |
| 48 | "a checkpoint with static activation ordering or desc_act=False." |
| 49 | ) |
| 50 | return desc_act |
| 51 | |
| 52 | |
| 53 | # Match dynamic rules with module name (prefix) and override quantize |
| 54 | # config if module (prefix) matches a rule |
| 55 | def override_config(config: AutoGPTQConfig, prefix: str): |
| 56 | weight_bits = get_dynamic_override(config, prefix, "bits", config.weight_bits) |
| 57 | if isinstance(weight_bits, int): |
| 58 | config.weight_bits = weight_bits |
| 59 | group_size = get_dynamic_override(config, prefix, "group_size", config.group_size) |
| 60 | if isinstance(group_size, int): |
| 61 | config.group_size = group_size |
| 62 | desc_act = get_dynamic_override(config, prefix, "desc_act", config.desc_act) |
| 63 | if isinstance(desc_act, bool): |
| 64 | config.desc_act = desc_act |
| 65 | |
| 66 | if config.desc_act and config.group_size == -1: |
| 67 | # Activation ordering is a no-op for channelwise quantization. |
| 68 | config.desc_act = False |
| 69 | |
| 70 | config.pack_factor = 32 // config.weight_bits # packed into int32 |
| 71 | assert isinstance(config, AutoGPTQConfig) |
| 72 | is_sym = get_dynamic_override(config, prefix, "sym", config.is_sym) |
| 73 | if isinstance(is_sym, bool): |
| 74 | config.is_sym = is_sym |
| 75 | |
| 76 | if (config.weight_bits, config.is_sym) not in config.TYPE_MAP: |
| 77 | raise ValueError( |
| 78 | "Unsupported quantization config: " |
| 79 | f"bits={config.weight_bits}, sym={config.is_sym}" |
| 80 | ) |
| 81 | |
| 82 | config.quant_type = config.TYPE_MAP[(config.weight_bits, config.is_sym)] |
| 83 | |
| 84 | |
| 85 | def get_dynamic_override( |
| 86 | config: AutoGPTQConfig, |
| 87 | layer_name: str, |
| 88 | key: str | None = None, |
| 89 | default_value: int | bool | None = None, |
| 90 | ) -> dict | int | bool | None: |
| 91 | for pattern, pattern_dict in config.dynamic.items(): |
| 92 | # Negative match: matched modules are excluded from quantized init |
| 93 | if pattern.startswith("-:"): |
| 94 | if re.match(pattern.removeprefix("-:"), layer_name): |
| 95 | return False |
| 96 | # Positive match: matched modules have quant properties overrides |
| 97 | # base quant config |
| 98 | elif re.match(pattern.removeprefix("+:"), layer_name): |
| 99 | if key is None: |
| 100 | return pattern_dict |
| 101 | else: |
| 102 | return pattern_dict.get(key, default_value) |
| 103 | return default_value |
| 104 | |
| 105 | |
| 106 | def flatten_list(lst: list[Any]) -> list[Any]: |
| 107 | output = [] |
| 108 | |
| 109 | def _flatten(lst: list[Any]): |
| 110 | for i in lst: |
| 111 | if isinstance(i, list): |
| 112 | _flatten(i) |
| 113 | else: |
| 114 | output.append(i) |
| 115 | |
| 116 | _flatten(lst) |
| 117 | return output |
| 118 | |
| 119 | |
| 120 | def is_layer_gptq_quantized( |
| 121 | prefix: str, |
| 122 | quantized_layers: list[str], |
| 123 | fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), |
| 124 | ) -> bool: |
| 125 | # prefix: model.layers.0.self_attn.q_proj |
| 126 | # proj_name: q_proj |
| 127 | |
| 128 | # GPTQ's `modules_in_block_to_quantize`: |
| 129 | # Substr: ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"] |
| 130 | # Full prefix ["model.layers.0.self_attn.q_proj"] |
| 131 | |
| 132 | proj_name = prefix.split(".")[-1] |
| 133 | |
| 134 | quantized_layers = flatten_list(quantized_layers) |
| 135 | |
| 136 | # Fused layers like gate_up_proj or qkv_proj will not be fused |
| 137 | # in the safetensors checkpoint. So, we convert the name |
| 138 | # from the fused version to unfused + check to make sure that |
| 139 | # each shard of the fused layer has the same scheme. |
| 140 | if proj_name in fused_mapping: |
| 141 | shard_prefixes = [ |
| 142 | prefix.replace(proj_name, shard_proj_name) |
| 143 | for shard_proj_name in fused_mapping[proj_name] |
| 144 | ] |
| 145 | |
| 146 | is_quantized = None |
| 147 | for shard_prefix in shard_prefixes: |
| 148 | is_shard_quantized = any( |
| 149 | layer in shard_prefix for layer in quantized_layers |
| 150 | ) |
| 151 | |
| 152 | if is_quantized is None: |
| 153 | is_quantized = is_shard_quantized |
| 154 | elif is_shard_quantized != is_quantized: |
| 155 | raise ValueError( |
| 156 | f"Detected some but not all shards of {prefix} " |
| 157 | "are quantized. All shards of fused layers " |
| 158 | "to have the same precision." |
| 159 | ) |
| 160 | else: |
| 161 | is_quantized = any(layer in prefix for layer in quantized_layers) |
| 162 | |
| 163 | assert is_quantized is not None |
| 164 | return is_quantized |
| 165 | |
| 166 | |
| 167 | def get_linear_quant_method( |
| 168 | config: AutoGPTQConfig, |
| 169 | layer: torch.nn.Module, |
| 170 | prefix: str, |
| 171 | linear_method_cls: type, |
| 172 | ): |
| 173 | cloned_config = deepcopy(config) |
| 174 | parallel_lm_head_quantized = ( |
| 175 | isinstance(layer, ParallelLMHead) and cloned_config.lm_head_quantized |
| 176 | ) |
| 177 | if isinstance(layer, LinearBase) or parallel_lm_head_quantized: |
| 178 | is_layer_quantized = is_layer_gptq_quantized( |
| 179 | prefix=prefix, |
| 180 | quantized_layers=cloned_config.modules_in_block_to_quantize, |
| 181 | fused_mapping=cloned_config.packed_modules_mapping, |
| 182 | ) |
| 183 | # False = skip module, None = no override, else = Positive match |
| 184 | if get_dynamic_override( # noqa: E712 |
| 185 | cloned_config, # noqa: E712 |
| 186 | layer_name=prefix, |
| 187 | ) == False or (not is_layer_quantized): # noqa: E712 |
| 188 | if parallel_lm_head_quantized: |
| 189 | return UnquantizedEmbeddingMethod() |
| 190 | return UnquantizedLinearMethod() |
| 191 | |
| 192 | if prefix: |
| 193 | # Dynamic per module/layer rules may override base config |
| 194 | override_config(cloned_config, prefix=prefix) |
| 195 | |
| 196 | return linear_method_cls(cloned_config) |
| 197 | return None |