gptq.py 11.5 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 Any, Optional, Union
CHU Tianxiang's avatar
CHU Tianxiang committed
8
9
10
11

import torch
from torch.nn.parameter import Parameter

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


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,
38
        lm_head_quantized: bool,
39
        dynamic: dict[str, dict[str, Union[int, bool]]],
40
        autoround_version: str = "",
CHU Tianxiang's avatar
CHU Tianxiang committed
41
    ) -> None:
42
43
        # GPTQModel use `dynamic` config property to allow per module
        # quantization config so each module can be individually optimized.
44
        # Format is dict[str, dict] where key is a regex string that can
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
        # 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
        # }
65
        super().__init__()
66
67
        self.dynamic = dynamic

CHU Tianxiang's avatar
CHU Tianxiang committed
68
69
70
        self.weight_bits = weight_bits
        self.group_size = group_size
        self.desc_act = desc_act
71
        self.lm_head_quantized = lm_head_quantized
72
73
        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
74
            raise ValueError(
75
76
                "Currently, only 2/3/4/8-bit weight quantization is "
                f"supported for GPTQ, but got {self.weight_bits} bits.")
CHU Tianxiang's avatar
CHU Tianxiang committed
77

78
79
80
        # used to identify GPTQ model quantized by autoround
        self.autoround_version = autoround_version

CHU Tianxiang's avatar
CHU Tianxiang committed
81
82
83
    def __repr__(self) -> str:
        return (f"GPTQConfig(weight_bits={self.weight_bits}, "
                f"group_size={self.group_size}, "
84
                f"desc_act={self.desc_act}), "
85
86
                f"lm_head_quantized={self.lm_head_quantized}), "
                f"dynamic={self.dynamic}")
CHU Tianxiang's avatar
CHU Tianxiang committed
87
88

    @classmethod
89
    def get_name(cls) -> QuantizationMethods:
CHU Tianxiang's avatar
CHU Tianxiang committed
90
91
92
        return "gptq"

    @classmethod
93
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
CHU Tianxiang's avatar
CHU Tianxiang committed
94
95
96
97
98
99
100
101
        return [torch.half]

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

    @classmethod
102
    def get_config_filenames(cls) -> list[str]:
CHU Tianxiang's avatar
CHU Tianxiang committed
103
104
105
        return ["quantize_config.json"]

    @classmethod
106
    def from_config(cls, config: dict[str, Any]) -> "GPTQConfig":
107
108
109
        dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})
        dynamic = {} if dynamic is None else dynamic

CHU Tianxiang's avatar
CHU Tianxiang committed
110
111
112
        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"])
113
114
        lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
                                                 default=False)
115
116
        autoround_version = cls.get_from_keys_or(config, ["autoround_version"],
                                                 default="")
117
        return cls(weight_bits, group_size, desc_act, lm_head_quantized,
118
                   dynamic, autoround_version)
CHU Tianxiang's avatar
CHU Tianxiang committed
119

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
    ) -> Optional[Union["GPTQLinearMethod", "QuantizeMethodBase"]]:
        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,
            }
            return MoeWNA16Config.from_config(config).get_quant_method(
                layer, prefix)

137
        return get_linear_quant_method(self, layer, prefix, GPTQLinearMethod)
CHU Tianxiang's avatar
CHU Tianxiang committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158


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,
159
        layer: torch.nn.Module,
CHU Tianxiang's avatar
CHU Tianxiang committed
160
        input_size_per_partition: int,
161
        output_partition_sizes: list[int],
CHU Tianxiang's avatar
CHU Tianxiang committed
162
163
164
        input_size: int,
        output_size: int,
        params_dtype: torch.dtype,
165
166
        **extra_weight_attrs,
    ):
CHU Tianxiang's avatar
CHU Tianxiang committed
167
        del output_size  # Unused.
168
        weight_loader = extra_weight_attrs.get("weight_loader")
CHU Tianxiang's avatar
CHU Tianxiang committed
169
170
171
172
173
        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 "
                "tensor parallel size.")
James Fleming's avatar
James Fleming committed
174
        output_size_per_partition = sum(output_partition_sizes)
175
176
        if (output_size_per_partition % self.quant_config.pack_factor.numerator
                != 0):
CHU Tianxiang's avatar
CHU Tianxiang committed
177
178
179
180
181
182
183
184
185
186
187
188
            raise ValueError(
                "The output size is not aligned with the quantized "
                "weight shape. This can be caused by too large "
                "tensor parallel size.")

        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
189
190
        if (input_size != input_size_per_partition
                and self.quant_config.group_size != -1):
CHU Tianxiang's avatar
CHU Tianxiang committed
191
192
193
194
195
196
197
198
            # 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

199
200
        qweight = PackedvLLMParameter(
            data=torch.empty(
CHU Tianxiang's avatar
CHU Tianxiang committed
201
202
203
204
                input_size_per_partition // self.quant_config.pack_factor,
                output_size_per_partition,
                dtype=torch.int32,
            ),
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
            input_dim=0,
            output_dim=1,
            packed_dim=0,
            packed_factor=self.quant_config.pack_factor,
            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)
        qzeros_args = {
            "data":
CHU Tianxiang's avatar
CHU Tianxiang committed
222
223
224
225
226
            torch.empty(
                scale_and_zero_size,
                output_size_per_partition // self.quant_config.pack_factor,
                dtype=torch.int32,
            ),
227
228
229
230
231
            "weight_loader":
            weight_loader
        }
        weight_scale_args = {
            "data":
CHU Tianxiang's avatar
CHU Tianxiang committed
232
233
234
235
236
            torch.empty(
                scale_and_zero_size,
                output_size_per_partition,
                dtype=params_dtype,
            ),
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
            "weight_loader":
            weight_loader
        }
        if scale_and_zero_input_dim is None:
            scales = ChannelQuantScaleParameter(output_dim=1,
                                                **weight_scale_args)
            qzeros = PackedColumnParameter(
                output_dim=1,
                packed_dim=1,
                packed_factor=self.quant_config.pack_factor,
                **qzeros_args)

        else:
            scales = GroupQuantScaleParameter(output_dim=1,
                                              input_dim=0,
                                              **weight_scale_args)
            qzeros = PackedvLLMParameter(
                input_dim=0,
                output_dim=1,
                packed_dim=1,
                packed_factor=self.quant_config.pack_factor,
                **qzeros_args)
259
260
261
262
263
264
265

        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
266

267
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
268
269
270
271
        # 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)
272
        layer.scales = Parameter(layer.scales.data, requires_grad=False)
273

CHU Tianxiang's avatar
CHU Tianxiang committed
274
275
        # exllama needs to shuffle the weight after the weight is loaded
        # here we do the shuffle on first forward pass
276
        if layer.exllama_state == ExllamaState.UNINITIALIZED:
CHU Tianxiang's avatar
CHU Tianxiang committed
277
            if self.quant_config.desc_act:
278
                layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
CHU Tianxiang's avatar
CHU Tianxiang committed
279
            else:
280
                layer.g_idx.data = torch.empty((0, ),
281
                                               dtype=torch.int,
282
283
284
                                               device=layer.g_idx.device)
            layer.exllama_state = ExllamaState.READY
            ops.gptq_shuffle(layer.qweight, layer.g_idx,
285
                             self.quant_config.weight_bits)
286
287
288
289
290
291
292
293

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

294
295
296
        output = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
                               layer.scales, layer.g_idx,
                               layer.exllama_state == ExllamaState.READY,
297
                               self.quant_config.weight_bits)
CHU Tianxiang's avatar
CHU Tianxiang committed
298
        if bias is not None:
299
            output.add_(bias)
CHU Tianxiang's avatar
CHU Tianxiang committed
300
        return output.reshape(out_shape)