Revert "fixing over padding and GPTQ padding bug"

This reverts commit 5bf70c1.

Breaks llama-70B + int4 + TP

Author: Chillee

model.py

@@ -4,18 +4,15 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 from torch.nn import functional as F
-from math import gcd
-from functools import reduce
 
 
-def find_multiple(n: int, *args: Tuple[int]) -> int:
-    k = reduce(lambda x,y: x*y//gcd(x,y), args+(1,))
+def find_multiple(n: int, k: int) -> int:
     if n % k == 0:
         return n
     return n + k - (n % k)

quantize.py

@@ -17,7 +17,7 @@
 except:
     pass
 
-from model import Transformer, find_multiple
+from model import Transformer
 
 ##### Quantization Primitives ######
 
@@ -376,27 +376,29 @@ def linear_forward_int4(x, weight_int4pack, scales_and_zeros, out_features, grou
 def _check_linear_int4_k(k, groupsize = 1, inner_k_tiles = 1):
     return k % groupsize == 0 and k % (inner_k_tiles * 16) == 0
 
-def _calc_padded_size_linear_int4(k, groupsize = 1, inner_k_tiles = 1):
-    return find_multiple(k, groupsize, inner_k_tiles*16)
-
-def replace_linear_int4(module, groupsize, inner_k_tiles, padding_allowed):
+def replace_linear_int4(module, groupsize, inner_k_tiles, padding):
     for name, child in module.named_children():
         if isinstance(child, nn.Linear):
-            if _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles) or padding_allowed:
+            if _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles):
                 setattr(module, name, WeightOnlyInt4Linear(
                     child.in_features, child.out_features, bias=False,
-                    groupsize=groupsize, inner_k_tiles=inner_k_tiles,
+                    groupsize=groupsize, inner_k_tiles=inner_k_tiles, padding=False,
+                ))
+            elif padding:
+                setattr(module, name, WeightOnlyInt4Linear(
+                    child.in_features, child.out_features, bias=False,
+                    groupsize=groupsize, inner_k_tiles=inner_k_tiles, padding=True,
                 ))
         else:
-            replace_linear_int4(child, groupsize, inner_k_tiles, padding_allowed)
+            replace_linear_int4(child, groupsize, inner_k_tiles, padding)
 
 
 class WeightOnlyInt4QuantHandler:
-    def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding_allowed=True):
+    def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding=True):
         self.mod = mod
         self.groupsize = groupsize
         self.inner_k_tiles = inner_k_tiles
-        self.padding_allowed = padding_allowed
+        self.padding = padding
         assert groupsize in [32, 64, 128, 256]
         assert inner_k_tiles in [2, 4, 8]
 
@@ -418,9 +420,11 @@ def create_quantized_state_dict(self, use_cuda = True):
 
                 weight = mod.weight.data
                 if not _check_linear_int4_k(in_features, self.groupsize, self.inner_k_tiles):
-                    if self.padding_allowed:
+                    if self.padding:
+                        from model import find_multiple
+                        import torch.nn.functional as F
                         print(f"warning: {fqn} is padded to satisfy in_features % 1024 == 0")
-                        padded_in_features = _calc_padded_size_linear_int4(in_features, 1024)
+                        padded_in_features = find_multiple(in_features, 1024)
                         weight = F.pad(weight, pad=(0, padded_in_features - in_features))
                     else:
                         print(f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " +
@@ -435,30 +439,31 @@ def create_quantized_state_dict(self, use_cuda = True):
         return cur_state_dict
 
     def convert_for_runtime(self):
-        replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding_allowed)
+        replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding)
         return self.mod
 
 class WeightOnlyInt4GPTQQuantHandler(GPTQQuantHandler):
-    def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding_allowed=True):
+    def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding=True):
+        from model import find_multiple
         self.mod = mod
         self.groupsize = groupsize
         self.inner_k_tiles = inner_k_tiles
-        self.padding_allowed = padding_allowed
+        self.padding = padding
         self.get_qparams_func = lambda w: get_group_qparams(w, 4, groupsize)
         self.quantize_func = lambda w, qparams: \
             group_quantize_tensor_from_qparams(w, qparams[0], qparams[1], 4, groupsize)
         self.dequantize_func = lambda q, qparams: \
             group_dequantize_tensor_from_qparams(q, qparams[0], qparams[1], 4, groupsize).float()
         self.combine_qparams_list_func = lambda qparams_list: \
             [torch.cat(x, dim=1) for x in zip(*qparams_list)]
-        # skip unless padding_allowed=True or its correctly sized
+        # skip unless padding=True or its correctly sized
         self.skip_layer_func = lambda linear_weight: not (
-            _check_linear_int4_k(linear_weight.shape[-1], groupsize, inner_k_tiles) or padding_allowed
+            _check_linear_int4_k(linear_weight.shape[-1], groupsize, inner_k_tiles) or padding
         )
         # we need to do the padding here, both for q and the qparams if necessary
         def make_names_and_values_dict_func(q, qparams):
             k = q.shape[1]
-            new_k = _calc_padded_size_linear_int4(k, groupsize, inner_k_tiles)
+            new_k = find_multiple(k, 1024)
             # how much we need to pad the weight
             delta_k = new_k - q.shape[1]
             final_q = torch.ops.aten._convert_weight_to_int4pack(F.pad(q, pad=(0, delta_k)), inner_k_tiles)
@@ -472,7 +477,7 @@ def make_names_and_values_dict_func(q, qparams):
 
 
     def convert_for_runtime(self):
-        replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding_allowed)
+        replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding)
         return self.mod
 
 class WeightOnlyInt4Linear(torch.nn.Module):
@@ -483,16 +488,17 @@ class WeightOnlyInt4Linear(torch.nn.Module):
 
     def __init__(
             self, in_features: int, out_features: int,
-            bias=True, device=None, dtype=None, groupsize: int = 128, inner_k_tiles: int = 8,
+            bias=True, device=None, dtype=None, groupsize: int = 128, inner_k_tiles: int = 8, padding: bool = True,
     ) -> None:
         super().__init__()
+        self.padding = padding
+        if padding:
+            from model import find_multiple
+            self.origin_in_features = in_features
+            in_features = find_multiple(in_features, 1024)
 
-        # always pad if needed since it becomes a noop at runtime if not needed
-        self.origin_in_features = in_features
-        in_features = _calc_padded_size_linear_int4(in_features, groupsize, inner_k_tiles)
         self.in_features = in_features
         self.out_features = out_features
-
         assert not bias, "require bias=False"
         self.groupsize = groupsize
         self.inner_k_tiles = inner_k_tiles
@@ -510,7 +516,9 @@ def __init__(
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         input = input.to(torch.bfloat16)
-        input = F.pad(input, pad=(0, self.in_features - self.origin_in_features))
+        if self.padding:
+            import torch.nn.functional as F
+            input = F.pad(input, pad=(0, self.in_features - self.origin_in_features))
         return linear_forward_int4(
             input,
             self.weight, self.scales_and_zeros, self.out_features, self.groupsize