| 48 | return model, has_been_replaced |
| 49 | |
| 50 | class QLinear(nn.Linear): |
| 51 | def __init__(self, input_features, output_features, bias=True, compute_dtype=torch.bfloat16, quant_type="ste-n2f3", q_group_size=128, device=None): |
| 52 | super().__init__(input_features, output_features, bias, device) |
| 53 | |
| 54 | if quant_type == "ste-n2f3": |
| 55 | self.weight_quantizer = SteN2F3Quantizer(q_group_size=q_group_size) |
| 56 | elif quant_type == "int2-asym": |
| 57 | self.weight_quantizer = SteInt2AsymQuantizer(q_group_size=q_group_size) |
| 58 | else: |
| 59 | raise ValueError(f"Has no support {quant_type}. Valid quant_type:[ste-n2f3, int2-asym]") |
| 60 | # self.quant_type = quant_type |
| 61 | self.compute_dtype = compute_dtype |
| 62 | |
| 63 | def forward(self, x: torch.Tensor): |
| 64 | if self.bias is not None and self.bias.dtype != x.dtype: |
| 65 | self.bias.data = self.bias.data.to(x.dtype) |
| 66 | |
| 67 | inp_dtype = x.dtype |
| 68 | |
| 69 | if self.compute_dtype is not None: |
| 70 | x = x.to(self.compute_dtype) |
| 71 | |
| 72 | bias = None if self.bias is None else self.bias.to(self.compute_dtype) |
| 73 | out = None |
| 74 | |
| 75 | quantize_weight = self.weight_quantizer(self.weight.to(self.compute_dtype)) |
| 76 | out = F.linear(x, quantize_weight, bias).to(inp_dtype) |
| 77 | |
| 78 | return out |
no outgoing calls
no test coverage detected