base.py 15.5 KB
Newer Older
1
import os
Casper Hansen's avatar
Casper Hansen committed
2
import gc
3
import json
Casper Hansen's avatar
Casper Hansen committed
4
import torch
Casper Hansen's avatar
Casper Hansen committed
5
import logging
Casper Hansen's avatar
Casper Hansen committed
6
7
import functools
import torch.nn as nn
Casper Hansen's avatar
Casper Hansen committed
8
from tqdm import tqdm
Casper Hansen's avatar
Casper Hansen committed
9
10
from collections import defaultdict

11
from awq.modules.act import ScaledActivation
12
from huggingface_hub import snapshot_download
Casper Hansen's avatar
Casper Hansen committed
13
from awq.utils.calib_data import get_calib_dataset
14
from awq.quantize.quantizer import pseudo_quantize_tensor
15
from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
Casper Hansen's avatar
Casper Hansen committed
16
17
from awq.quantize.auto_clip import auto_clip_block, apply_clip
from awq.quantize.auto_scale import auto_scale_block, apply_scale
18
19
from transformers import AutoModelForCausalLM, AutoConfig, PreTrainedModel
from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
Casper Hansen's avatar
Casper Hansen committed
20
from awq.utils.module import append_str_prefix, get_op_name, get_named_linears, set_op_by_name
Casper Hansen's avatar
Casper Hansen committed
21

22
class BaseAWQForCausalLM(nn.Module):
23
    def __init__(self, model, model_type, is_quantized, quant_config):
24
        super().__init__()
25
26
27
28
        self.model:PreTrainedModel = model
        self.model_type:str = model_type
        self.is_quantized:bool = is_quantized
        self.search_result = None
29
        self.quant_config:dict = quant_config
30
31
32
33
34
35
    
    def to(self, device: str):
        return self.model.to(device)
    
    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)
Casper Hansen's avatar
Casper Hansen committed
36
37
38
39
    
    def generate(self, *args, **kwargs):
        with torch.inference_mode():
            return self.model.generate(*args, **kwargs)
40

Casper Hansen's avatar
Casper Hansen committed
41
    @torch.no_grad()
42
    def quantize(self, tokenizer=None, quant_config={}, n_samples=128, seqlen=512,
43
                       auto_scale=True, mse_range=True, run_search=True, run_quant=True,
Casper Hansen's avatar
Casper Hansen committed
44
                       calib_data="pileval"):
45
        self.quant_config = quant_config
46
47
48
49
50
        quant_config["version"] = "GEMM" if 'version' not in quant_config.keys() else quant_config["version"]
        if quant_config["version"] == "GEMM":
            logging.warning('Deprecated model weight format. Re-quantize '
                            'your weights again with version="GEMV" for a speedup. '
                            'In the next AutoAWQ version, GEMM will be deprecated.')
51

Casper Hansen's avatar
Casper Hansen committed
52
        if run_search:
53
            self.search_result = self._awq_search(tokenizer, quant_config, n_samples=n_samples, seqlen=seqlen,
Casper Hansen's avatar
Casper Hansen committed
54
55
56
                       auto_scale=auto_scale, mse_range=mse_range, calib_data=calib_data)
        
        if run_quant:
57
            self._awq_quant()
Casper Hansen's avatar
Casper Hansen committed
58
            self.is_quantized = True
Casper Hansen's avatar
Casper Hansen committed
59
    
qwopqwop200's avatar
qwopqwop200 committed
60
    @staticmethod
61
    def fuse_layers(model, quant_config):
qwopqwop200's avatar
qwopqwop200 committed
62
63
        pass
        
64
65
    def _awq_quant(self):
        assert self.quant_config["zero_point"], "We only support zero_point quantization now."
66
        layers = self.get_model_layers(self.model)
Casper's avatar
Casper committed
67

Casper Hansen's avatar
Casper Hansen committed
68
69
70
71
        # Run AWQ quantization
        for i in tqdm(range(len(layers)), desc="AWQ Quantization"):
            layer = layers[i]
            named_linears = get_named_linears(layer)
72
            self._scale_activations(self, layer)
Casper Hansen's avatar
Casper Hansen committed
73
74

            for name, module in named_linears.items():
Casper Hansen's avatar
Casper Hansen committed
75
                module.cuda()
76
77
78
79
80
81
82

                module.weight.data, scales, zeros = pseudo_quantize_tensor(
                    module.weight.data, 
                    get_scale_zp=True, 
                    **self.quant_config
                )

Casper Hansen's avatar
Casper Hansen committed
83
84
                scales = scales.t().contiguous()
                zeros = zeros.t().contiguous()
85

86
87
88
89
90
91
92
93
94
95
96
                if self.quant_config["version"] == 'GEMM':
                    q_linear_module = WQLinear_GEMM
                elif self.quant_config["version"] == 'GEMV':
                    q_linear_module = WQLinear_GEMV
                
                q_linear = q_linear_module.from_linear(
                    module,
                    self.quant_config['w_bit'],
                    self.quant_config['q_group_size'],
                    False,
                    scales,
97
98
99
                    zeros
                )

Casper Hansen's avatar
Casper Hansen committed
100
101
102
103
104
                module.cpu()
                q_linear.to(next(layer.parameters()).device)
                set_op_by_name(layer, name, q_linear)
                torch.cuda.empty_cache()
                gc.collect()
Casper Hansen's avatar
Casper Hansen committed
105
106
107
108
            
            torch.cuda.empty_cache()
            gc.collect()
    
109
    def _awq_search(self, tokenizer, quant_config, n_samples=128, seqlen=512,
Casper Hansen's avatar
Casper Hansen committed
110
                       auto_scale=True, mse_range=True, calib_data="pileval"):
111
        layers = self.get_model_layers(self.model)
Casper Hansen's avatar
Casper Hansen committed
112
113
114
115
116
117
118
119
120

        samples = get_calib_dataset(
            data=calib_data, tokenizer=tokenizer, n_samples=n_samples, block_size=seqlen)
        samples = torch.cat(samples, dim=0)

        inps = []
        layer_kwargs = {}

        layers[0] = layers[0].cuda()
121
        self.move_embed(self.model, "cuda")
Casper Hansen's avatar
Casper Hansen committed
122
123
124
125
126
127
128
129
130
        
        # 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

Casper's avatar
Casper committed
131
132
            def forward(self, hijacked_inputs, **kwargs):
                inps.append(hijacked_inputs)
Casper Hansen's avatar
Casper Hansen committed
133
134
135
136
137
138
                layer_kwargs.update(kwargs)
                raise ValueError  # early exit to break later inference

        # patch layer 0 to catch input and kwargs
        layers[0] = Catcher(layers[0])
        try:
139
            self.model(samples.to(next(self.model.parameters()).device))
Casper Hansen's avatar
Casper Hansen committed
140
141
142
143
144
145
146
        except ValueError:  # work with early exit
            pass
        del samples
        layers[0] = layers[0].module  # restore
        inps = inps[0]

        layers[0] = layers[0].cpu()
147
        self.move_embed(self.model, "cpu")
Casper Hansen's avatar
Casper Hansen committed
148
149
150
151
152
153
154
155
        
        gc.collect()
        torch.cuda.empty_cache()
        awq_results = {
            "scale": [],
            "clip": [],
        }

Casper Hansen's avatar
Casper Hansen committed
156
        # Run AWQ search layer by layer
157
        for i in tqdm(range(len(layers)), desc="AWQ Search"):
Casper Hansen's avatar
Casper Hansen committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
            layer = layers[i]
            layer = layer.cuda()
            named_linears = get_named_linears(layer)

            # 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)))
            inps = inps.to(next(layer.parameters()).device)  # in case multi-gpu
            # get output as next layer's input
            inps = layer(inps, **layer_kwargs)[0]
            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()}

            # Clear GPU memory
            torch.cuda.empty_cache()

            if auto_scale:  # if it applies, we should also modify the input_feat with scales
                scales_list = auto_scale_block(
                    self,
188
189
190
                    layer,
                    layer_kwargs,
                    quant_config=quant_config,
Casper Hansen's avatar
Casper Hansen committed
191
192
                    input_feat=input_feat,
                )
193

Casper Hansen's avatar
Casper Hansen committed
194
                apply_scale(layers[i], scales_list, input_feat_dict=input_feat)
195

Casper Hansen's avatar
Casper Hansen committed
196
                # append prefix to make names global
197
                awq_results["scale"] += append_str_prefix(scales_list, get_op_name(self.model, layer) + ".")
Casper Hansen's avatar
Casper Hansen committed
198
199
200
201
202

            # Clear GPU memory
            torch.cuda.empty_cache()
            
            if mse_range:
203
204
205
206
207
208
                clip_list = auto_clip_block(
                    layer,
                    quant_config=quant_config,
                    input_feat=input_feat
                )

Casper Hansen's avatar
Casper Hansen committed
209
210
                apply_clip(layer, clip_list)
                # append prefix to make names global
211
                awq_results["clip"] += append_str_prefix(clip_list, get_op_name(self.model, layer) + ".")
Casper Hansen's avatar
Casper Hansen committed
212
213
214
215
216
217

            layer = layer.cpu()
            # Haotian: check activation replacement
            del input_feat
            gc.collect()
            torch.cuda.empty_cache()
Casper Hansen's avatar
Casper Hansen committed
218
        
Casper Hansen's avatar
Casper Hansen committed
219
        return awq_results
Casper's avatar
Casper committed
220

221
    def save_quantized(self, save_dir):
222
223
224
225
226
227
228
229
230
231
232
233
234
235
        def _save_files(save_dir, model_name, model):
            class EmptyModule(nn.Module):
                def __init__(self): super(EmptyModule, self).__init__()
                def forward(self, x): return x

            # Save model fiels without search results
            self.model.save_pretrained(save_dir, state_dict=EmptyModule().state_dict())

            # Remove empty module
            os.remove(f'{save_dir}/pytorch_model.bin')

            # Save search results
            torch.save(model, f'{save_dir}/{model_name}')

236
237
238
239
            # Save config
            with open(f'{save_dir}/quant_config.json', 'w+') as file:
                file.write(json.dumps(self.quant_config, indent=4))

240
241
242
        save_dir = save_dir[:-1] if save_dir[-1] == '/' else save_dir

        # Save model
Casper Hansen's avatar
Casper Hansen committed
243
        if self.search_result is None or self.is_quantized:
Casper Hansen's avatar
Casper Hansen committed
244
            model_name = f'awq_model_w{self.quant_config["w_bit"]}_g{self.quant_config["q_group_size"]}.pt'
245
            _save_files(save_dir, model_name, self.model.state_dict())
246
247
        else:
            model_name = 'awq_model_search_result.pt'
248
249
            _save_files(save_dir, model_name, self.search_result)
        
250
251
    @classmethod
    def from_pretrained(self, model_path, model_type, torch_dtype: torch.dtype = torch.float16, 
Casper Hansen's avatar
Casper Hansen committed
252
                        trust_remote_code=True, safetensors=False):
253
254
255
        return self.from_quantized(
            model_path, 
            model_type, 
256
            model_filename='', 
257
            max_new_tokens=None,
258
259
260
            device='balanced', 
            torch_dtype=torch_dtype, 
            trust_remote_code=trust_remote_code, 
Casper Hansen's avatar
Casper Hansen committed
261
            safetensors=safetensors,
262
263
            is_quantized=False
        )
Casper's avatar
Casper committed
264

265
    @classmethod
266
    def from_quantized(self, model_path, model_type, model_filename, max_new_tokens=None,
267
                       device='balanced', torch_dtype=torch.float16, trust_remote_code=True, 
Casper Hansen's avatar
Casper Hansen committed
268
                       safetensors=False, is_quantized=True, fuse_layers=False, version='GEMM'):
269
        # [STEP 1] Download model if path is not a directory
270
        if not os.path.isdir(model_path):
271
272
273
274
275
276
277
            ignore_patterns = ["*msgpack*", "*h5*"]
            if safetensors:
                ignore_patterns.extend(["*.pt", "*.bin"])
            else:
                ignore_patterns.append("*safetensors*")

            model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns)
278
279
280
        
        # TODO: Better naming, model_filename becomes a directory
        model_filename = model_path + f'/{model_filename}'
281

282
        # [STEP 2] Load config and set sequence length
283
        # TODO: Create BaseAWQConfig class
284
285
286
287
        quant_config_path = f'{model_path}/quant_config.json'
        if os.path.exists(quant_config_path):
            with open(quant_config_path, 'r') as file:
                quant_config = json.loads(file.read())
288
289
290
            
            if "version" not in quant_config.keys():
                quant_config["version"] = version
291
292
        else:
            # Default config that works for most models
293
            quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version": version}
294
        
295
296
297
298
299
300
301
302
303
        # Load model config and set max generation length
        if max_new_tokens is None and hasattr(self, 'max_new_tokens_key'):
            config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
            config.max_new_tokens = getattr(config, self.max_new_tokens_key)
        else:
            max_new_tokens = 2048 if max_new_tokens is None else max_new_tokens
            config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
            config.max_new_tokens = max_new_tokens
        
304
        # [STEP 3] Load model
305
        with init_empty_weights():
306
307
            model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code)
        
308
        # Only need to replace layers if a model is AWQ quantized
309
310
        if is_quantized:
            # Prepare WQLinear layers, replace nn.Linear
311
            self._load_quantized_modules(self, model, quant_config, quant_config["version"])
312
313
        
        model.tie_weights()
314

315
316
317
318
319
320
        device_map = infer_auto_device_map(
            model,
            no_split_module_classes=[self.layer_type], 
            dtype=torch_dtype
        )

321
        # Load model weights
322
        if is_quantized:
323
324
325
326
327
328
            model = load_checkpoint_and_dispatch(
                model, 
                model_filename, 
                device_map=device_map, 
                no_split_module_classes=[self.layer_type]
            )
329

330
            if fuse_layers:
331
                self.fuse_layers(model, quant_config)
332

333
334
        else:
            # If not quantized, must load with AutoModelForCausalLM
335
336
337
338
            del model
            
            # Load model weights
            model = AutoModelForCausalLM.from_pretrained(
339
340
341
342
343
344
345
                model_filename, 
                device_map=device_map, 
                trust_remote_code=trust_remote_code, 
                offload_folder="offload", 
                offload_state_dict=True, 
                torch_dtype=torch_dtype, 
                use_safetensors=safetensors
346
347
            )
            model.eval()
348

349
        return self(model, model_type, is_quantized=is_quantized, quant_config=quant_config)
Casper's avatar
Casper committed
350

Casper Hansen's avatar
Casper Hansen committed
351
    def _load_quantized_modules(self, model, quant_config, version):
352
        # Real quantization of weights
353
        assert quant_config["zero_point"], "We only support zero_point quantization now."
Casper Hansen's avatar
Casper Hansen committed
354
355
356
357
358

        if version == 'GEMM':
            logging.warning('Deprecated model weight format. Re-quantize '
                            'your weights again with version="GEMV" for a speedup. '
                            'In the next AutoAWQ version, GEMM will be deprecated.')
359
360
        
        # Get blocks of model
361
        layers = self.get_model_layers(model)
362

363
364
        for i in tqdm(range(len(layers)), desc="Replacing layers..."):
            layer = layers[i]
365
366

            # Get every linear layer in a block
367
            named_linears = get_named_linears(layer)
368
369

            # Replace activation functions
370
            self._scale_activations(self, layer)
371

372
            # Replace nn.Linear with WQLinear
373
            for name, module in named_linears.items():
Casper Hansen's avatar
Casper Hansen committed
374
375
376
377
378
379
                if version == 'GEMM':
                    q_linear_module = WQLinear_GEMM
                elif version == 'GEMV':
                    q_linear_module = WQLinear_GEMV
                
                q_linear = q_linear_module.from_linear(
380
381
382
                    module,
                    quant_config['w_bit'],
                    quant_config['q_group_size'],
Casper Hansen's avatar
Casper Hansen committed
383
384
                    True
                )
385
386
387
388
389
390
                q_linear.to(next(layer.parameters()).device)
                set_op_by_name(layer, name, q_linear)
            
            torch.cuda.empty_cache()
            gc.collect()
    
391
    @staticmethod
392
    def _scale_activations(self, layer):
393
        scale_dict = self.get_act_for_scaling(layer)
394

395
396
397
        if scale_dict['is_scalable']:
            if not isinstance(scale_dict['scale_layer'], ScaledActivation):
                param = next(layer.parameters())
398

399
400
                # get activation scale
                scale_like = torch.ones(scale_dict['scale_shape'], dtype=param.dtype, device=param.device)
401

402
403
                # scale activation
                scaled_act = ScaledActivation(scale_dict['scale_layer'], scale_like)
Casper's avatar
Casper committed
404
                set_op_by_name(layer, scale_dict['scale_name'], scaled_act)