gptq.py 13.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

CHU Tianxiang's avatar
CHU Tianxiang committed
4
5
import enum
from enum import Enum
6
from fractions import Fraction
7
from typing import TYPE_CHECKING, Any, Union
CHU Tianxiang's avatar
CHU Tianxiang committed
8
9

import torch
10
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
CHU Tianxiang's avatar
CHU Tianxiang committed
11
12
from torch.nn.parameter import Parameter

13
from vllm import _custom_ops as ops
14
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
15
from vllm.model_executor.layers.linear import LinearMethodBase
CHU Tianxiang's avatar
CHU Tianxiang committed
16
from vllm.model_executor.layers.quantization.base_config import (
17
18
19
    QuantizationConfig,
    QuantizeMethodBase,
)
20
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
21
22
23
24
25
26
27
28
29
    get_linear_quant_method,
)
from vllm.model_executor.parameter import (
    ChannelQuantScaleParameter,
    GroupQuantScaleParameter,
    PackedColumnParameter,
    PackedvLLMParameter,
    RowvLLMParameter,
)
30
from vllm.transformers_utils.config import get_safetensors_params_metadata
31
from vllm.utils.collection_utils import is_list_of
CHU Tianxiang's avatar
CHU Tianxiang committed
32

33
34
if TYPE_CHECKING:
    from vllm.model_executor.layers.quantization import QuantizationMethods
35
    from vllm.model_executor.models.utils import WeightsMapper
36
37
38
else:
    QuantizationMethods = str

CHU Tianxiang's avatar
CHU Tianxiang committed
39
40
41
42
43
44
45
46
47
48
49
50

class GPTQConfig(QuantizationConfig):
    """Config class for GPTQ.

    Reference: https://arxiv.org/abs/2210.17323
    """

    def __init__(
        self,
        weight_bits: int,
        group_size: int,
        desc_act: bool,
51
        lm_head_quantized: bool,
52
        dynamic: dict[str, dict[str, int | bool]],
53
        autoround_version: str = "",
54
        modules_in_block_to_quantize: list[str] | None = None,
CHU Tianxiang's avatar
CHU Tianxiang committed
55
    ) -> None:
56
57
        # GPTQModel use `dynamic` config property to allow per module
        # quantization config so each module can be individually optimized.
58
        # Format is dict[str, dict] where key is a regex string that can
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
        # perform both positive ("+:" prefixed) or negative ("-:" prefixed)
        # matching of a module.
        # Default to positive match, override base quant config mode, if no
        # prefix is used. Value is in dict format of field key and override
        # value.
        # Negative matching will skip quantization init for this module
        # entirely:
        # non-quantized inference. More details and quantization examples can be
        # found at: https://github.com/ModelCloud/GPTQModel
        # Example:
        #  # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9
        #  # last 1/4 of the layers 16-21 has 8bit and group_size 64
        # dynamic = {
        #  #`.*\.` matches the layers_node prefix
        #  # positive match layer 10-15
        #  r"+:.*\.(?:1[0-5])\..*": {"bits": 8,},
        #  # positive match layer 16-21
        #  r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,},
        #  r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers
        # }
79
        super().__init__()
80
81
        self.dynamic = dynamic

CHU Tianxiang's avatar
CHU Tianxiang committed
82
83
84
        self.weight_bits = weight_bits
        self.group_size = group_size
        self.desc_act = desc_act
85
        self.lm_head_quantized = lm_head_quantized
86
87
        self.pack_factor = Fraction(32, self.weight_bits)
        if self.weight_bits not in [2, 3, 4, 8]:
CHU Tianxiang's avatar
CHU Tianxiang committed
88
            raise ValueError(
89
                "Currently, only 2/3/4/8-bit weight quantization is "
90
91
                f"supported for GPTQ, but got {self.weight_bits} bits."
            )
CHU Tianxiang's avatar
CHU Tianxiang committed
92

93
94
        self.modules_in_block_to_quantize = modules_in_block_to_quantize or []

95
96
97
        # used to identify GPTQ model quantized by autoround
        self.autoround_version = autoround_version

CHU Tianxiang's avatar
CHU Tianxiang committed
98
    def __repr__(self) -> str:
99
100
101
102
103
104
105
106
        return (
            f"GPTQConfig(weight_bits={self.weight_bits}, "
            f"group_size={self.group_size}, "
            f"desc_act={self.desc_act}), "
            f"lm_head_quantized={self.lm_head_quantized}, "
            f"dynamic={self.dynamic}, "
            f"modules_in_block_to_quantize={self.modules_in_block_to_quantize})"
        )
CHU Tianxiang's avatar
CHU Tianxiang committed
107
108

    @classmethod
109
    def get_name(cls) -> QuantizationMethods:
CHU Tianxiang's avatar
CHU Tianxiang committed
110
111
112
        return "gptq"

    @classmethod
113
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
CHU Tianxiang's avatar
CHU Tianxiang committed
114
115
116
117
118
119
120
121
        return [torch.half]

    @classmethod
    # Need to figure it out
    def get_min_capability(cls) -> int:
        return 60

    @classmethod
122
    def get_config_filenames(cls) -> list[str]:
CHU Tianxiang's avatar
CHU Tianxiang committed
123
124
125
        return ["quantize_config.json"]

    @classmethod
126
    def from_config(cls, config: dict[str, Any]) -> "GPTQConfig":
127
128
129
        dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})
        dynamic = {} if dynamic is None else dynamic

CHU Tianxiang's avatar
CHU Tianxiang committed
130
131
132
        weight_bits = cls.get_from_keys(config, ["bits"])
        group_size = cls.get_from_keys(config, ["group_size"])
        desc_act = cls.get_from_keys(config, ["desc_act"])
133
134
135
136
        lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
        autoround_version = cls.get_from_keys_or(
            config, ["autoround_version"], default=""
        )
137
        modules_in_block_to_quantize = cls.get_from_keys_or(
138
139
140
141
142
143
144
145
146
147
148
            config, ["modules_in_block_to_quantize"], default=None
        )
        return cls(
            weight_bits,
            group_size,
            desc_act,
            lm_head_quantized,
            dynamic,
            autoround_version,
            modules_in_block_to_quantize,
        )
CHU Tianxiang's avatar
CHU Tianxiang committed
149

150
151
    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
152
    ) -> Union["GPTQLinearMethod", "QuantizeMethodBase"] | None:
153
154
155
156
157
158
159
160
161
162
163
        if isinstance(layer, FusedMoE):
            # GPTQ MoE support: fall back to MoeWNA16 for broad compatibility
            from .moe_wna16 import MoeWNA16Config

            config = {
                "quant_method": "gptq",
                "bits": self.weight_bits,
                "group_size": self.group_size,
                "sym": True,  # GPTQ typically uses symmetric quantization
                "lm_head": False,
            }
164
            return MoeWNA16Config.from_config(config).get_quant_method(layer, prefix)
165

166
        return get_linear_quant_method(self, layer, prefix, GPTQLinearMethod)
CHU Tianxiang's avatar
CHU Tianxiang committed
167

168
    def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
169
170
        if self.modules_in_block_to_quantize is not None:
            self.modules_in_block_to_quantize = hf_to_vllm_mapper.apply_list(
171
172
                self.modules_in_block_to_quantize
            )
173

174
    def maybe_update_config(self, model_name: str, revision: str | None = None):
175
176
177
178
179
        if self.modules_in_block_to_quantize:
            if is_list_of(self.modules_in_block_to_quantize, list):
                # original modules_in_block_to_quantize: list[list[str]]
                # flatten original modules_in_block_to_quantize
                self.modules_in_block_to_quantize = [
180
181
                    item
                    for sublist in self.modules_in_block_to_quantize
182
183
184
185
186
                    for item in sublist
                ]
            return

        unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
187
        metadata = get_safetensors_params_metadata(model_name, revision=revision)
188
189
190
        quant_layers: set[str] = {
            param_name.rsplit(".", 1)[0]
            for param_name, info in metadata.items()
191
            if (dtype := info.get("dtype", None))
192
193
194
195
            and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
        }
        self.modules_in_block_to_quantize = list(quant_layers)

CHU Tianxiang's avatar
CHU Tianxiang committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214

class ExllamaState(Enum):
    UNUSED = enum.auto()
    UNINITIALIZED = enum.auto()
    READY = enum.auto()


class GPTQLinearMethod(LinearMethodBase):
    """Linear method for GPTQ.

    Args:
        quant_config: The GPTQ quantization config.
    """

    def __init__(self, quant_config: GPTQConfig):
        self.quant_config = quant_config

    def create_weights(
        self,
215
        layer: torch.nn.Module,
CHU Tianxiang's avatar
CHU Tianxiang committed
216
        input_size_per_partition: int,
217
        output_partition_sizes: list[int],
CHU Tianxiang's avatar
CHU Tianxiang committed
218
219
220
        input_size: int,
        output_size: int,
        params_dtype: torch.dtype,
221
222
        **extra_weight_attrs,
    ):
CHU Tianxiang's avatar
CHU Tianxiang committed
223
        del output_size  # Unused.
224
        weight_loader = extra_weight_attrs.get("weight_loader")
CHU Tianxiang's avatar
CHU Tianxiang committed
225
226
227
228
        if input_size_per_partition % self.quant_config.group_size != 0:
            raise ValueError(
                "The input size is not aligned with the quantized "
                "weight shape. This can be caused by too large "
229
230
                "tensor parallel size."
            )
James Fleming's avatar
James Fleming committed
231
        output_size_per_partition = sum(output_partition_sizes)
232
        if output_size_per_partition % self.quant_config.pack_factor.numerator != 0:
CHU Tianxiang's avatar
CHU Tianxiang committed
233
234
235
            raise ValueError(
                "The output size is not aligned with the quantized "
                "weight shape. This can be caused by too large "
236
237
                "tensor parallel size."
            )
CHU Tianxiang's avatar
CHU Tianxiang committed
238
239
240
241
242
243
244
245

        if self.quant_config.group_size != -1:
            group_size = self.quant_config.group_size
        else:
            group_size = input_size
        exllama_state = ExllamaState.UNINITIALIZED
        scale_and_zero_size = input_size // group_size
        scale_and_zero_input_dim = None
246
247
248
249
        if (
            input_size != input_size_per_partition
            and self.quant_config.group_size != -1
        ):
CHU Tianxiang's avatar
CHU Tianxiang committed
250
251
252
253
254
255
256
257
            # For act-order models, we cannot use Exllama for row parallel layer
            if self.quant_config.desc_act:
                exllama_state = ExllamaState.UNUSED
            else:
                # we need to partition qzeros and scales for exllama kernel
                scale_and_zero_size = input_size_per_partition // group_size
                scale_and_zero_input_dim = 0

258
259
        qweight = PackedvLLMParameter(
            data=torch.empty(
CHU Tianxiang's avatar
CHU Tianxiang committed
260
261
262
263
                input_size_per_partition // self.quant_config.pack_factor,
                output_size_per_partition,
                dtype=torch.int32,
            ),
264
265
266
267
            input_dim=0,
            output_dim=1,
            packed_dim=0,
            packed_factor=self.quant_config.pack_factor,
268
269
270
271
272
273
274
275
276
277
278
279
280
281
            weight_loader=weight_loader,
        )

        g_idx = RowvLLMParameter(
            data=torch.tensor(
                [
                    i // self.quant_config.group_size
                    for i in range(input_size_per_partition)
                ],
                dtype=torch.int32,
            ),
            input_dim=0,
            weight_loader=weight_loader,
        )
282
        qzeros_args = {
283
            "data": torch.empty(
CHU Tianxiang's avatar
CHU Tianxiang committed
284
285
286
287
                scale_and_zero_size,
                output_size_per_partition // self.quant_config.pack_factor,
                dtype=torch.int32,
            ),
288
            "weight_loader": weight_loader,
289
290
        }
        weight_scale_args = {
291
            "data": torch.empty(
CHU Tianxiang's avatar
CHU Tianxiang committed
292
293
294
295
                scale_and_zero_size,
                output_size_per_partition,
                dtype=params_dtype,
            ),
296
            "weight_loader": weight_loader,
297
298
        }
        if scale_and_zero_input_dim is None:
299
            scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
300
301
302
303
            qzeros = PackedColumnParameter(
                output_dim=1,
                packed_dim=1,
                packed_factor=self.quant_config.pack_factor,
304
305
                **qzeros_args,
            )
306
307

        else:
308
309
310
            scales = GroupQuantScaleParameter(
                output_dim=1, input_dim=0, **weight_scale_args
            )
311
312
313
314
315
            qzeros = PackedvLLMParameter(
                input_dim=0,
                output_dim=1,
                packed_dim=1,
                packed_factor=self.quant_config.pack_factor,
316
317
                **qzeros_args,
            )
318
319
320
321
322
323
324

        layer.register_parameter("qweight", qweight)
        layer.register_parameter("g_idx", g_idx)
        layer.register_parameter("qzeros", qzeros)
        layer.register_parameter("scales", scales)

        layer.exllama_state = exllama_state
CHU Tianxiang's avatar
CHU Tianxiang committed
325

326
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
327
328
329
330
        # for torch.compile
        layer.qzeros = Parameter(layer.qzeros.data, requires_grad=False)
        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
        layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
331
        layer.scales = Parameter(layer.scales.data, requires_grad=False)
332

CHU Tianxiang's avatar
CHU Tianxiang committed
333
334
        # exllama needs to shuffle the weight after the weight is loaded
        # here we do the shuffle on first forward pass
335
        if layer.exllama_state == ExllamaState.UNINITIALIZED:
CHU Tianxiang's avatar
CHU Tianxiang committed
336
            if self.quant_config.desc_act:
337
                layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
CHU Tianxiang's avatar
CHU Tianxiang committed
338
            else:
339
340
341
                layer.g_idx.data = torch.empty(
                    (0,), dtype=torch.int, device=layer.g_idx.device
                )
342
            layer.exllama_state = ExllamaState.READY
343
344
345
346
347
348
            ops.gptq_shuffle(layer.qweight, layer.g_idx, self.quant_config.weight_bits)

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
349
        bias: torch.Tensor | None = None,
350
351
    ) -> torch.Tensor:
        out_shape = x.shape[:-1] + (layer.qweight.shape[-1],)
352
353
        reshaped_x = x.reshape(-1, x.shape[-1])

354
355
356
357
358
359
360
361
362
        output = ops.gptq_gemm(
            reshaped_x,
            layer.qweight,
            layer.qzeros,
            layer.scales,
            layer.g_idx,
            layer.exllama_state == ExllamaState.READY,
            self.quant_config.weight_bits,
        )
CHU Tianxiang's avatar
CHU Tianxiang committed
363
        if bias is not None:
364
            output.add_(bias)
CHU Tianxiang's avatar
CHU Tianxiang committed
365
        return output.reshape(out_shape)