Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| 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) | |