Upload modeling.py with huggingface_hub

modeling.py

@@ -0,0 +1,107 @@
+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)
+