WisdomShell/CodeShell-7B-Chat

1# coding=utf-82# Copyright 2023 WisdomShell Inc. All Rights Reserved.3 4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16try:17    import bitsandbytes as bnb18    from bitsandbytes.nn.modules import Params4bit, Int8Params19except ImportError:20    pass21import torch22 23def Params4bitCuda(self, device):24    self.data = self.data.cuda(device)25    self.quant_state[0] = self.quant_state[0].cuda(device)26    self.quant_state[4][0] = self.quant_state[4][0].cuda(device)27    self.quant_state[4][1][0] = self.quant_state[4][1][0].cuda(device)28    self.quant_state[4][1][1] = self.quant_state[4][1][1].cuda(device)29 30    self.quant_state[6] = self.quant_state[6].cuda(device)31    return self32 33class Linear4bitOnline(torch.nn.Module):34    def __init__(self, weight, bias, quant_type):35        super().__init__()36        self.weight = Params4bit(37            weight.data, requires_grad=False, compress_statistics=True, quant_type=quant_type38        )39        self.compute_dtype = None40        #self.weight.cuda(weight.device)41        self.bias = bias42 43    def forward(self, x: torch.Tensor):44        # weights are cast automatically as Int8Params, but the bias has to be cast manually45        if self.bias is not None and self.bias.dtype != x.dtype:46            self.bias.data = self.bias.data.to(x.dtype)47 48        if getattr(self.weight, "quant_state", None) is None:49            print(50                "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first."51            )52        inp_dtype = x.dtype53        if self.compute_dtype is not None:54            x = x.to(self.compute_dtype)55 56        bias = None if self.bias is None else self.bias.to(self.compute_dtype)57        out = bnb.matmul_4bit(58            x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state59        )60 61        out = out.to(inp_dtype)62 63        return out64 65class Linear8bitLtOnline(torch.nn.Module):66    def __init__(67        self,68        weight,69        bias,70        has_fp16_weights=True,71        memory_efficient_backward=False,72        threshold=0.0,73        index=None,74    ):75        super().__init__()76        assert (77            not memory_efficient_backward78        ), "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0"79        self.state = bnb.MatmulLtState()80        self.index = index81 82        # Necessary for stacked layers83        self.state.threshold = threshold84        self.state.has_fp16_weights = has_fp16_weights85        self.state.memory_efficient_backward = memory_efficient_backward86        if threshold > 0.0 and not has_fp16_weights:87            self.state.use_pool = True88 89        self.weight = Int8Params(90            weight.data,91            has_fp16_weights=has_fp16_weights,92            requires_grad=has_fp16_weights,93        )94        self.bias = bias95 96    def init_8bit_state(self):97        self.state.CB = self.weight.CB98        self.state.SCB = self.weight.SCB99        self.weight.CB = None100        self.weight.SCB = None101 102    def forward(self, x: torch.Tensor):103        self.state.is_training = self.training104        if self.weight.CB is not None:105            self.init_8bit_state()106 107        # weights are cast automatically as Int8Params, but the bias has to be cast manually108        if self.bias is not None and self.bias.dtype != x.dtype:109            self.bias.data = self.bias.data.to(x.dtype)110 111        out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)112 113        if not self.state.has_fp16_weights:114            if self.state.CB is not None and self.state.CxB is not None:115                # we converted 8-bit row major to turing/ampere format in the first inference pass116                # we no longer need the row-major weight117                del self.state.CB118                self.weight.data = self.state.CxB119        return out120 121def quantize_online(model, bits: int):122    def quant(weight, bias=None):123        if bits == 8:124            linear = Linear8bitLtOnline(125                weight,126                bias,127                has_fp16_weights=False,128                threshold=6.0,129            )130            if bias is not None:131                linear.bias = torch.nn.Parameter(bias)132        elif bits == 4:133            linear = Linear4bitOnline(134                weight,135                bias,136                quant_type="nf4", #fp4/nf4137            )138        else:139            raise ValueError("quantize only support 4/8 bit")140        return linear141 142    def auto_quant(layer):143        if hasattr(layer,"bias"):144            linear = quant(layer.weight,bias=layer.bias)145        else:146            linear = quant(layer.weight)147        return linear148 149    for i,layer in enumerate(model.transformer.h):150        layer.mlp.c_fc = auto_quant(layer.mlp.c_fc)151        layer.mlp.c_proj = auto_quant(layer.mlp.c_proj)152 153        layer.attn.c_attn=auto_quant(layer.attn.c_attn)154        layer.attn.c_proj=auto_quant(layer.attn.c_proj)155 156    return model