lvkaokao
update codes.
5a7ab71
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
@dataclasses.dataclass
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)