| | 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)
|
| |
|
