| import torch | |
| import json | |
| import os | |
| import struct | |
| import numpy as np | |
| from transformers import AutoModelForCausalLM, AutoConfig | |
| from collections import OrderedDict | |
| from tqdm import tqdm | |
| def quantize_tensor_symmetric(tensor, bits_per_param): | |
| """ | |
| Symmetric linear quantization with zero-point preservation. | |
| Args: | |
| tensor: Input float tensor | |
| bits_per_param: Bit width for quantization (2,3,4,5,6,8,12,16,32) | |
| Returns: | |
| quantized_signed: Signed integer tensor | |
| scale: Quantization scale factor | |
| level: Maximum absolute level L = 2^(bits-1)-1 | |
| dtype_name: String descriptor | |
| """ | |
| if bits_per_param == 32: | |
| return tensor.to(torch.float32), None, None, "float32" | |
| elif bits_per_param == 16: | |
| return tensor.to(torch.float16), None, None, "float16" | |
| max_abs = tensor.abs().max().item() | |
| if max_abs == 0: | |
| max_abs = 1e-8 | |
| L = 2 ** (bits_per_param - 1) - 1 | |
| scale = max_abs / L | |
| quantized = (tensor / scale).round().clamp(-L, L) | |
| dtype_name = f"int{bits_per_param}_sym" | |
| storage_dtype = torch.int8 if bits_per_param <= 8 else torch.int16 | |
| return quantized.to(storage_dtype), scale, L, dtype_name | |
| def pack_2bit(values): | |
| """Pack 4 unsigned 2-bit values into 1 byte.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint8) | |
| pad = (4 - len(flat) % 4) % 4 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)]) | |
| flat = flat.view(-1, 4) | |
| packed = ((flat[:, 0] & 0x03) | | |
| ((flat[:, 1] & 0x03) << 2) | | |
| ((flat[:, 2] & 0x03) << 4) | | |
| ((flat[:, 3] & 0x03) << 6)) | |
| return packed, shape, pad | |
| def pack_3bit(values): | |
| """Pack 8 unsigned 3-bit values into 3 bytes.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint8) | |
| pad = (8 - len(flat) % 8) % 8 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)]) | |
| flat = flat.view(-1, 8) | |
| packed = torch.zeros(len(flat) * 3, dtype=torch.uint8) | |
| packed[0::3] = ((flat[:, 0] & 0x07) | | |
| ((flat[:, 1] & 0x07) << 3) | | |
| ((flat[:, 2] & 0x03) << 6)) | |
| packed[1::3] = (((flat[:, 2] >> 2) & 0x01) | | |
| ((flat[:, 3] & 0x07) << 1) | | |
| ((flat[:, 4] & 0x07) << 4) | | |
| ((flat[:, 5] & 0x01) << 7)) | |
| packed[2::3] = (((flat[:, 5] >> 1) & 0x03) | | |
| ((flat[:, 6] & 0x07) << 2) | | |
| ((flat[:, 7] & 0x07) << 5)) | |
| return packed, shape, pad | |
| def pack_4bit(values): | |
| """Pack 2 unsigned 4-bit values into 1 byte.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint8) | |
| pad = (2 - len(flat) % 2) % 2 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)]) | |
| flat = flat.view(-1, 2) | |
| packed = ((flat[:, 0] & 0x0F) | | |
| ((flat[:, 1] & 0x0F) << 4)) | |
| return packed, shape, pad | |
| def pack_5bit(values): | |
| """Pack 8 unsigned 5-bit values into 5 bytes.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint8) | |
| pad = (8 - len(flat) % 8) % 8 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)]) | |
| flat = flat.view(-1, 8) | |
| packed = torch.zeros(len(flat) * 5, dtype=torch.uint8) | |
| packed[0::5] = ((flat[:, 0] & 0x1F) | | |
| ((flat[:, 1] & 0x07) << 5)) | |
| packed[1::5] = (((flat[:, 1] >> 3) & 0x03) | | |
| ((flat[:, 2] & 0x1F) << 2) | | |
| ((flat[:, 3] & 0x01) << 7)) | |
| packed[2::5] = (((flat[:, 3] >> 1) & 0x0F) | | |
| ((flat[:, 4] & 0x0F) << 4)) | |
| packed[3::5] = (((flat[:, 4] >> 4) & 0x01) | | |
| ((flat[:, 5] & 0x1F) << 1) | | |
| ((flat[:, 6] & 0x03) << 6)) | |
| packed[4::5] = (((flat[:, 6] >> 2) & 0x07) | | |
| ((flat[:, 7] & 0x1F) << 3)) | |
| return packed, shape, pad | |
| def pack_6bit(values): | |
| """Pack 4 unsigned 6-bit values into 3 bytes.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint8) | |
| pad = (4 - len(flat) % 4) % 4 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)]) | |
| flat = flat.view(-1, 4) | |
| packed = torch.zeros(len(flat) * 3, dtype=torch.uint8) | |
| packed[0::3] = ((flat[:, 0] & 0x3F) | | |
| ((flat[:, 1] & 0x03) << 6)) | |
| packed[1::3] = (((flat[:, 1] >> 2) & 0x0F) | | |
| ((flat[:, 2] & 0x0F) << 4)) | |
| packed[2::3] = (((flat[:, 2] >> 4) & 0x03) | | |
| ((flat[:, 3] & 0x3F) << 2)) | |
| return packed, shape, pad | |
| def pack_12bit(values): | |
| """Pack 2 unsigned 12-bit values into 3 bytes.""" | |
| shape = values.shape | |
| flat = values.flatten().to(torch.uint16) | |
| pad = (2 - len(flat) % 2) % 2 | |
| if pad: | |
| flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint16)]) | |
| flat = flat.view(-1, 2) | |
| a = flat[:, 0].numpy().astype(np.uint16) | |
| b = flat[:, 1].numpy().astype(np.uint16) | |
| byte0 = a & 0xFF | |
| byte1 = ((a >> 8) & 0x0F) | ((b & 0x0F) << 4) | |
| byte2 = (b >> 4) & 0xFF | |
| packed = np.stack([byte0, byte1, byte2], axis=1).flatten().astype(np.uint8) | |
| return torch.from_numpy(packed), shape, pad | |
| def convert_to_dq(model_path, output_path, bits_per_param=8): | |
| """ | |
| Convert HuggingFace model to .dq quantized format. | |
| Args: | |
| model_path: Path to local HuggingFace model directory | |
| output_path: Output path for .dq file | |
| bits_per_param: Target bit width (2, 3, 4, 5, 6, 8, 12, 16, 32) | |
| Format specification: | |
| Magic: 'DQMD' (4 bytes) | |
| Version: uint32 (4 bytes) | |
| Global metadata: length (uint32) + JSON bytes | |
| Tensor headers: length (uint32) + JSON bytes | |
| Tensor data: raw bytes per tensor in header order | |
| """ | |
| config = AutoConfig.from_pretrained(model_path) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_path, | |
| torch_dtype=torch.float32, | |
| low_cpu_mem_usage=True | |
| ) | |
| state_dict = model.state_dict() | |
| dq_data = { | |
| 'config': config.to_dict(), | |
| 'bits_per_param': bits_per_param, | |
| 'tensors': OrderedDict(), | |
| 'metadata': { | |
| 'model_path': model_path, | |
| 'architecture': config.architectures[0] if hasattr(config, 'architectures') else 'unknown' | |
| } | |
| } | |
| pack_functions = { | |
| 2: pack_2bit, | |
| 3: pack_3bit, | |
| 4: pack_4bit, | |
| 5: pack_5bit, | |
| 6: pack_6bit, | |
| 12: pack_12bit, | |
| } | |
| total_params = 0 | |
| with tqdm(total=len(state_dict), desc="Quantizing") as pbar: | |
| for name, tensor in state_dict.items(): | |
| total_params += tensor.numel() | |
| quantized_signed, scale, level, dtype_name = quantize_tensor_symmetric( | |
| tensor, bits_per_param | |
| ) | |
| tensor_info = { | |
| 'shape': list(tensor.shape), | |
| 'dtype': dtype_name, | |
| 'scale': scale, | |
| 'level': level, | |
| } | |
| if bits_per_param in pack_functions: | |
| unsigned = quantized_signed.to(torch.int32) + level | |
| if bits_per_param <= 8: | |
| unsigned = unsigned.to(torch.uint8) | |
| else: | |
| unsigned = unsigned.to(torch.uint16) | |
| packed, original_shape, pad_len = pack_functions[bits_per_param](unsigned) | |
| tensor_info['data'] = packed.numpy().tobytes() | |
| tensor_info['packed'] = True | |
| tensor_info['packed_shape'] = list(original_shape) | |
| tensor_info['pad'] = pad_len | |
| del unsigned, packed | |
| elif bits_per_param == 8: | |
| unsigned = (quantized_signed.to(torch.int32) + level).to(torch.uint8) | |
| tensor_info['data'] = unsigned.numpy().tobytes() | |
| tensor_info['packed'] = False | |
| del unsigned | |
| elif bits_per_param in (16, 32): | |
| tensor_info['data'] = quantized_signed.numpy().tobytes() | |
| tensor_info['packed'] = False | |
| dq_data['tensors'][name] = tensor_info | |
| del tensor, quantized_signed | |
| pbar.update(1) | |
| serializable_tensors = OrderedDict() | |
| for name, info in dq_data['tensors'].items(): | |
| serializable_tensors[name] = {k: v for k, v in info.items() if k != 'data'} | |
| serializable_tensors[name]['data_size'] = len(info['data']) | |
| with open(output_path, 'wb') as f: | |
| f.write(b'DQMD') | |
| f.write(struct.pack('<I', 1)) | |
| global_meta = { | |
| 'config': dq_data['config'], | |
| 'bits_per_param': dq_data['bits_per_param'], | |
| 'metadata': dq_data['metadata'] | |
| } | |
| global_meta_json = json.dumps(global_meta).encode('utf-8') | |
| f.write(struct.pack('<I', len(global_meta_json))) | |
| f.write(global_meta_json) | |
| tensor_headers_json = json.dumps(serializable_tensors).encode('utf-8') | |
| f.write(struct.pack('<I', len(tensor_headers_json))) | |
| f.write(tensor_headers_json) | |
| for info in dq_data['tensors'].values(): | |
| f.write(info['data']) | |
| original_size_mb = total_params * 4 / 1024 ** 2 | |
| compressed_size_mb = os.path.getsize(output_path) / 1024 ** 2 | |
| ratio = (total_params * 4) / os.path.getsize(output_path) | |
| return { | |
| 'total_params': total_params, | |
| 'bits_per_param': bits_per_param, | |
| 'original_size_mb': original_size_mb, | |
| 'compressed_size_mb': compressed_size_mb, | |
| 'compression_ratio': ratio, | |
| 'output_path': output_path | |
| } | |
| if __name__ == "__main__": | |
| result = convert_to_dq("modelHF", "modelDQ.dq", bits_per_param=8) | |
| print(f"Conversion complete: {result['compression_ratio']:.2f}x compression " | |
| f"({result['original_size_mb']:.1f}MB -> {result['compressed_size_mb']:.1f}MB)") |