(
self,
in_features: int,
out_features: int,
weight_bit_width: int,
weight: torch.Tensor = None,
bias: torch.Tensor = None,
*args,
**kwargs
)
| 31 | |
| 32 | class QuantizedLinear(torch.nn.Module): |
| 33 | def __init__( |
| 34 | self, |
| 35 | in_features: int, |
| 36 | out_features: int, |
| 37 | weight_bit_width: int, |
| 38 | weight: torch.Tensor = None, |
| 39 | bias: torch.Tensor = None, |
| 40 | *args, |
| 41 | **kwargs |
| 42 | ): |
| 43 | super(QuantizedLinear, self).__init__() |
| 44 | |
| 45 | self.in_features = in_features |
| 46 | self.out_features = out_features |
| 47 | self.weight_bit_width = weight_bit_width |
| 48 | |
| 49 | if weight is None: |
| 50 | self.weight = torch.empty( |
| 51 | self.out_features, self.in_features * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"] |
| 52 | ) |
| 53 | self.weight_scale = torch.empty(self.out_features, dtype=kwargs["params_dtype"], device=kwargs["device"]) |
| 54 | else: |
| 55 | self.weight_scale = (weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half() |
| 56 | self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8) |
| 57 | if weight_bit_width == 4: |
| 58 | self.weight = compress_int4_weight(self.weight) |
| 59 | |
| 60 | if bias is None: |
| 61 | self.register_parameter('bias', None) |
| 62 | else: |
| 63 | self.bias = bias |
| 64 | |
| 65 | self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False) |
| 66 | self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False) |
| 67 | |
| 68 | def forward(self, input_): |
| 69 | # Matrix multiply. |
no test coverage detected