Method forward

codegeex/quantization/quantize_oneflow.py:89–101

(self, input_)

            self.weight_zero = Parameter(self.weight_zero.to(kwargs["device"]), requires_grad=False)

    def forward(self, input_):
        # Matrix multiply.
        output = torch._C.fused_linear_with_groupwise_quantized_weight(input_, 
                                                                        w=self.weight, 
                                                                        w_scale=self.weight_scale, 
                                                                        w_zero=self.weight_zero, 
                                                                        b=self.bias if self.bias is not None else None, 
                                                                        num_bits=self.weight_bit_width,
                                                                        symmetric=self.symmetric,
                                                                        group_dim=self.group_dim,
                                                                        group_size=self.group_size)
        
        return output

def quantize_oneflow(model, weight_bit_width):
    """Replace fp16 linear with quantized linear"""