Fixing starcoder based models with 15B (#118)

awq/models/gpt_bigcode.py

@@ -21,6 +21,7 @@ class GptBigCodeAWQForCausalLM(BaseAWQForCausalLM):
     @staticmethod
     def move_embed(model: GPTBigCodeForCausalLM, device):
         model.transformer.wte = model.transformer.wte.to(device)
+        model.transformer.wpe = model.transformer.wpe.to(device)
         model.transformer.drop = model.transformer.drop.to(device)
 
     @staticmethod

awq/quantize/quantizer.py

@@ -318,6 +318,9 @@ class AwqQuantizer:
         self.awq_model.move_embed(self.model, "cpu")
         
         clear_memory()
+        
+        if "attention_mask" in layer_kwargs.keys():
+            layer_kwargs["attention_mask"] = layer_kwargs["attention_mask"].to("cuda")
 
         return modules, layer_kwargs, inps
     

awq/quantize/scale.py

@@ -108,4 +108,4 @@ def scale_gelu_fc(gelu: allowed_act_fns, fc: nn.Linear, scales: torch.Tensor):
     fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))
 
     for p in fc.parameters():
-        assert torch.isnan(p).sum() == 0
+        assert torch.isnan(p).sum() == 0
\ No newline at end of file