| import math | |
| import torch | |
| from torch import nn | |
| def weight_quant(weight, num_bits=1): | |
| dtype = weight.dtype | |
| weight = weight.float() | |
| s = 1 / weight.abs().mean().clamp(min=1e-5) | |
| result = (weight * s).round().clamp(-1, 1) / s | |
| return result.type(dtype) | |
| def activation_quant(x, num_bits=8): | |
| dtype = x.dtype | |
| x = x.float() | |
| Qn = -2 ** (num_bits - 1) | |
| Qp = 2 ** (num_bits - 1) - 1 | |
| s = Qp / x.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5) | |
| result = (x * s).round().clamp(Qn, Qp) / s | |
| return result.type(dtype) | |
| class BitLinear(nn.Linear): | |
| def __init__(self, | |
| *kargs, | |
| weight_bits=1, | |
| input_bits=8, | |
| **kwargs | |
| ): | |
| super(BitLinear, self).__init__(*kargs, **kwargs) | |
| """ | |
| RMSNorm is placed outside BitLinear | |
| """ | |
| self.weight_bits = weight_bits | |
| self.input_bits = input_bits | |
| def forward(self, input): | |
| quant_input = input + (activation_quant(input, self.input_bits) - input).detach() | |
| quant_weight = self.weight + (weight_quant(self.weight, self.weight_bits) - self.weight).detach() | |
| out = nn.functional.linear(quant_input, quant_weight) | |
| if not self.bias is None: | |
| out += self.bias.view(1, -1).expand_as(out) | |
| return out |