File size: 3,017 Bytes
9c41926 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import torch
import numpy as np
def ternary_quantize(x: torch.Tensor):
scale = x.abs().max()
if scale == 0:
scale = 1.0
normalized = x / scale
ternary = torch.sign(normalized)
ternary[ternary == 0] = 1
return ternary.to(torch.int8), scale
def ternary_pack(t: torch.Tensor) -> torch.Tensor:
encoded = t.to(torch.int8) + 1
n = encoded.numel()
pad = (5 - n % 5) % 5
if pad:
encoded = torch.cat([encoded.flatten(), torch.zeros(pad, dtype=torch.int8, device=encoded.device)])
weights = torch.tensor([81, 27, 9, 3, 1], dtype=torch.int32, device=encoded.device)
packed = (encoded.reshape(-1, 5).to(torch.int32) * weights).sum(dim=1)
return packed.to(torch.uint8)
def ternary_unpack(packed: torch.Tensor, original_shape: tuple) -> torch.Tensor:
weights = torch.tensor([81, 27, 9, 3, 1], dtype=torch.int32, device=packed.device)
expanded = packed.to(torch.int32).unsqueeze(-1) // weights % 3
flat = (expanded - 1).flatten()
return flat[:np.prod(original_shape)].reshape(original_shape).to(torch.int8)
def sigma_quantize(s: torch.Tensor, num_bits: int = 2):
max_val = s.abs().max()
if max_val == 0:
max_val = 1.0
qmax = 2 ** (num_bits - 1) - 1
scale = max_val / qmax
quantized = (s / scale).round().clamp(-qmax, qmax)
return quantized.to(torch.int8), scale
def sigma_dequantize(q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
return q.float() * scale
def quantize_factor(U: torch.Tensor, S: torch.Tensor, Vt: torch.Tensor, sigma_bits: int = 2):
U_q, U_scale = ternary_quantize(U)
Vt_q, Vt_scale = ternary_quantize(Vt)
S_q, S_scale = sigma_quantize(S, sigma_bits)
return {
'U': U_q, 'U_scale': U_scale,
'Vt': Vt_q, 'Vt_scale': Vt_scale,
'S': S_q, 'S_scale': S_scale,
}
def pack_factor(data: dict):
return {
'U_packed': ternary_pack(data['U']),
'U_scale': data['U_scale'],
'Vt_packed': ternary_pack(data['Vt']),
'Vt_scale': data['Vt_scale'],
'S': data['S'],
'S_scale': data['S_scale'],
'U_shape': data['U'].shape,
'Vt_shape': data['Vt'].shape,
}
def unpack_factor(data: dict):
if 'U_packed' in data:
U = ternary_unpack(data['U_packed'], data['U_shape'])
Vt = ternary_unpack(data['Vt_packed'], data['Vt_shape'])
else:
U = data['U']
Vt = data['Vt']
return {
'U': U,
'U_scale': data['U_scale'],
'Vt': Vt,
'Vt_scale': data['Vt_scale'],
'S': data['S'],
'S_scale': data['S_scale'],
}
def dequantize_factor(data: dict) -> torch.Tensor:
if 'U_packed' in data:
U = ternary_unpack(data['U_packed'], data['U_shape']).float()
Vt = ternary_unpack(data['Vt_packed'], data['Vt_shape']).float()
else:
U = data['U'].float()
Vt = data['Vt'].float()
S = sigma_dequantize(data['S'], data['S_scale'])
return torch.matmul(U * S.unsqueeze(0), Vt)
|