|
import dataclasses |
|
|
|
import torch |
|
from torch import Tensor |
|
import torch.nn as nn |
|
from torch.nn import functional as F |
|
|
|
|
|
@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, bias, device): |
|
super().__init__() |
|
|
|
self.weight = compress(weight.data.to(device), default_compression_config) |
|
self.bias = bias |
|
|
|
def forward(self, input: Tensor) -> Tensor: |
|
weight = decompress(self.weight, default_compression_config) |
|
return F.linear(input, 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 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_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) |
|
|
|
|
|
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, |
|
) |
|
|
|
|
|
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) |
|
|
|
|
|
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) |
|
|
