File size: 6,479 Bytes
2051791 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import torch
def pack_4bit_to_int8(quantized_4bit_weights):
# Ensure the quantized weights are uint8
quantized_4bit_weights = quantized_4bit_weights.to(torch.uint8)
# Reshape to ensure even number of elements per row
if quantized_4bit_weights.size(1) % 2 != 0:
quantized_4bit_weights = torch.cat((quantized_4bit_weights, torch.zeros((quantized_4bit_weights.size(0), 1), dtype=torch.uint8)), dim=1)
# Packing two 4-bit values into one int8
packed_weights = (quantized_4bit_weights[:, ::2] << 4) | (quantized_4bit_weights[:, 1::2] & 0xF)
return packed_weights
def unpack_int8_to_4bit(packed_weights):
# Unpack two 4-bit values from one int8
high_bits = (packed_weights >> 4) & 0xF
low_bits = packed_weights & 0xF
# Interleave the high and low bits
unpacked_weights = torch.stack((high_bits, low_bits), dim=2).flatten(1)
return unpacked_weights
def dequantize_4bit_to_fp16(quantized, maxima, minima, nbits, blocksize):
shape = quantized.shape
if quantized.numel() % (blocksize * blocksize) == 0:
quantized = quantized.reshape(-1, blocksize * blocksize)
scale = (2 ** nbits - 1) / (maxima - minima)
dequantized = (quantized / scale) + minima
dequantized = dequantized.reshape(shape)
else:
scale = (2 ** nbits - 1) / (maxima - minima)
dequantized = (quantized / scale) + minima
return dequantized
def round_to_nearest_pole(x, poles):
# Compute the absolute differences between each element of x and all poles
differences = torch.abs(x.unsqueeze(-1) - poles)
# Find the index of the minimum difference along the poles dimension
nearest_indices = torch.argmin(differences, dim=-1)
# Return the corresponding nearest pole values for each element in x
nearest_values = poles[nearest_indices]
return nearest_values
def quantize_blockwise(weight, nbits, blocksize, clip_method, clip_threshold, scale_shift=False, use_normal_float=False):
# weight should be param.data
shape = weight.shape
dtype = weight.dtype
_num_outliers = 0
if weight.numel() % (blocksize) == 0:
weight = weight.reshape(-1, blocksize)
# block wise
if clip_method != "no":
if clip_method == "block_percentage":
num_top_elements = max(int(weight.size(1) * clip_threshold + 1), 1) # per block
threshold = torch.topk(weight.abs(), num_top_elements, dim=1).values[:, -1].unsqueeze(-1)
elif clip_method == "tensor_percentage":
num_top_elements = max(int(weight.numel() * clip_threshold), 1)
threshold = torch.topk(weight.view(-1).abs(), num_top_elements).values[-1]
elif clip_method == "zscore":
means = weight.abs().mean(dim=1, keepdim=True)
stds = weight.abs().std(dim=1, keepdim=True)
threshold = clip_threshold * stds + means
elif clip_method == "iqr":
q1 = weight.abs().float().quantile(0.25, dim=1, keepdim=True)
q3 = weight.abs().float().quantile(0.75, dim=1, keepdim=True)
threshold = q3 + clip_threshold * (q3 - q1) # 1.5 * IQR
threshold = threshold.to(weight.dtype)
else:
raise ValueError(f"Unknown clip method: {clip_method}")
_num_outliers = int(torch.sum(weight.abs() > threshold ))
weight = torch.clamp(weight, -threshold, threshold)
minima, _ = weight.min(dim=1, keepdims=True)
maxima, _ = weight.max(dim=1, keepdims=True)
else:
if clip_method != "no":
if clip_method == "tensor_percentage" or clip_method == "block_percentage":
num_top_elements = max(int(weight.numel() * clip_threshold), 1)
threshold = torch.topk(weight.view(-1).abs(), num_top_elements).values[-1]
elif clip_method == "zscore":
mean = weight.abs().mean()
std = weight.abs().std()
threshold = clip_threshold * std + mean
elif clip_method == "iqr":
q1 = weight.abs().float().quantile(0.25)
q3 = weight.abs().float().quantile(0.75)
threshold = q3 + clip_threshold * (q3 - q1) # 1.5 * IQR
# count clipped outliers
_num_outliers = int(torch.sum(weight.abs() > threshold))
weight = torch.clamp(weight, -threshold, threshold)
minima, maxima = weight.min(), weight.max()
if use_normal_float: # NF4 or NF3
NF4_LEVELS = [
-1.0,
-0.6961928009986877,
-0.5250730514526367,
-0.39491748809814453,
-0.28444138169288635,
-0.18477343022823334,
-0.09105003625154495,
0.0,
0.07958029955625534,
0.16093020141124725,
0.24611230194568634,
0.33791524171829224,
0.44070982933044434,
0.5626170039176941,
0.7229568362236023,
1.0,
]
NF3_LEVELS = [-1, -0.5350227355957031, -
0.2469314038753510, 0, 0.1833375245332718, 0.3819939494132996, 0.6229856610298157, 1]
if nbits == 4:
quantization_levels = NF4_LEVELS
elif nbits == 3:
quantization_levels = NF3_LEVELS
else:
raise ValueError("Normal Float Quantization only suuports 4 and 3 bits now.")
quantization_levels = torch.tensor(quantization_levels).to(weight.device)
scale = 2 / (maxima - minima) # scale to [0, 2]
weight.sub_(minima).mul_(scale).sub_(1.0) # shift to [-1, 1]
weight.copy_(round_to_nearest_pole(weight, quantization_levels))
weight.add_(1).div_(scale).add_(minima)
else: # INT
if scale_shift:
# mapping weights to [-0.4999, 15.4999] and then round to nearest integers
scale = (2 ** nbits - 0.01) / (maxima - minima)
weight.sub_(minima).mul_(scale).sub_(0.49)
weight.round_()
weight.add_(0.49).div_(scale).add_(minima)
else:
# mapping weights to [0, 15] and then round to nearest integers
scale = (2 ** nbits - 1) / (maxima - minima)
weight.sub_(minima).mul_(scale)
weight.round_()
weight.div_(scale).add_(minima)
dequantized = weight.reshape(shape).to(dtype)
return dequantized, _num_outliers
|