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

12
from awq.modules.act import ScaledActivation
13
from huggingface_hub import snapshot_download
Casper Hansen's avatar
Casper Hansen committed
14
from awq.utils.calib_data import get_calib_dataset
15
from awq.quantize.quantizer import pseudo_quantize_tensor
16
from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
Casper Hansen's avatar
Casper Hansen committed
17
18
from awq.quantize.auto_clip import auto_clip_block, apply_clip
from awq.quantize.auto_scale import auto_scale_block, apply_scale
19
20
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
21
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
22

23
class BaseAWQForCausalLM(nn.Module):
24
    def __init__(self, model, model_type, is_quantized, quant_config):
25
        super().__init__()
26
27
28
29
        self.model:PreTrainedModel = model
        self.model_type:str = model_type
        self.is_quantized:bool = is_quantized
        self.search_result = None
30
        self.quant_config:dict = quant_config
31
32
33
34
35
36
    
    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
37
38
39
40
    
    def generate(self, *args, **kwargs):
        with torch.inference_mode():
            return self.model.generate(*args, **kwargs)
41

Casper Hansen's avatar
Casper Hansen committed
42
    @torch.no_grad()
43
    def quantize(self, tokenizer=None, quant_config={}, n_samples=128, seqlen=512,
44
                       auto_scale=True, mse_range=True, run_search=True, run_quant=True,
45
                       calib_data: Union[str, List[str]]="pileval"):
46
        self.quant_config = quant_config
47
        quant_config["version"] = "GEMM" if 'version' not in quant_config.keys() else quant_config["version"]
48

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

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

            for name, module in named_linears.items():
Casper Hansen's avatar
Casper Hansen committed
72
                module.cuda()
73
74
75
76

                module.weight.data, scales, zeros = pseudo_quantize_tensor(
                    module.weight.data, 
                    get_scale_zp=True, 
77
78
                    w_bit=self.quant_config["w_bit"], 
                    q_group_size=self.quant_config["q_group_size"]
79
80
                )

81
                if self.quant_config["version"] == 'GEMM':
82
83
                    scales = scales.t().contiguous()
                    zeros = zeros.t().contiguous()
84
85
86
87
88
89
90
91
92
93
                    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,
94
95
96
                    zeros
                )

Casper Hansen's avatar
Casper Hansen committed
97
98
99
100
101
                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
102
103
104
105
            
            torch.cuda.empty_cache()
            gc.collect()
    
106
    def _awq_search(self, tokenizer, quant_config, n_samples=128, seqlen=512,
107
                       auto_scale=True, mse_range=True, calib_data:Union[str, List[str]]="pileval"):
108
        layers = self.get_model_layers(self.model)
Casper Hansen's avatar
Casper Hansen committed
109
110
111
112
113
114
115
116
117

        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()
118
        self.move_embed(self.model, "cuda")
Casper Hansen's avatar
Casper Hansen committed
119
120
121
122
123
124
125
126
127
        
        # 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
128
129
            def forward(self, hijacked_inputs, **kwargs):
                inps.append(hijacked_inputs)
Casper Hansen's avatar
Casper Hansen committed
130
131
132
133
134
135
                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:
136
            self.model(samples.to(next(self.model.parameters()).device))
Casper Hansen's avatar
Casper Hansen committed
137
138
139
140
141
142
143
        except ValueError:  # work with early exit
            pass
        del samples
        layers[0] = layers[0].module  # restore
        inps = inps[0]

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

Casper Hansen's avatar
Casper Hansen committed
153
        # Run AWQ search layer by layer
154
        for i in tqdm(range(len(layers)), desc="AWQ Search"):
Casper Hansen's avatar
Casper Hansen committed
155
156
157
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
            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,
185
186
187
                    layer,
                    layer_kwargs,
                    quant_config=quant_config,
Casper Hansen's avatar
Casper Hansen committed
188
189
                    input_feat=input_feat,
                )
190

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

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

            # Clear GPU memory
            torch.cuda.empty_cache()
            
            if mse_range:
200
201
202
203
204
205
                clip_list = auto_clip_block(
                    layer,
                    quant_config=quant_config,
                    input_feat=input_feat
                )

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

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

218
    def save_quantized(self, save_dir):
219
220
221
222
223
224
225
226
227
228
229
230
231
232
        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}')

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

237
238
239
        save_dir = save_dir[:-1] if save_dir[-1] == '/' else save_dir

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

262
    @classmethod
263
    def from_quantized(self, model_path, model_type, model_filename, max_new_tokens=None,
264
                       device='balanced', torch_dtype=torch.float16, trust_remote_code=True, 
Casper Hansen's avatar
Casper Hansen committed
265
                       safetensors=False, is_quantized=True, fuse_layers=False, version='GEMM'):
266
        # [STEP 1] Download model if path is not a directory
267
        if not os.path.isdir(model_path):
268
269
270
271
272
273
274
            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)
275
276
277
        
        # TODO: Better naming, model_filename becomes a directory
        model_filename = model_path + f'/{model_filename}'
278

279
        # [STEP 2] Load config and set sequence length
280
        # TODO: Create BaseAWQConfig class
281
282
283
284
        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())
285
286
287
            
            if "version" not in quant_config.keys():
                quant_config["version"] = version
288
289
        else:
            # Default config that works for most models
290
            quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version": version}
291
        
292
293
294
295
296
297
298
299
300
        # 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
        
301
        # [STEP 3] Load model
302
        with init_empty_weights():
303
304
            model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code)
        
305
        # Only need to replace layers if a model is AWQ quantized
306
307
        if is_quantized:
            # Prepare WQLinear layers, replace nn.Linear
308
            self._load_quantized_modules(self, model, quant_config, quant_config["version"])
309
310
        
        model.tie_weights()
311

312
313
314
315
316
317
        device_map = infer_auto_device_map(
            model,
            no_split_module_classes=[self.layer_type], 
            dtype=torch_dtype
        )

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

327
            if fuse_layers:
328
                self.fuse_layers(model, quant_config)
329

330
331
        else:
            # If not quantized, must load with AutoModelForCausalLM
332
333
334
335
            del model
            
            # Load model weights
            model = AutoModelForCausalLM.from_pretrained(
336
337
338
339
340
341
342
                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
343
344
            )
            model.eval()
345

346
        return self(model, model_type, is_quantized=is_quantized, quant_config=quant_config)
Casper's avatar
Casper committed
347

Casper Hansen's avatar
Casper Hansen committed
348
    def _load_quantized_modules(self, model, quant_config, version):
349
        # Real quantization of weights
350
        assert quant_config["zero_point"], "We only support zero_point quantization now."
351
352
        
        # Get blocks of model
353
        layers = self.get_model_layers(model)
354

355
356
        for i in tqdm(range(len(layers)), desc="Replacing layers..."):
            layer = layers[i]
357
358

            # Get every linear layer in a block
359
            named_linears = get_named_linears(layer)
360
361

            # Replace activation functions
362
            self._scale_activations(self, layer)
363

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

387
388
389
        if scale_dict['is_scalable']:
            if not isinstance(scale_dict['scale_layer'], ScaledActivation):
                param = next(layer.parameters())
390

391
392
                # get activation scale
                scale_like = torch.ones(scale_dict['scale_shape'], dtype=param.dtype, device=param.device)
393

394
395
                # scale activation
                scaled_act = ScaledActivation(scale_dict['scale_layer'], scale_like)
Casper's avatar
Casper committed
396
                set_op_by_name(layer, scale_dict['scale_name'], scaled_act)