quantizer.py 12.8 KB
Newer Older
Ji Lin's avatar
Ji Lin committed
1
import torch
Casper's avatar
Casper committed
2
3
4
5
6
7
8
9
10
11
12
13
14
import logging
import functools
import torch.nn as nn
from tqdm import tqdm
from collections import defaultdict
from awq.utils.utils import clear_memory
from awq.utils.calib_data import get_calib_dataset
from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
from awq.quantize.apply_quantized import apply_scale, apply_clip
from awq.utils.module import append_str_prefix, get_op_name, get_named_linears, set_op_by_name


class AwqQuantizer:
Casper Hansen's avatar
Casper Hansen committed
15
16
    def __init__(self, awq_model, model, tokenizer, w_bit, group_size, version, calib_data, split, text_column) -> None:
        self.awq_model = awq_model
Casper's avatar
Casper committed
17
18
19
20
21
22
23
24
        self.model = model
        self.tokenizer = tokenizer
        self.w_bit = w_bit
        self.group_size = group_size
        self.version = version
        self.calib_data = calib_data
        self.split = split
        self.text_column = text_column
Casper Hansen's avatar
Casper Hansen committed
25
        self.modules, self.module_kwargs, self.inps = self.init_quant()
26
    
Casper's avatar
Casper committed
27
28
29
30
31
32
33
34
    def pseudo_quantize_tensor(self, w: torch.Tensor, get_scale_zp=False):
        org_w_shape = w.shape
        if self.group_size > 0:
            assert org_w_shape[-1] % self.group_size == 0
            w = w.reshape(-1, self.group_size)
        assert w.dim() == 2

        # zero point quantization
Ji Lin's avatar
Ji Lin committed
35
36
        max_val = w.amax(dim=1, keepdim=True)
        min_val = w.amin(dim=1, keepdim=True)
Casper's avatar
Casper committed
37
        max_int = 2 ** self.w_bit - 1
Ji Lin's avatar
Ji Lin committed
38
39
40
        min_int = 0
        scales = (max_val - min_val).clamp(min=1e-5) / max_int
        zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
Casper's avatar
Casper committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54

        assert torch.isnan(scales).sum() == 0
        assert torch.isnan(w).sum() == 0

        w = (torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros) * scales
        assert torch.isnan(w).sum() == 0

        w = w.reshape(org_w_shape)

        if get_scale_zp:
            return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1)
        else:
            return w
    
Casper Hansen's avatar
Casper Hansen committed
55
56
    def quantize(self):
        for i in tqdm(range(len(self.modules)), desc="AWQ"):
Casper's avatar
Casper committed
57
58
59
60
61
62
63
            # [STEP 1]: Get layer, extract linear modules, extract input features
            self.modules[i] = self.modules[i].cuda()
            named_linears = get_named_linears(self.modules[i])
            input_feat = self._get_input_feat(self.modules[i], named_linears)
            clear_memory()

            # [STEP 2]: Compute and apply scale list
Casper Hansen's avatar
Casper Hansen committed
64
            module_config: list[dict] = self.awq_model.get_layers_for_scaling(
Casper's avatar
Casper committed
65
66
                self.modules[i], input_feat, self.module_kwargs
            )
Casper Hansen's avatar
Casper Hansen committed
67
            scales_list = [self._search_best_scale(self.modules[i], **layer) for layer in module_config]
Casper's avatar
Casper committed
68
69
70
71
            apply_scale(self.modules[i], scales_list, input_feat_dict=input_feat)
            scales_list = append_str_prefix(scales_list, get_op_name(self.model, self.modules[i]) + ".")

            # [STEP 3]: Compute and apply clipping list
Casper Hansen's avatar
Casper Hansen committed
72
73
74
            clip_list = self._search_best_clip(self.modules[i], named_linears, input_feat)
            apply_clip(self.modules[i], clip_list)
            clip_list = append_str_prefix(clip_list, get_op_name(self.model, self.modules[i]) + ".")
Casper's avatar
Casper committed
75
76
77

            # [STEP 4]: Quantize weights
            for name, linear_layer in named_linears.items():
Casper Hansen's avatar
Casper Hansen committed
78
                # NOTE: small regression in perplexity if linear layer uses .cpu().float()
79
80
                linear_layer = linear_layer.cuda().half()

Casper's avatar
Casper committed
81
                linear_layer.weight.data, scales, zeros = self.pseudo_quantize_tensor(
82
                    linear_layer.weight.data, 
Casper's avatar
Casper committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
                    get_scale_zp=True
                )

                if self.version == 'GEMM':
                    scales = scales.t().contiguous()
                    zeros = zeros.t().contiguous()
                    q_linear_module = WQLinear_GEMM

                elif self.version  == 'GEMV':
                    q_linear_module = WQLinear_GEMV
                
                q_linear = q_linear_module.from_linear(
                    linear=linear_layer,
                    w_bit=self.w_bit,
                    group_size=self.group_size,
                    init_only=False,
                    scales=scales,
                    zeros=zeros
                )

                linear_layer.cpu()
                q_linear.to(next(self.modules[i].parameters()).device)
                set_op_by_name(self.modules[i], name, q_linear)
                clear_memory()
            
            clear_memory()

    @torch.no_grad()
Casper Hansen's avatar
Casper Hansen committed
111
112
113
114
115
116
117
118
    def _search_best_scale(self, module, prev_op, layers: list[nn.Linear], inp: torch.Tensor, module2inspect=None, kwargs={}):
        if module2inspect is None:
            assert len(layers) == 1
            module2inspect = layers[0]
        
        if "use_cache" in kwargs:
            kwargs.pop("use_cache")
        
Casper's avatar
Casper committed
119
        # Put x on the right device
Casper Hansen's avatar
Casper Hansen committed
120
        inp = inp.to(next(module2inspect.parameters()).device)
Casper's avatar
Casper committed
121
122

        # [STEP 1]: Compute maximum of weight
Casper Hansen's avatar
Casper Hansen committed
123
124
        weight = torch.cat([_m.weight for _m in layers], dim=0)
        org_shape = weight.shape
Casper's avatar
Casper committed
125
        weight = weight.view(-1, self.group_size)
Casper Hansen's avatar
Casper Hansen committed
126
127
128
        w_scale = weight.abs() / weight.abs().amax(dim=1, keepdim=True)
        w_scale = w_scale.view(org_shape)
        w_max = w_scale.mean(0)
Casper's avatar
Casper committed
129
130
131
        clear_memory(weight)

        # [STEP 2]: Compute maximum of x
Casper Hansen's avatar
Casper Hansen committed
132
        x_max = inp.abs().view(-1, inp.shape[-1]).mean(0)
Casper's avatar
Casper committed
133

Casper Hansen's avatar
Casper Hansen committed
134
        # [STEP 3]: Compute output of module
Casper's avatar
Casper committed
135
        with torch.no_grad():
Casper Hansen's avatar
Casper Hansen committed
136
            org_out = module2inspect(inp, **kwargs)
Casper's avatar
Casper committed
137
138
139
140
141
            if isinstance(org_out, tuple):
                org_out = org_out[0]
        
        # [STEP 4]: Compute loss
        best_scales = self._compute_best_scale(
Casper Hansen's avatar
Casper Hansen committed
142
143
            inp, w_max, x_max, module2inspect,
            layers, org_out, kwargs
Casper's avatar
Casper committed
144
145
        )
        
Casper Hansen's avatar
Casper Hansen committed
146
        return (get_op_name(module, prev_op), tuple([get_op_name(module, m) for m in layers]), best_scales)
Casper's avatar
Casper committed
147

Casper Hansen's avatar
Casper Hansen committed
148
    def _compute_best_scale(self, x, w_max, x_max, module2inspect, linears2scale: list[nn.Linear], org_out, kwargs={}):
Casper's avatar
Casper committed
149
150
151
        """
        Compute loss and select best scales

Casper's avatar
Casper committed
152
        L(s) = || Q(W * s) (s^-1 * X) - W * X ||
Casper's avatar
Casper committed
153
154
155
156
157
158
159
160
161
162
163
        Q: weight quantization function | pseudo_quantize_tensor(W * s)
        X: inputs from calib dataset    | X
        W: original weights in FP16     | layer
        s: per channel scaling factor   | s^-1 * X
        """
        n_grid = 20
        history = []
        best_ratio = -1
        best_scales = None
        best_error = float('inf')

Casper Hansen's avatar
Casper Hansen committed
164
        org_sd = {k: v.cpu() for k, v in module2inspect.state_dict().items()}
Casper's avatar
Casper committed
165
166
167
168
169
        
        device = x.device
        x_max = x_max.view(-1).to(device)
        w_max = w_max.view(-1).to(device)
        
Casper's avatar
Casper committed
170
171
        for ratio in range(n_grid):
            # create new scales
Casper's avatar
Casper committed
172
173
            ratio = ratio / n_grid
            scales = (x_max.pow(ratio) / w_max.pow(1-ratio)).clamp(min=1e-4)
Casper's avatar
Casper committed
174
            scales = scales / (scales.max() * scales.min()).sqrt()
Casper's avatar
Casper committed
175
176
            
            scales_view = scales.view(1, -1).to(device)
Casper's avatar
Casper committed
177
            for fc in linears2scale:
Casper's avatar
Casper committed
178
179
                fc.weight.mul_(scales_view)
                fc.weight.data = self.pseudo_quantize_tensor(fc.weight.data) / scales_view
Casper's avatar
Casper committed
180

Casper Hansen's avatar
Casper Hansen committed
181
            out = module2inspect(x, **kwargs)
Casper's avatar
Casper committed
182
183
184
185
186
187
            if isinstance(out, tuple):
                out = out[0]

            # measure loss and check if better than best
            loss = (org_out - out).float().pow(2).mean().item() # NOTE: float prevents overflow
            history.append(loss)
Casper's avatar
Casper committed
188
            if loss < best_error:
Casper's avatar
Casper committed
189
190
                best_error = loss
                best_ratio = ratio
Casper's avatar
Casper committed
191
                best_scales = scales.clone()
Casper Hansen's avatar
Casper Hansen committed
192
            module2inspect.load_state_dict(org_sd)
Casper's avatar
Casper committed
193

Casper's avatar
Casper committed
194
195
196
197
198
199
        if best_ratio == -1:
            logging.debug(history)
            raise Exception

        assert torch.isnan(best_scales).sum() == 0, best_scales

Casper Hansen's avatar
Casper Hansen committed
200
        return best_scales.detach().cpu()
Casper's avatar
Casper committed
201

Casper Hansen's avatar
Casper Hansen committed
202
203
204
205
    @torch.no_grad()
    def _search_best_clip(self, layer, named_linears, input_feat):
        clip_list = []
        avoid_clipping = ["q_", "k_", "query", "key", "Wqkv"]
Casper's avatar
Casper committed
206

Casper Hansen's avatar
Casper Hansen committed
207
208
209
210
211
212
213
214
215
216
        for name in named_linears:
            # due to qk bmm, it is hard to clip precisely
            if any([_ in name for _ in avoid_clipping]):
                continue

            named_linears[name].cuda()
            max_val = self._compute_best_clip(named_linears[name].weight, input_feat[name])
            clip_list.append((name, max_val))

            named_linears[name].cpu()
Casper Hansen's avatar
Casper Hansen committed
217
218
        
        return clip_list
Casper Hansen's avatar
Casper Hansen committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235

    @torch.no_grad()
    def _compute_best_clip(self, w: torch.Tensor, input_feat: torch.Tensor, n_grid=20, max_shrink=0.5, n_sample_token=512):
        assert w.dim() == 2
        org_w_shape = w.shape
        # w           [co, ci]      -> [co, 1, n_group, group size]
        # input_feat  [n_token, ci] -> [1, n_token, n_group, group size]
        group_size = self.group_size if self.group_size > 0 else w.shape[1]
        input_feat = input_feat.view(-1, input_feat.shape[-1])
        input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size)
        input_feat = input_feat[:, 0::input_feat.shape[1] // n_sample_token]
        w = w.reshape(w.shape[0], 1, -1, group_size)

        oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64  # prevent OOM
        assert w.shape[0] % oc_batch_size == 0
        w_all = w
        best_max_val_all = []
Casper's avatar
Casper committed
236

Casper Hansen's avatar
Casper Hansen committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
        for i_b in range(w.shape[0] // oc_batch_size):
            w = w_all[i_b * oc_batch_size: (i_b + 1) * oc_batch_size]

            org_max_val = w.abs().amax(dim=-1, keepdim=True)  # co, 1, n_group, 1

            best_max_val = org_max_val.clone()
            min_errs = torch.ones_like(org_max_val) * 1e9
            input_feat = input_feat.to(w.device)
            org_out = (input_feat * w).sum(dim=-1)  # co, n_token, n_group

            for i_s in range(int(max_shrink * n_grid)):
                max_val = org_max_val * (1 - i_s / n_grid)
                min_val = - max_val
                cur_w = torch.clamp(w, min_val, max_val)
                q_w = self.pseudo_quantize_tensor(cur_w)
                cur_out = (input_feat * q_w).sum(dim=-1)

                # co, 1, n_group, 1
                err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape)
                del cur_w
                del cur_out
                cur_best_idx = err < min_errs
                min_errs[cur_best_idx] = err[cur_best_idx]
                best_max_val[cur_best_idx] = max_val[cur_best_idx]
            best_max_val_all.append(best_max_val)

        best_max_val = torch.cat(best_max_val_all, dim=0)

        clear_memory(input_feat)
        clear_memory(org_out)

        return best_max_val.squeeze(1)

    def init_quant(self, n_samples=128, seqlen=512):
        modules = self.awq_model.get_model_layers(self.model)
Casper's avatar
Casper committed
272
273
274
275
276
277
278
279
280
        samples = get_calib_dataset(
            data=self.calib_data, tokenizer=self.tokenizer, n_samples=n_samples, block_size=seqlen,
            split=self.split, text_column=self.text_column
        )
        samples = torch.cat(samples, dim=0)

        inps = []
        layer_kwargs = {}

Casper Hansen's avatar
Casper Hansen committed
281
282
        modules[0] = modules[0].cuda()
        self.awq_model.move_embed(self.model, "cuda")
Casper's avatar
Casper committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
        
        # get input and kwargs to layer 0
        # with_kwargs is only supported in PyTorch 2.0
        # use this Catcher hack for now
        class Catcher(nn.Module):
            def __init__(self, module):
                super().__init__()
                self.module = module

            def forward(self, hijacked_inputs, **kwargs):
                inps.append(hijacked_inputs)
                layer_kwargs.update(kwargs)
                raise ValueError  # early exit to break later inference

        # patch layer 0 to catch input and kwargs
Casper Hansen's avatar
Casper Hansen committed
298
        modules[0] = Catcher(modules[0])
Casper's avatar
Casper committed
299
300
301
302
303
        try:
            self.model(samples.to(next(self.model.parameters()).device))
        except ValueError:  # work with early exit
            pass
        del samples
Casper Hansen's avatar
Casper Hansen committed
304
        modules[0] = modules[0].module  # restore
Casper's avatar
Casper committed
305
306
        inps = inps[0]

Casper Hansen's avatar
Casper Hansen committed
307
308
        modules[0] = modules[0].cpu()
        self.awq_model.move_embed(self.model, "cpu")
Casper's avatar
Casper committed
309
310
311
        
        clear_memory()

Casper Hansen's avatar
Casper Hansen committed
312
        return modules, layer_kwargs, inps
Casper's avatar
Casper committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
    
    def _get_input_feat(self, layer, named_linears):
        # firstly, get input features of all linear layers
        def cache_input_hook(m, x, y, name, feat_dict):
            x = x[0]
            x = x.detach().cpu()
            feat_dict[name].append(x)

        input_feat = defaultdict(list)
        handles = []
        for name in named_linears:
            handles.append(named_linears[name].register_forward_hook(
                functools.partial(cache_input_hook, name=name,
                                feat_dict=input_feat)))
Casper Hansen's avatar
Casper Hansen committed
327
        self.inps = self.inps.to(next(layer.parameters()).device)  # in case multi-gpu
Casper's avatar
Casper committed
328
        # get output as next layer's input
Casper Hansen's avatar
Casper Hansen committed
329
        self.inps = layer(self.inps, **self.module_kwargs)[0]
Casper's avatar
Casper committed
330
331
332
333
334
335
        for h in handles:
            h.remove()
        # now solve for scaling and clipping
        input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()}
        
        return input_feat