| import argparse |
| import copy |
| import hashlib |
| import importlib.metadata |
| import json |
| import random |
| import shutil |
| import time |
| from pathlib import Path |
|
|
| import modelopt.torch.quantization as mtq |
| import pyarrow.parquet as parquet |
| import torch |
| from huggingface_hub import HfApi, hf_hub_download, save_torch_state_dict |
| from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer |
| from packed_int4 import PackedLinear, pack_weight |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--method', choices=['awq', 'gptq'], required=True) |
| parser.add_argument('--source', required=True) |
| parser.add_argument('--output', required=True) |
| parser.add_argument('--samples', type=int, default=128) |
| args = parser.parse_args() |
| source, output = Path(args.source), Path(args.output) |
| output.mkdir(parents=True, exist_ok=False) |
| torch.set_num_threads(4) |
| torch.manual_seed(20261020) |
| tokenizer = AutoTokenizer.from_pretrained(source, trust_remote_code=True, local_files_only=True) |
| tokenizer.chat_template = (Path(__file__).parent / 'chat_template.jinja').read_text() |
| dataset = 'Salesforce/wikitext' |
| revision = HfApi(token=False).dataset_info(dataset).sha |
| data_path = hf_hub_download(dataset, 'wikitext-2-raw-v1/train-00000-of-00001.parquet', |
| repo_type='dataset', revision=revision) |
| text = '\n\n'.join(parquet.read_table(data_path, columns=['text'])['text'].to_pylist()) |
| ids = tokenizer(text, return_tensors='pt', add_special_tokens=False).input_ids |
| rng = random.Random(20261020) |
| starts = [rng.randrange(ids.shape[1] - 512) for _ in range(args.samples)] |
| batches = [ids[:, start:start + 512].clone() for start in starts] |
| calibration = {'dataset': dataset, 'revision': revision, 'split': 'train', 'samples': args.samples, |
| 'sequence_length': 512, 'seed': 20261020, 'token_offsets': starts, |
| 'token_ids_sha256': hashlib.sha256(torch.cat(batches).numpy().tobytes()).hexdigest()} |
| (output / 'calibration.json').write_text(json.dumps(calibration, indent=2) + '\n') |
| config = AutoConfig.from_pretrained(source, trust_remote_code=True, local_files_only=True) |
| config.backend = 'torch' |
| config.prefill_chunk_size = 512 |
| model = AutoModelForCausalLM.from_pretrained( |
| source, config=config, trust_remote_code=True, local_files_only=True, |
| dtype=torch.bfloat16, device_map={'': 0}).eval().requires_grad_(False) |
| model.generation_config.eos_token_id = [248046, 248044] |
| model.generation_config.pad_token_id = 248044 |
| modules = {name: module for name, module in model.named_modules() |
| if isinstance(module, torch.nn.Linear) and name != 'lm_head' |
| and min(module.weight.shape) >= 128 and '.admission_' not in name} |
| assert len(modules) == 256 |
| recipe = copy.deepcopy(mtq.INT4_AWQ_CFG) |
| for name, module in model.named_modules(): |
| if isinstance(module, torch.nn.Linear) and name not in modules: |
| recipe['quant_cfg'].append({'quantizer_name': name + '.*', 'enable': False}) |
| if args.method == 'gptq': |
| recipe['algorithm'] = {'method': 'gptq', 'block_size': 128, 'perc_damp': 0.01} |
|
|
| def forward_loop(current): |
| for index, batch in enumerate(batches): |
| current(input_ids=batch.to('cuda'), use_cache=False, logits_to_keep=1) |
| if index % 32 == 0: |
| print(json.dumps({'stage': 'calibration', 'batch': index, 'samples': args.samples}), flush=True) |
|
|
| started = time.perf_counter() |
| print(json.dumps({'stage': 'quantizing', 'method': args.method, 'modules': len(modules)}), flush=True) |
| with torch.inference_mode(): |
| mtq.quantize(model, recipe, forward_loop=forward_loop) |
| manifest = {'method': args.method, 'bits': 4, 'group_size': 128, |
| 'activation_dtype': 'bfloat16', 'scale_dtype': 'float32', |
| 'zero_point': 8, 'awq_input_scaling': args.method == 'awq', 'modules': {}} |
| state = {} |
| maximum_weight_error = 0.0 |
| maximum_kernel_relative_error = 0.0 |
| for index, (name, module) in enumerate(modules.items()): |
| weight = module.weight_quantizer(module.weight).to(torch.bfloat16) |
| rows, columns = weight.shape |
| scale = (module.weight_quantizer.amax.float().reshape(rows, columns // 128) / 7).clamp_min(1e-12) |
| codes = (weight.float().reshape(rows, columns // 128, 128) / scale[..., None]).round().clamp(-8, 7).reshape(rows, columns).to(torch.int32) + 8 |
| input_scale = getattr(module.input_quantizer, '_pre_quant_scale', None) |
| layer = PackedLinear(columns, rows, args.method == 'awq', input_scale is not None, module.bias is not None).to('cuda') |
| layer.qweight = pack_weight(codes, args.method == 'awq') |
| layer.scales = scale.t().contiguous() |
| layer.qzeros.fill_(-2004318072 if args.method == 'awq' else 2004318071) |
| if layer.g_idx is not None: |
| layer.g_idx = torch.arange(columns, device='cuda', dtype=torch.int32) // 128 |
| if input_scale is not None: |
| layer.input_scale = input_scale.to(torch.bfloat16).reshape(-1).contiguous() |
| if module.bias is not None: |
| layer.bias = module.bias.detach() |
| dense = layer.dequantize() |
| error = (dense.float() - weight.float()).abs().max().item() |
| maximum_weight_error = max(maximum_weight_error, error) |
| inputs = torch.randn((3, columns), device='cuda', dtype=torch.bfloat16) |
| expected_inputs = inputs if layer.input_scale is None else inputs * layer.input_scale |
| expected = torch.nn.functional.linear(expected_inputs, dense, layer.bias) |
| actual = layer(inputs) |
| relative = (expected.float() - actual.float()).abs().max().item() / max(expected.float().abs().max().item(), 1e-12) |
| assert relative < 0.025, (name, relative) |
| maximum_kernel_relative_error = max(maximum_kernel_relative_error, relative) |
| manifest['modules'][name] = {'shape': [rows, columns], 'input_scale': input_scale is not None, 'bias': module.bias is not None} |
| state.update({name + '.' + field: value.detach().cpu().contiguous() for field, value in layer.state_dict().items()}) |
| if index % 32 == 0: |
| print(json.dumps({'stage': 'packing', 'module': index, 'weight_error': error, 'kernel_relative_error': relative}), flush=True) |
| for name, tensor in model.state_dict().items(): |
| if name == 'lm_head.weight' or 'quantizer' in name or name.removesuffix('.weight') in modules: |
| continue |
| state[name] = tensor.detach().cpu().contiguous() |
| save_torch_state_dict(state, output, max_shard_size='3GB', safe_serialization=True) |
| model.config.quantization_config = { |
| 'quant_method': args.method, 'bits': 4, 'group_size': 128, |
| **({'version': 'gemm', 'zero_point': True} if args.method == 'awq' |
| else {'checkpoint_format': 'gptq', 'desc_act': False, 'sym': True, 'damp_percent': 0.01}), |
| 'edgerock_runtime': 'packed_int4.load_int4', |
| } |
| model.config.save_pretrained(output) |
| model.generation_config.save_pretrained(output) |
| tokenizer.save_pretrained(output) |
| for name in ('configuration_edgerock.py', 'modeling_edgerock.py', 'LICENSE', 'LICENSE-CODE', 'NOTICE'): |
| shutil.copy2(source / name, output / name) |
| for name in ('packed_int4.py', 'quantize.py', 'validate.py', 'chat_template.jinja'): |
| shutil.copy2(Path(__file__).parent / name, output / name) |
| (output / 'edgerock_quantization.json').write_text(json.dumps(manifest, indent=2) + '\n') |
| versions = {name: importlib.metadata.version(name) for name in |
| ('torch', 'transformers', 'nvidia-modelopt', 'triton', 'accelerate', 'safetensors', 'huggingface-hub')} |
| conversion = {'method': args.method, 'recipe': recipe, 'quantized_modules': len(modules), |
| 'maximum_weight_repacking_error': maximum_weight_error, |
| 'maximum_kernel_relative_error': maximum_kernel_relative_error, |
| 'weight_bytes': sum(path.stat().st_size for path in output.glob('model*.safetensors')), |
| 'seconds': time.perf_counter() - started, 'versions': versions, |
| 'hardware': torch.cuda.get_device_name(0)} |
| (output / 'conversion.json').write_text(json.dumps(conversion, indent=2) + '\n') |
| print(json.dumps({'stage': 'conversion_complete', 'weight_bytes': conversion['weight_bytes'], |
| 'seconds': conversion['seconds']}), flush=True) |
|
|
