Fix padding not used correctly in exllama v2 layer (#626)

padding not used correctly in exllama
AutoGPTQ · Apr 5, 2024 · b4b801c · b4b801c
1 parent 866b4c8
commit b4b801c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/auto_gptq/nn_modules/qlinear/qlinear_exllama.py b/auto_gptq/nn_modules/qlinear/qlinear_exllama.py
@@ -58,8 +58,11 @@ def __init__(self, bits, group_size, infeatures, outfeatures, bias, trainable=Fa
         if trainable:
             raise NotImplementedError("Exllama kernel does not support training.")
 
+        self.padding = -outfeatures % 32
+        self.outfeatures = outfeatures + self.padding
+        outfeatures = self.outfeatures
+
         self.infeatures = infeatures
-        self.outfeatures = outfeatures
         self.bits = bits
         self.group_size = group_size if group_size != -1 else infeatures
         self.trainable = trainable

diff --git a/auto_gptq/nn_modules/qlinear/qlinear_exllamav2.py b/auto_gptq/nn_modules/qlinear/qlinear_exllamav2.py
@@ -121,10 +121,12 @@ def __init__(self, bits, group_size, infeatures, outfeatures, bias, trainable=Fa
 
         self.q_handle = None
         self.q_tensors = None
+
         self.padding = -outfeatures % 32
+        self.outfeatures = outfeatures + self.padding
+        outfeatures = self.outfeatures
 
         self.infeatures = infeatures
-        self.outfeatures = outfeatures + self.padding
         self.bits = bits
         self.group_size = group_size if group_size != -1 else infeatures
         self.trainable = trainable