GPTQ integration

src/transformers/modeling_utils.py

@@ -64,6 +64,7 @@
     download_url,
     has_file,
     is_accelerate_available,
+    is_auto_gptq_available,
     is_bitsandbytes_available,
     is_offline_mode,
     is_optimum_available,
@@ -75,7 +76,7 @@
 )
 from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files
 from .utils.import_utils import ENV_VARS_TRUE_VALUES, is_sagemaker_mp_enabled
-from .utils.quantization_config import BitsAndBytesConfig
+from .utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
 from .utils.versions import require_version_core
 
 
@@ -2344,15 +2345,35 @@ def from_pretrained(
         else:
             model_kwargs = kwargs
 
+        # get the quantization method inside the config of the model if it exists
+        quantization_method = None
+        if hasattr(config, "quantization_config"):
+            if config.quantization_config.get("load_in_8bit", False):
+                quantization_method = QuantizationMethod.BITS_AND_BYTES
+            else:
+                quantization_method = config.quantization_config.get("quant_method", None)
+
+        quantizer = None
+        if quantization_method == QuantizationMethod.GPTQ:
+            if not (is_optimum_available() and is_auto_gptq_available()):
+                raise ImportError(
+                    "Loading GTPQ quantized model requires optimum library : `pip install optimum` and auto-gptq library 'pip install auto-gptq'"
+                )
+            else:
+                # Need to protect the import
+                from optimum.gptq import GPTQQuantizer
+            quantizer = GPTQQuantizer.from_dict(config.quantization_config)
+            torch_dtype = config.torch_dtype
+
         if is_8bit_serializable and quantization_config is not None and load_in_8bit:
-            if hasattr(config, "quantization_config"):
+            if quantization_method == QuantizationMethod.BITS_AND_BYTES:
                 logger.warning(
                     "You passed `quantization_config` to `from_pretrained` but the model you're loading already has a"
                     " `quantization_config` attribute. The `quantization_config` attribute will be overwritten with the"
                     " one you passed to `from_pretrained`."
                 )
             config.quantization_config = quantization_config
-        elif is_8bit_serializable and not load_in_8bit and hasattr(config, "quantization_config"):
+        elif is_8bit_serializable and not load_in_8bit and quantization_method == QuantizationMethod.BITS_AND_BYTES:
             quantization_config = config.quantization_config
             if isinstance(quantization_config, dict):
                 quantization_config = BitsAndBytesConfig.from_dict(quantization_config, return_unused_kwargs=False)
@@ -2382,7 +2403,9 @@ def from_pretrained(
                     if low_cpu_mem_usage is None:
                         low_cpu_mem_usage = True
 
-        elif not is_8bit_serializable and not load_in_8bit and hasattr(config, "quantization_config"):
+        elif (
+            not is_8bit_serializable and not load_in_8bit and quantization_method == QuantizationMethod.BITS_AND_BYTES
+        ):
             logger.warning(
                 "Detected the presence of a `quantization_config` attribute in the model's configuration but you don't have the correct"
                 " `bitsandbytes` version to support int8 serialization. Please install the latest version of `bitsandbytes` with "
@@ -2767,6 +2790,8 @@ def from_pretrained(
                 "All non-linear modules will be loaded in full precision."
                 " If you want to load the other modules in other precision, please specify a `torch_dtype` attribute."
             )
+        if quantization_method == QuantizationMethod.GPTQ:
+            model = quantizer.convert_model(model)
 
         if isinstance(device_map, str):
             special_dtypes = {}

src/transformers/testing_utils.py

@@ -51,6 +51,7 @@
 from .utils import (
     is_accelerate_available,
     is_apex_available,
+    is_auto_gptq_available,
     is_bitsandbytes_available,
     is_bs4_available,
     is_cython_available,
@@ -770,6 +771,13 @@ def require_optimum(test_case):
     return unittest.skipUnless(is_optimum_available(), "test requires optimum")(test_case)
 
 
+def require_auto_gptq(test_case):
+    """
+    Decorator for auto_gptq dependency
+    """
+    return unittest.skipUnless(is_auto_gptq_available(), "test requires auto-gptq")(test_case)
+
+
 def require_phonemizer(test_case):
     """
     Decorator marking a test that requires phonemizer

src/transformers/utils/__init__.py

@@ -103,6 +103,7 @@
     get_torch_version,
     is_accelerate_available,
     is_apex_available,
+    is_auto_gptq_available,
     is_bitsandbytes_available,
     is_bs4_available,
     is_coloredlogs_available,

src/transformers/utils/import_utils.py

@@ -98,6 +98,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[
 _onnx_available = _is_package_available("onnx")
 _openai_available = _is_package_available("openai")
 _optimum_available = _is_package_available("optimum")
+_auto_gptq_available = _is_package_available("auto_gptq")
 _pandas_available = _is_package_available("pandas")
 _peft_available = _is_package_available("peft")
 _phonemizer_available = _is_package_available("phonemizer")
@@ -554,6 +555,10 @@ def is_optimum_available():
     return _optimum_available
 
 
+def is_auto_gptq_available():
+    return _auto_gptq_available
+
+
 def is_optimum_neuron_available():
     return _optimum_available and _is_package_available("optimum.neuron")
 

src/transformers/utils/quantization_config.py

@@ -19,6 +19,7 @@
 import json
 import os
 from dataclasses import dataclass
+from enum import Enum
 from typing import Any, Dict, Union
 
 from packaging import version
@@ -33,6 +34,11 @@
 logger = logging.get_logger(__name__)
 
 
+class QuantizationMethod(Enum):
+    BITS_AND_BYTES = "bitsandbytes"
+    GPTQ = "gtpq"
+
+
 @dataclass
 class BitsAndBytesConfig:
     """
@@ -97,6 +103,7 @@ def __init__(
         bnb_4bit_use_double_quant=False,
         **kwargs,
     ):
+        self.quant_method = QuantizationMethod.BITS_AND_BYTES
         self.load_in_8bit = load_in_8bit
         self.load_in_4bit = load_in_4bit
         self.llm_int8_threshold = llm_int8_threshold

tests/quantization/gptq/test_gptq.py

@@ -0,0 +1,145 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tempfile
+import unittest
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.testing_utils import (
+    is_torch_available,
+    require_accelerate,
+    require_auto_gptq,
+    require_optimum,
+    require_torch_gpu,
+    require_torch_multi_gpu,
+    slow,
+)
+
+
+if is_torch_available():
+    import torch
+
+
+@slow
+@require_optimum
+@require_auto_gptq
+@require_torch_gpu
+class GTPQTest(unittest.TestCase):
+    model_name = "bigscience/bloom-560m"
+
+    input_text = "Hello my name is"
+    EXPECTED_OUTPUT = "Hello my name is John and I am a professional photographer. I"
+
+    # this seems a little small considering that we are doing 4bit quant but we have a small model and ww don't quantize the embeddings
+    EXPECTED_RELATIVE_DIFFERENCE = 1.664253062
+
+    bits = 4
+    group_size = 128
+    desc_act = False
+
+    dataset = [
+        "auto-gptq is an easy-to-use model quantization library with user-friendly apis, based on GPTQ algorithm."
+    ]
+
+    device_map = None
+
+    # called only once for all test in this class
+    @classmethod
+    def setUpClass(cls):
+        from optimum.gptq import GPTQQuantizer
+
+        """
+        Setup quantized model
+        """
+        cls.model_fp16 = AutoModelForCausalLM.from_pretrained(
+            cls.model_name, torch_dtype=torch.float16, device_map=cls.device_map
+        )
+        cls.mem_fp16 = cls.model_fp16.get_memory_footprint()
+
+        cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name, use_fast=True)
+        cls.quantizer = GPTQQuantizer(bits=cls.bits, group_size=cls.group_size, desc_act=cls.desc_act)
+
+        cls.quantized_model = cls.quantizer.quantize_model(cls.model_fp16, cls.tokenizer, cls.dataset)
+
+    def test_memory_footprint(self):
+        r"""
+        A simple test to check if the model conversion has been done correctly by checking on the
+        memory footprint of the converted model
+        """
+
+        mem_quantized = self.quantized_model.get_memory_footprint()
+
+        self.assertAlmostEqual(self.mem_fp16 / mem_quantized, self.EXPECTED_RELATIVE_DIFFERENCE)
+
+    def test_quantized_layers_class(self):
+        """
+        Simple test to check if the model conversion has been done correctly by checking on
+        the class type of the linear layers of the converted models
+        """
+        from auto_gptq.utils.import_utils import dynamically_import_QuantLinear
+
+        QuantLinear = dynamically_import_QuantLinear(
+            use_triton=False, desc_act=self.desc_act, group_size=self.group_size
+        )
+        self.assertTrue(self.quantized_model.transformer.h[0].mlp.dense_4h_to_h.__class__ == QuantLinear)
+
+    def check_inference_correctness(self, model):
+        r"""
+        Test the generation quality of the quantized model and see that we are matching the expected output.
+        Given that we are operating on small numbers + the testing model is relatively small, we might not get
+        the same output across GPUs. So we'll generate few tokens (5-10) and check their output.
+        """
+        # Check that inference pass works on the model
+        encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
+
+        # Check the exactness of the results
+        output_parallel = model.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
+
+        # Get the generation
+        self.assertEqual(self.tokenizer.decode(output_parallel[0], skip_special_tokens=True), self.EXPECTED_OUTPUT)
+
+    def test_generate_quality(self):
+        """
+        Simple test to check the quality of the model by comapring the the generated tokens with the expected tokens
+        """
+        if self.device_map is None:
+            self.check_inference_correctness(self.quantized_model.to(0))
+        else:
+            self.check_inference_correctness(self.quantized_model)
+
+    def test_serialization(self):
+        """
+        Test the serialization of the model and the loading of the quantized weights works
+        """
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            self.quantized_model.to("cpu").save_pretrained(tmpdirname)
+            quantized_model_from_saved = AutoModelForCausalLM.from_pretrained(tmpdirname).to(0)
+            self.check_inference_correctness(quantized_model_from_saved)
+
+    @require_accelerate
+    def test_serialization_big_model_inference(self):
+        """
+        Test the serialization of the model and the loading of the quantized weights with big model inference
+        """
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            self.quantized_model.to("cpu").save_pretrained(tmpdirname)
+            quantized_model_from_saved = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map="auto")
+            self.check_inference_correctness(quantized_model_from_saved)
+
+
+@require_accelerate
+@require_torch_multi_gpu
+class GTPQTestDeviceMap(GTPQTest):
+    device_map = "auto"