(
self,
in_features: int,
out_features: int,
weight_bit_width: int,
weight: torch.Tensor = None,
bias: torch.Tensor = None,
*args,
**kwargs
)
| 52 | |
| 53 | class QuantizedLinear(torch.nn.Module): |
| 54 | def __init__( |
| 55 | self, |
| 56 | in_features: int, |
| 57 | out_features: int, |
| 58 | weight_bit_width: int, |
| 59 | weight: torch.Tensor = None, |
| 60 | bias: torch.Tensor = None, |
| 61 | *args, |
| 62 | **kwargs |
| 63 | ): |
| 64 | super(QuantizedLinear, self).__init__() |
| 65 | |
| 66 | self.in_features = in_features |
| 67 | self.out_features = out_features |
| 68 | self.weight_bit_width = weight_bit_width |
| 69 | self.symmetric = True |
| 70 | self.group_dim = 1 |
| 71 | self.group_size = in_features |
| 72 | |
| 73 | self.weight, self.weight_scale, self.weight_zero = _quantize( |
| 74 | self.weight_bit_width, self.symmetric, weight, self.group_dim, self.group_size, torch.int8 |
| 75 | ) |
| 76 | if bias is None: |
| 77 | self.register_parameter('bias', None) |
| 78 | else: |
| 79 | self.bias = bias |
| 80 | self.bias = self.bias.to(kwargs["device"]) |
| 81 | |
| 82 | self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False) |
| 83 | self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False) |
| 84 | if self.bias is not None: |
| 85 | self.bias = Parameter(self.bias.to(kwargs["device"]), requires_grad=False) |
| 86 | if self.weight_zero is not None: |
| 87 | self.weight_zero = Parameter(self.weight_zero.to(kwargs["device"]), requires_grad=False) |
| 88 | |
| 89 | def forward(self, input_): |
| 90 | # Matrix multiply. |
nothing calls this directly
no test coverage detected