Fix import versions for GPTQ (#105)

Author: jerryzh168

test/quantization/model.py

@@ -12,6 +12,7 @@
 from torch.nn import functional as F
 
 def prepare_inputs_for_model(inps):
+    inps = inps.squeeze(0)
     # setup inputs in correct format
     max_new_tokens = 1
     T = inps.size(0)

torchao/quantization/GPTQ.py

@@ -20,7 +20,7 @@
 # from model import Transformer  # pyre-ignore[21]
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
-from .utils import TORCH_VERSION_AFTER_2_4
+from .utils import TORCH_VERSION_AFTER_2_3
 from typing import Any, Dict, Tuple, Optional
 from .unified import Quantizer
 from functools import reduce
@@ -89,7 +89,7 @@ def __init__(
             # for model
             self.input_prep_func = (
                 input_prep_func if input_prep_func is not None
-                else lambda x: x
+                else lambda x: (x,)
             )
 
             self.pad_calibration_inputs = pad_calibration_inputs
@@ -180,6 +180,7 @@ def _model_call(self, inps):
             else:
                 inps = F.pad(inps, (self.pad_token, self.calibration_seq_length - T))
 
+            inps = inps.unsqueeze(0)
             model_in = self.input_prep_func(inps)
 
             self.add_input(model_in)
@@ -546,7 +547,7 @@ def faster_quant(self, H, W):
         return Q, DQ.to(orig_dtype), all_qparams
 
 
-if TORCH_VERSION_AFTER_2_4:
+if TORCH_VERSION_AFTER_2_3:
     from .quant_primitives import (
         get_group_qparams_symmetric,
         group_quantize_tensor_symmetric,

torchao/quantization/quant_api.py

@@ -46,10 +46,14 @@
     from .GPTQ import (
         Int8DynActInt4WeightQuantizer,
         Int8DynActInt4WeightGPTQQuantizer,
+        Int4WeightQuantizer,
+        Int4WeightGPTQQuantizer,
     )
     __all__ += [
         "Int8DynActInt4WeightQuantizer",
         "Int8DynActInt4WeightGPTQQuantizer",
+        "Int4WeightQuantizer",
+        "Int4WeightGPTQQuantizer",
     ]
 
 
@@ -196,6 +200,3 @@ def replace_conv2d_1x1(conv):
     _replace_with_custom_fn_if_matches_filter(
         model, replace_conv2d_1x1, filter_fn=filter_fn
     )
-
-if TORCH_VERSION_AFTER_2_3:
-    from .GPTQ import Int8DynActInt4WeightQuantizer, Int8DynActInt4WeightGPTQQuantizer