Update src/pipeline.py

src/pipeline.py

@@ -17,14 +17,19 @@ import os
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01"
 Pipeline = None
 
+def w8_a16_forward(weight, input, scales, bias=None):
+    casted_weights = weight.to(input.dtype)
+    output = F.linear(input, casted_weights) * scales # overhead
+    if bias is not None:
+        output = output + bias
+    return output
+    
 class W8A16LinearLayer(nn.Module):
     def __init__(self, in_features, out_features, bias=True, dtype=torch.float32):
         super().__init__()
-        
         self.register_buffer(
             "int8_weights",
             torch.randint(-128, 127, (out_features, in_features), dtype=torch.int8))
-        
         self.register_buffer("scales", torch.randn((out_features), dtype=dtype))
         
         if bias:
@@ -33,14 +38,12 @@ class W8A16LinearLayer(nn.Module):
     
     def quantize(self, weights):
         w_fp32 = weights.clone().to(torch.float32)
-
         scales = w_fp32.abs().max(dim=-1).values / 127
         scales = scales.to(weights.dtype)
-
         int8_weights = torch.round(weights/scales.unsqueeze(1)).to(torch.int8)
-
         self.int8_weights = int8_weights
-        self.scales = scales            self.bias = None
+        self.scales = scales
+        self.bias = None
 
     def forward(self, input):
         return w8_a16_forward(self.int8_weights, input, self.scales, self.bias)
@@ -51,12 +54,9 @@ def replace_linear_with_target_and_quantize(module, target_class, module_name_to
         any([x == name for x in module_name_to_exclude]):
             old_bias = child.bias
             old_weight = child.weight
-
             new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype)
             setattr(module, name, new_module)
-
             getattr(module, name).quantize(old_weight)
-            
             if old_bias is not None:
               getattr(module, name).bias = old_bias
         else: