WisdomShell/CodeShell-7B-Chat-int4

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    if self.quant_state is not None:26        self.quant_state[0] = self.quant_state[0].cuda(device)27        self.quant_state[6] = self.quant_state[6].cuda(device)28    return self29 30def Params4bitTo(self, *args, **kwargs):31    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)32 33    if (device is not None and device.type == "cuda" and self.data.device.type == "cpu"):34        return self.cuda(device)35    else:36        if self.quant_state is not None:37            # make sure the quantization state is on the right device38            self.quant_state[0] = self.quant_state[0].to(device)39            self.quant_state[6] = self.quant_state[6].to(device)40        new_param = Params4bit(self.to(device=device, dtype=dtype, non_blocking=non_blocking),41                                requires_grad=self.requires_grad, quant_state=self.quant_state,42                                blocksize=self.blocksize, compress_statistics=self.compress_statistics,43                                quant_type=self.quant_type)44 45    return new_param46 47class Linear4bitOnline(torch.nn.Module):48    def __init__(self, weight, bias, quant_type):49        super().__init__()50        self.weight = Params4bit(51            weight.data, requires_grad=False, compress_statistics=True, quant_type=quant_type52        )53        self.compute_dtype = None54        #self.weight.cuda(weight.device)55        self.bias = bias56 57    def forward(self, x: torch.Tensor):58        # weights are cast automatically as Int8Params, but the bias has to be cast manually59        if self.bias is not None and self.bias.dtype != x.dtype:60            self.bias.data = self.bias.data.to(x.dtype)61 62        if getattr(self.weight, "quant_state", None) is None:63            print(64                "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first."65            )66        inp_dtype = x.dtype67        if self.compute_dtype is not None:68            x = x.to(self.compute_dtype)69 70        bias = None if self.bias is None else self.bias.to(self.compute_dtype)71        out = bnb.matmul_4bit(72            x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state73        )74 75        out = out.to(inp_dtype)76 77        return out78 79class Linear8bitLtOnline(torch.nn.Module):80    def __init__(81        self,82        weight,83        bias,84        has_fp16_weights=True,85        memory_efficient_backward=False,86        threshold=0.0,87        index=None,88    ):89        super().__init__()90        assert (91            not memory_efficient_backward92        ), "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0"93        self.state = bnb.MatmulLtState()94        self.index = index95 96        # Necessary for stacked layers97        self.state.threshold = threshold98        self.state.has_fp16_weights = has_fp16_weights99        self.state.memory_efficient_backward = memory_efficient_backward100        if threshold > 0.0 and not has_fp16_weights:101            self.state.use_pool = True102 103        self.weight = Int8Params(104            weight.data,105            has_fp16_weights=has_fp16_weights,106            requires_grad=has_fp16_weights,107        )108        self.bias = bias109 110    def init_8bit_state(self):111        self.state.CB = self.weight.CB112        self.state.SCB = self.weight.SCB113        self.weight.CB = None114        self.weight.SCB = None115 116    def forward(self, x: torch.Tensor):117        self.state.is_training = self.training118        if self.weight.CB is not None:119            self.init_8bit_state()120 121        # weights are cast automatically as Int8Params, but the bias has to be cast manually122        if self.bias is not None and self.bias.dtype != x.dtype:123            self.bias.data = self.bias.data.to(x.dtype)124 125        out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)126 127        if not self.state.has_fp16_weights:128            if self.state.CB is not None and self.state.CxB is not None:129                # we converted 8-bit row major to turing/ampere format in the first inference pass130                # we no longer need the row-major weight131                del self.state.CB132                self.weight.data = self.state.CxB133        return out134 135def quantize_online(model, bits: int):136    def quant(weight, bias=None):137        if bits == 8:138            linear = Linear8bitLtOnline(139                weight,140                bias,141                has_fp16_weights=False,142                threshold=6.0,143            )144            if bias is not None:145                linear.bias = torch.nn.Parameter(bias)146        elif bits == 4:147            linear = Linear4bitOnline(148                weight,149                bias,150                quant_type="nf4", #fp4/nf4151            )152        else:153            raise ValueError("quantize only support 4/8 bit")154        return linear155 156    def auto_quant(layer):157        if hasattr(layer,"bias"):158            linear = quant(layer.weight,bias=layer.bias)159        else:160            linear = quant(layer.weight)161        return linear162 163    for i,layer in enumerate(model.transformer.h):164        layer.mlp.c_fc = auto_quant(layer.mlp.c_fc)165        layer.mlp.c_proj = auto_quant(layer.mlp.c_proj)166 167        layer.attn.c_attn=auto_quant(layer.attn.c_attn)168        layer.attn.c_proj=auto_quant(layer.attn.c_proj)169 170    return model171 172 173general_weight_dict = {174    "transformer.wte.weight": False,175    "transformer.ln_f.weight": False,176    "transformer.ln_f.bias": False,177    "lm_head.weight": False,178}179 180layer_weight_dict = {181    "transformer.h.{i}.ln_1.weight": False,182    "transformer.h.{i}.ln_1.bias": False,183    "transformer.h.{i}.attn.c_attn.weight": True,184    "transformer.h.{i}.attn.c_attn.bias": False,185    "transformer.h.{i}.attn.c_proj.weight": True,186    "transformer.h.{i}.attn.c_proj.bias": False,187    "transformer.h.{i}.attn.rotary_emb.inv_freq": False,188    "transformer.h.{i}.ln_2.weight": False,189    "transformer.h.{i}.ln_2.bias": False,190    "transformer.h.{i}.mlp.c_fc.weight": True,191    "transformer.h.{i}.mlp.c_fc.bias": False,192    "transformer.h.{i}.mlp.c_proj.weight": True,193    "transformer.h.{i}.mlp.c_proj.bias": False,194}195num_dict = {str(i):i for i in range(100)}196 197def set_value(model, name, state_dict, is_4bit):198    keys = name.split('.')199    parent = model200    for key in keys[:-1]:201        if key in num_dict:202            parent = parent[num_dict[key]]203        else:204            parent = getattr(parent, key)205    if is_4bit:206        weight_data = state_dict[f'{name}.data']207        weight_quant_state = state_dict[f'{name}.quant_state']208        assert weight_data is not None, name209        assert weight_quant_state is not None, name210        setattr(parent, keys[-1], Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state))211    else:212        setattr(parent, keys[-1], state_dict[name])213 214def quantize_offline(model):215    for i, layer in enumerate(model.transformer.h):216        layer.mlp.c_fc = bnb.nn.Linear4bit(217                            layer.mlp.c_fc.weight.shape[1],218                            layer.mlp.c_fc.weight.shape[0],219                            False,220                            torch.bfloat16,221                            compress_statistics=True,222                            quant_type="nf4",223                        )224        layer.mlp.c_proj = bnb.nn.Linear4bit(225                            layer.mlp.c_proj.weight.shape[1],226                            layer.mlp.c_proj.weight.shape[0],227                            False,228                            torch.bfloat16,229                            compress_statistics=True,230                            quant_type="nf4",231                        )232 233        layer.attn.c_attn = bnb.nn.Linear4bit(234                            layer.attn.c_attn.weight.shape[1],235                            layer.attn.c_attn.weight.shape[0],236                            False,237                            torch.bfloat16,238                            compress_statistics=True,239                            quant_type="nf4",240                        )241        layer.attn.c_proj = bnb.nn.Linear4bit(242                            layer.attn.c_proj.weight.shape[1],243                            layer.attn.c_proj.weight.shape[0],244                            False,245                            torch.bfloat16,246                            compress_statistics=True,247                            quant_type="nf4",248                        )249    return model250 251def load_state_dict_for_qunantied_model(model, state_dict):252    #replace Params4bit.cuda with Params4bitCuda253    Params4bit.cuda = Params4bitCuda254    Params4bit.to = Params4bitTo255 256    for name, is_4bit in general_weight_dict.items():257        set_value(model, name, state_dict, is_4bit)258                259    for layer_i in range(len(model.transformer.h)):260        for name, is_4bit in layer_weight_dict.items():261            name = name.replace('{i}', str(layer_i))262            set_value(model, name, state_dict, is_4bit)263    return model