| import os | |
| os.environ["PATH"] = "/usr/local/cuda/bin:" + os.environ["PATH"] | |
| os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" | |
| import bitblas | |
| import torch | |
| import torch.nn as nn | |
| from transformers import BertConfig, BertModel, PreTrainedModel, PretrainedConfig,AutoModel,AutoConfig,BertPreTrainedModel | |
| class bitlinear(bitblas.Linear): | |
| def __init__( | |
| self, | |
| in_features: int, | |
| out_features: int, | |
| bias: bool = False, | |
| A_dtype: str = "float16", | |
| W_dtype: str = "int2", | |
| accum_dtype: str = "float16", | |
| out_dtype: str = "float16", | |
| group_size: int = -1, | |
| with_scaling: bool = False, | |
| with_zeros: bool = False, | |
| zeros_mode: str = None, | |
| opt_M: list = [1, 16, 32, 64, 128, 256, 512], | |
| fast_decoding: bool = True, | |
| alpha: torch.dtype = torch.float16, | |
| b:torch.Tensor=None | |
| ): | |
| super().__init__( | |
| in_features=in_features, | |
| out_features=out_features, | |
| bias=bias, | |
| A_dtype=A_dtype, | |
| W_dtype=W_dtype, | |
| accum_dtype=accum_dtype, | |
| out_dtype=out_dtype, | |
| group_size=group_size, | |
| with_scaling=with_scaling, | |
| with_zeros=with_zeros, | |
| zeros_mode=zeros_mode, | |
| opt_M=opt_M, | |
| fast_decoding=fast_decoding, | |
| ) | |
| self.alpha = nn.Parameter(alpha,requires_grad=False) | |
| self.b = nn.Parameter(b,requires_grad=False) | |
| def forward(self, A: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor: | |
| out = super().forward(A, out) | |
| out *= self.alpha | |
| if self.b is not None: | |
| out += self.b.view(1, -1).expand_as(out) | |
| return out.to(torch.float32) | |
| class TernaryBertConfig(BertConfig): | |
| model_type = "ternarybert" | |
| def __init__(self, **kwargs): | |
| super().__init__(**kwargs) | |
| class TernaryBert(PreTrainedModel): | |
| #config_class = TernaryBertConfig | |
| config_class = BertConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.bert = BertModel(config) | |
| self.replace_linear2bitblas(self.bert) | |
| #def forward(self, input_ids, attention_mask=None,token_type_ids=None): | |
| # return self.bert(input_ids, attention_mask=attention_mask,token_type_ids=token_type_ids) | |
| def forward(self, **kwargs): | |
| return self.bert(**kwargs) | |
| def convert_to_bitlinear(self,layer): | |
| bitlayer = bitlinear( | |
| in_features=layer.in_features, | |
| out_features=layer.out_features, | |
| bias=False, | |
| A_dtype="float16", # activation A dtype | |
| W_dtype="int2", # weight W dtype | |
| accum_dtype="float16", # accumulation dtype | |
| out_dtype="float16", # output dtype | |
| # configs for weight only quantization | |
| group_size=-1, # setting for grouped quantization | |
| with_scaling=False, # setting for scaling factor | |
| with_zeros=False, # setting for zeros | |
| zeros_mode=None, # setting for how to calculating zeros | |
| # Target optimization var for dynamic symbolic. | |
| # For detailed information please checkout docs/PythonAPI.md | |
| # By default, the optimization var is [1, 16, 32, 64, 128, 256, 512] | |
| opt_M=[1, 16, 32, 64, 128, 256, 512], | |
| fast_decoding=True, | |
| alpha=torch.tensor(1.).to(torch.float16), | |
| b = layer.bias.data.to(torch.float16) | |
| ) | |
| return bitlayer | |
| def replace_linear2bitblas(self,model): | |
| for name, module in model.named_children(): | |
| if isinstance(module, nn.Linear): | |
| new_layer = self.convert_to_bitlinear(module) | |
| setattr(model, name, new_layer) | |
| elif len(list(module.children())) > 0: | |
| self.replace_linear2bitblas(module) | |