| 51 | ) |
| 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. |
| 91 | output = torch._C.fused_linear_with_groupwise_quantized_weight(input_, |
| 92 | w=self.weight, |
| 93 | w_scale=self.weight_scale, |
| 94 | w_zero=self.weight_zero, |
| 95 | b=self.bias if self.bias is not None else None, |
| 96 | num_bits=self.weight_bit_width, |
| 97 | symmetric=self.symmetric, |
| 98 | group_dim=self.group_dim, |
| 99 | group_size=self.group_size) |
| 100 | |
| 101 | return output |
| 102 | |
| 103 | def quantize_oneflow(model, weight_bit_width): |
| 104 | """Replace fp16 linear with quantized linear""" |