Revert "Add support for Int8DynActInt4WeightQuantizer (#66)"

This reverts commit 12b9194.

Author: cpuhrsch

test/quantization/test_quant_api.py

@@ -24,8 +24,6 @@
     Quantizer,
     TwoStepQuantizer,
     Int8DynActInt4WeightGPTQQuantizer,
-    Int8DynActInt4WeightQuantizer,
-    Int8DynActInt4WeightLinear,
 )
 from pathlib import Path
 from sentencepiece import SentencePieceProcessor
@@ -87,8 +85,8 @@ def quantize(self, model: torch.nn.Module) -> torch.nn.Module:
 class M(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.linear1 = torch.nn.Linear(64, 32).to(torch.float)
-        self.linear2 = torch.nn.Linear(32, 64).to(torch.float)
+        self.linear1 = torch.nn.Linear(5, 5).to(torch.float)
+        self.linear2 = torch.nn.Linear(5, 5).to(torch.float)
 
     def forward(self, x):
         x = self.linear1(x)
@@ -134,15 +132,8 @@ def test_dynamic_quant_gpu_unified_api_eager_mode_impl(self):
         compiled = m(*example_inputs)
         torch.testing.assert_close(quantized, compiled, atol=0, rtol=0)
 
-    def test_8da4w_quantizer(self):
-        quantizer = Int8DynActInt4WeightQuantizer(group_size=32)
-        m = M().eval()
-        m = quantizer.quantize(m)
-        assert isinstance(m.linear1, Int8DynActInt4WeightLinear)
-        assert isinstance(m.linear2, Int8DynActInt4WeightLinear)
-
     @unittest.skip("skipping until we get checkpoints for gpt-fast")
-    def test_gptq_quantizer(self):
+    def test_gptq(self):
         # should be similar to TorchCompileDynamicQuantizer
         precision = torch.bfloat16
         device = "cpu"

torchao/quantization/GPTQ.py

@@ -17,7 +17,6 @@
 import torch.nn.functional as F
 # from model import Transformer
 from torch.utils._pytree import tree_flatten, tree_unflatten
-import logging
 
 # pyre-fixme[5]: Global expression must be annotated.
 aten = torch.ops.aten
@@ -62,7 +61,7 @@ def model_forward(model, x, input_pos):
         get_task_dict = tasks.get_task_dict
         evaluate = evaluator.evaluate
 else:
-    logging.info("lm_eval is not installed, GPTQ may not be usable")
+    print("lm_eval is not installed, GPTQ may not be usable")
 
 # pyre-fixme[3]: Return type must be annotated.
 def setup_cache_padded_seq_input_pos_max_seq_length_for_prefill(

torchao/quantization/quant_api.py

@@ -33,10 +33,8 @@
 from .quant_primitives import (
     get_group_qparams_symmetric,
     per_token_dynamic_quant,
-    group_quantize_tensor_symmetric,
 )
-from typing import Dict, Tuple, Any
-import logging
+from typing import Dict, Tuple
 
 __all__ = [
     "apply_weight_only_int8_quant",
@@ -52,18 +50,21 @@
 ############################# Unified Quantization APIs ##############################
 # API 1, single quantize call to create a quantized model with quantized state_dict
 class Quantizer:
-    def quantize(self, model: torch.nn.Module, *args: Any, **kwargs: Any) -> torch.nn.Module:
+    # pyre-fixme[2]: Parameter must be annotated.
+    def quantize(self, model: torch.nn.Module, *args, **kwargs) -> torch.nn.Module:
         # pyre-fixme[7]: Expected `Module` but got implicit return value of `None`.
         pass
 
 
 # API 2, flow that needs calibration or training
 class TwoStepQuantizer:
-    def prepare(self, model: torch.nn.Module, *args: Any, **kwargs: Any) -> torch.nn.Module:
+    # pyre-fixme[2]: Parameter must be annotated.
+    def prepare(self, model: torch.nn.Module, *args, **kwargs) -> torch.nn.Module:
         # pyre-fixme[7]: Expected `Module` but got implicit return value of `None`.
         pass
 
-    def convert(self, model: torch.nn.Module, *args: Any, **kwargs: Any) -> torch.nn.Module:
+    # pyre-fixme[2]: Parameter must be annotated.
+    def convert(self, model: torch.nn.Module, *args, **kwargs) -> torch.nn.Module:
         # pyre-fixme[7]: Expected `Module` but got implicit return value of `None`.
         pass
 
@@ -249,7 +250,7 @@ def replace_conv2d_1x1(conv):
         MultiInput,
     )
 else:
-    logging.info("lm_eval not available, skip defining GPTQQuantizer")
+    print("lm_eval not available, skip defining GPTQQuantizer")
 
 
 class GPTQQuantizer(Quantizer):
@@ -431,7 +432,11 @@ def _convert_for_runtime(self, model: torch.nn.Module) -> "nn.Module":
 
     @torch.no_grad()
     # pyre-fixme[14]: `quantize` overrides method defined in `Quantizer` inconsistently.
-    def quantize(self, model: torch.nn.Module, **kwargs: Any) -> torch.nn.Module:
+    def quantize(
+        self,
+        # pyre-fixme[2]: Parameter must be annotated.
+        model,
+    ) -> torch.nn.Module:
         state_dict = self._create_quantized_state_dict(
             model,
             # pyre-fixme[16]: `GPTQQuantizer` has no attribute `tokenizer`.
@@ -665,91 +670,6 @@ def replace_linear_8da4w(
             )
 
 
-class Int8DynActInt4WeightQuantizer(Quantizer):
-    def __init__(
-        self,
-        group_size: int = 256,
-        padding_allowed: bool = False,
-        precision: torch.dtype = torch.float32,
-        scales_precision: torch.dtype = torch.float32,
-    ) -> None:
-        self.group_size: int = group_size
-        self.padding_allowed: bool = padding_allowed
-        self.precision: torch.dtype = precision
-        self.scales_precision: torch.dtype = scales_precision
-        # assert group_size in [32, 64, 128, 256]
-
-    @torch.no_grad()
-    def _create_quantized_state_dict(self, model: torch.nn.Module) -> Dict[str, torch.Tensor]:
-        cur_state_dict = model.state_dict()
-        for fqn, mod in model.named_modules():
-            if isinstance(mod, torch.nn.Linear):
-                assert mod.bias is not None
-                out_features = mod.out_features
-                in_features = mod.in_features
-                # assert out_features % 8 == 0, "require out_features % 8 == 0"
-                print(f"linear: {fqn}, in={in_features}, out={out_features}")
-
-                assert (
-                    in_features % self.group_size == 0
-                ), f"require in_features:{in_features} % self.group_size:{self.group_size} == 0"
-
-                weight = mod.weight.data
-                """
-                if not _check_linear_int4_k(
-                    in_features, self.group_size
-                ):
-                    if self.padding_allowed:
-                        print(
-                            f"warning: {fqn} is padded to satisfy in_features % 1024 == 0"
-                        )
-                        padded_in_features = _calc_padded_size_linear_int4(
-                            in_features, self.group_size
-                        )
-                        weight = F.pad(
-                            weight, pad=(0, padded_in_features - in_features)
-                        )
-                    else:
-                        raise RuntimeError(
-                            f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, "
-                            + "and that group_size"
-                        )
-                """
-                (
-                    weight_int8,
-                    scales,
-                    zeros,
-                ) = group_quantize_tensor_symmetric(
-                    weight.to(self.precision),
-                    4,  # n_bit
-                    self.group_size,
-                    self.scales_precision,
-                )
-                cur_state_dict[f"{fqn}.weight"] = weight_int8.to("cpu")
-                cur_state_dict[f"{fqn}.scales"] = scales.to("cpu")
-                cur_state_dict[f"{fqn}.zeros"] = zeros.to("cpu")
-                # TODO: support bias?
-
-        return cur_state_dict
-
-    def _convert_for_runtime(self, model: torch.nn.Module) -> torch.nn.Module:
-        replace_linear_8da4w(
-            model,
-            self.group_size,
-            self.padding_allowed,
-            self.precision,
-            self.scales_precision,
-        )
-        return model
-
-    def quantize(self, model: torch.nn.Module, *args: Any, **kwargs: Any) -> torch.nn.Module:
-        state_dict = self._create_quantized_state_dict(model)
-        model = self._convert_for_runtime(model)
-        # TODO: make it strict
-        model.load_state_dict(state_dict, strict=False)
-        return model
-
-
 class Int8DynActInt4WeightGPTQQuantizer(GPTQQuantizer):
     # pyre-fixme[3]: Return type must be annotated.
     def __init__(