Spaces:
Sleeping
Sleeping
File size: 3,443 Bytes
4e316d6 | 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 | """
INT8 Absmax Weight Quantization.
Maps float weights to 8-bit integers using symmetric scaling:
scale = max(|W|) / 127
W_int8 = round(W / scale)
W_reconstructed = W_int8 * scale
Per-tensor mode uses one scale for the entire weight matrix.
Per-channel mode uses one scale per output channel (row), which
preserves more precision since each row is scaled independently —
a single outlier row no longer degrades the entire tensor.
INT8 halves memory vs float16/bfloat16. The tradeoff is a small
accuracy loss from rounding, measurable as a perplexity increase.
"""
import torch
import torch.nn as nn
def absmax_quantize(tensor, per_channel=False):
"""
Quantize a float tensor to INT8 using absmax scaling.
Args:
tensor: Float tensor (typically a weight matrix, shape: out_features x in_features)
per_channel: If True, compute one scale per row (output channel)
Returns:
(quantized_int8_tensor, scale)
"""
if per_channel:
scale = tensor.abs().amax(dim=1, keepdim=True) / 127
else:
scale = tensor.abs().max() / 127
scale = torch.clamp(scale, min=1e-8)
quantized = torch.clamp(torch.round(tensor / scale), -127, 127).to(torch.int8)
return quantized, scale
def absmax_dequantize(quantized, scale):
"""Reconstruct float tensor from INT8 + scale."""
return quantized.float() * scale
class QuantizedLinear(nn.Module):
"""Drop-in replacement for nn.Linear that stores INT8 weights and dequantizes on forward."""
def __init__(self, in_features, out_features, weight_int8, scale, bias=None):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.register_buffer("weight_int8", weight_int8)
self.register_buffer("scale", scale)
if bias is not None:
self.register_buffer("bias", bias)
else:
self.bias = None
@classmethod
def from_linear(cls, linear, per_channel=False):
"""Create a QuantizedLinear from an existing nn.Linear."""
weight_int8, scale = absmax_quantize(linear.weight.data.float(), per_channel=per_channel)
bias = linear.bias.data if linear.bias is not None else None
return cls(linear.in_features, linear.out_features, weight_int8, scale, bias)
def forward(self, x):
w = absmax_dequantize(self.weight_int8, self.scale).to(x.dtype)
return nn.functional.linear(x, w, self.bias)
def quantize_model(model, per_channel=False):
"""Walk the model and replace all nn.Linear layers with QuantizedLinear.
Skips Linear layers whose weights are tied to an Embedding (e.g. a
shared output head), since quantizing them would break the tie and
produce a mismatch between embedding lookups and output projections.
"""
tied_ptrs = {
m.weight.data_ptr()
for m in model.modules()
if isinstance(m, nn.Embedding)
}
_quantize_recursive(model, per_channel, tied_ptrs)
def _quantize_recursive(module, per_channel, tied_ptrs=frozenset()):
for name, child in module.named_children():
if isinstance(child, nn.Linear):
if child.weight.data_ptr() in tied_ptrs:
continue # skip weight-tied layers
setattr(module, name, QuantizedLinear.from_linear(child, per_channel=per_channel))
else:
_quantize_recursive(child, per_channel, tied_ptrs)
|