| 4 | |
| 5 | |
| 6 | def quantization(config, |
| 7 | quantization_model_name, |
| 8 | origin_model_name, |
| 9 | quantization_type: str = 'awq', |
| 10 | cuda_prefix: str = 'CUDA_VISIBLE_DEVICES=0'): |
| 11 | model_path = config.get('model_path') |
| 12 | log_path = config.get('log_path') |
| 13 | origin_model_path = os.path.join(config.get('model_path'), origin_model_name) |
| 14 | quantization_model_path = os.path.join(model_path, quantization_model_name) |
| 15 | quantization_log = os.path.join( |
| 16 | log_path, '_'.join(['quantization', quantization_type, |
| 17 | quantization_model_name.split('/')[1]]) + '.log') |
| 18 | |
| 19 | if quantization_type == 'awq': |
| 20 | quantization_cmd = ' '.join( |
| 21 | ['lmdeploy lite auto_awq', origin_model_path, '--work-dir', quantization_model_path]) |
| 22 | elif quantization_type == 'gptq': |
| 23 | quantization_cmd = ' '.join( |
| 24 | ['lmdeploy lite auto_gptq', origin_model_path, '--work-dir', quantization_model_path]) |
| 25 | elif quantization_type == 'w8a8': |
| 26 | quantization_cmd = ' '.join( |
| 27 | ['lmdeploy lite smooth_quant', origin_model_path, '--work-dir', quantization_model_path]) |
| 28 | else: |
| 29 | return False, 'quantization type should in [awq, gptq, w8a8], \ |
| 30 | now the type is ' + quantization_type |
| 31 | |
| 32 | # Add device option if specified in environment |
| 33 | device = os.environ.get('DEVICE', '') |
| 34 | if device == 'ascend': |
| 35 | quantization_cmd += ' --device npu ' |
| 36 | |
| 37 | if cuda_prefix is not None: |
| 38 | quantization_cmd = ' '.join([cuda_prefix, quantization_cmd]) |
| 39 | |
| 40 | if 'llama-3' in origin_model_name.lower(): |
| 41 | quantization_cmd += ' --search-scale' |
| 42 | |
| 43 | if quantization_type == 'gptq' or str(config.get('env_tag')) == '3090' or str(config.get('env_tag')) == '5080': |
| 44 | quantization_cmd += ' --batch-size 8' |
| 45 | else: |
| 46 | quantization_cmd += ' --batch-size 32' |
| 47 | |
| 48 | quantization_cmd += ' --trust-remote-code' |
| 49 | |
| 50 | with open(quantization_log, 'w') as f: |
| 51 | # remove existing folder |
| 52 | subprocess.run([' '.join(['rm -rf', quantization_model_path])], |
| 53 | stdout=f, |
| 54 | stderr=f, |
| 55 | shell=True, |
| 56 | text=True, |
| 57 | encoding='utf-8') |
| 58 | |
| 59 | f.writelines('reproduce command quantization_cmd: ' + quantization_cmd + '\n') |
| 60 | print('reproduce command quantization_cmd: ' + quantization_cmd) |
| 61 | # quantization |
| 62 | quantizationRes = subprocess.run([quantization_cmd], |
| 63 | stdout=f, |