Spaces:
Runtime error
Runtime error
import dataclasses | |
import gc | |
import glob | |
import os | |
from accelerate import init_empty_weights | |
from accelerate.utils import set_module_tensor_to_device | |
import torch | |
from torch import Tensor | |
import torch.nn as nn | |
from torch.nn import functional as F | |
from tqdm import tqdm | |
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig | |
class CompressionConfig: | |
"""Group-wise quantization.""" | |
num_bits: int | |
group_size: int | |
group_dim: int | |
symmetric: bool | |
enabled: bool = True | |
default_compression_config = CompressionConfig( | |
num_bits=8, group_size=256, group_dim=1, symmetric=True, enabled=True | |
) | |
class CLinear(nn.Module): | |
"""Compressed Linear Layer.""" | |
def __init__(self, weight=None, bias=None, device=None): | |
super().__init__() | |
if weight is None: | |
self.weight = None | |
elif isinstance(weight, Tensor): | |
self.weight = compress(weight.data.to(device), default_compression_config) | |
else: | |
self.weight = weight | |
self.bias = bias | |
def forward(self, input: Tensor) -> Tensor: | |
weight = decompress(self.weight, default_compression_config) | |
return F.linear(input.to(weight.dtype), weight, self.bias) | |
def compress_module(module, target_device): | |
for attr_str in dir(module): | |
target_attr = getattr(module, attr_str) | |
if type(target_attr) == torch.nn.Linear: | |
setattr( | |
module, | |
attr_str, | |
CLinear(target_attr.weight, target_attr.bias, target_device), | |
) | |
for name, child in module.named_children(): | |
compress_module(child, target_device) | |
def get_compressed_list(module, prefix=''): | |
compressed_list = [] | |
for attr_str in dir(module): | |
target_attr = getattr(module, attr_str) | |
if type(target_attr) == torch.nn.Linear: | |
full_name = f"{prefix}.{attr_str}.weight" if prefix else f"{attr_str}.weight" | |
compressed_list.append(full_name) | |
for name, child in module.named_children(): | |
child_prefix = f"{prefix}.{name}" if prefix else name | |
for each in get_compressed_list(child, child_prefix): | |
compressed_list.append(each) | |
return compressed_list | |
def apply_compressed_weight(module, compressed_state_dict, target_device, prefix=''): | |
for attr_str in dir(module): | |
target_attr = getattr(module, attr_str) | |
if type(target_attr) == torch.nn.Linear: | |
full_name = f"{prefix}.{attr_str}.weight" if prefix else f"{attr_str}.weight" | |
setattr(module, attr_str, | |
CLinear(compressed_state_dict[full_name], target_attr.bias, target_device)) | |
for name, child in module.named_children(): | |
child_prefix = f"{prefix}.{name}" if prefix else name | |
apply_compressed_weight(child, compressed_state_dict, target_device, child_prefix) | |
def load_compress_model(model_path, device, torch_dtype): | |
# partially load model | |
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) | |
base_pattern = os.path.join(model_path, "pytorch_model-*.bin") | |
files = glob.glob(base_pattern) | |
with init_empty_weights(): | |
config = AutoConfig.from_pretrained(model_path, low_cpu_mem_usage=True, | |
torch_dtype=torch_dtype) | |
model = AutoModelForCausalLM.from_config(config) | |
linear_weights = get_compressed_list(model) | |
compressed_state_dict = {} | |
for filename in tqdm(files): | |
tmp_state_dict = torch.load(filename) | |
for name in tmp_state_dict: | |
if name in linear_weights: | |
tensor = tmp_state_dict[name].to(device).data.to(torch_dtype) | |
compressed_state_dict[name] = compress(tensor, default_compression_config) | |
else: | |
compressed_state_dict[name] = tmp_state_dict[name].to(device) | |
tmp_state_dict[name] = None | |
tensor = None | |
gc.collect() | |
torch.cuda.empty_cache() | |
for name in model.state_dict(): | |
if name not in linear_weights: | |
set_module_tensor_to_device(model, name, device, value=compressed_state_dict[name]) | |
apply_compressed_weight(model, compressed_state_dict, device) | |
model.to(device) | |
return model, tokenizer | |
def compress(tensor, config): | |
"""Simulate group-wise quantization.""" | |
if not config.enabled: | |
return tensor | |
group_size, num_bits, group_dim, symmetric = ( | |
config.group_size, | |
config.num_bits, | |
config.group_dim, | |
config.symmetric, | |
) | |
assert num_bits <= 8 | |
original_shape = tensor.shape | |
num_groups = (original_shape[group_dim] + group_size - 1) // group_size | |
new_shape = ( | |
original_shape[:group_dim] | |
+ (num_groups, group_size) | |
+ original_shape[group_dim + 1 :] | |
) | |
# Pad | |
pad_len = (group_size - original_shape[group_dim] % group_size) % group_size | |
if pad_len != 0: | |
pad_shape = ( | |
original_shape[:group_dim] + (pad_len,) + original_shape[group_dim + 1 :] | |
) | |
tensor = torch.cat( | |
[tensor, torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device)], | |
dim=group_dim, | |
) | |
data = tensor.view(new_shape) | |
# Quantize | |
if symmetric: | |
B = 2 ** (num_bits - 1) - 1 | |
scale = B / torch.max(data.abs(), dim=group_dim + 1, keepdim=True)[0] | |
data = data * scale | |
data = data.clamp_(-B, B).round_().to(torch.int8) | |
return data, scale, original_shape | |
else: | |
B = 2**num_bits - 1 | |
mn = torch.min(data, dim=group_dim + 1, keepdim=True)[0] | |
mx = torch.max(data, dim=group_dim + 1, keepdim=True)[0] | |
scale = B / (mx - mn) | |
data = data - mn | |
data.mul_(scale) | |
data = data.clamp_(0, B).round_().to(torch.uint8) | |
return data, mn, scale, original_shape | |
def decompress(packed_data, config): | |
"""Simulate group-wise dequantization.""" | |
if not config.enabled: | |
return packed_data | |
group_size, num_bits, group_dim, symmetric = ( | |
config.group_size, | |
config.num_bits, | |
config.group_dim, | |
config.symmetric, | |
) | |
# Dequantize | |
if symmetric: | |
data, scale, original_shape = packed_data | |
data = data / scale | |
else: | |
data, mn, scale, original_shape = packed_data | |
data = data / scale | |
data.add_(mn) | |
# Unpad | |
pad_len = (group_size - original_shape[group_dim] % group_size) % group_size | |
if pad_len: | |
padded_original_shape = ( | |
original_shape[:group_dim] | |
+ (original_shape[group_dim] + pad_len,) | |
+ original_shape[group_dim + 1 :] | |
) | |
data = data.reshape(padded_original_shape) | |
indices = [slice(0, x) for x in original_shape] | |
return data[indices].contiguous() | |
else: | |
return data.view(original_shape) | |