[core / Quantization ] AWQ integration

src/transformers/__init__.py

@@ -772,7 +772,7 @@
         "is_vision_available",
         "logging",
     ],
-    "utils.quantization_config": ["BitsAndBytesConfig", "GPTQConfig"],
+    "utils.quantization_config": ["AWQConfig", "BitsAndBytesConfig", "GPTQConfig"],
 }
 
 # sentencepiece-backed objects
@@ -4923,7 +4923,7 @@
     )
 
     # bitsandbytes config
-    from .utils.quantization_config import BitsAndBytesConfig, GPTQConfig
+    from .utils.quantization_config import AWQConfig, BitsAndBytesConfig, GPTQConfig
 
     try:
         if not is_sentencepiece_available():

src/transformers/integrations/__init__.py

@@ -17,6 +17,7 @@
 
 
 _import_structure = {
+    "awq": ["replace_with_awq_linear"],
     "bitsandbytes": [
         "get_keys_to_not_convert",
         "replace_8bit_linear",
@@ -77,6 +78,7 @@
 }
 
 if TYPE_CHECKING:
+    from .awq import replace_with_awq_linear
     from .bitsandbytes import (
         get_keys_to_not_convert,
         replace_8bit_linear,

src/transformers/integrations/awq.py

@@ -0,0 +1,72 @@
+# 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.
+from ..utils import is_accelerate_available, is_auto_awq_available
+
+
+if is_auto_awq_available():
+    import torch.nn as nn
+    from awq.modules.linear import WQLinear_GEMM
+
+if is_accelerate_available():
+    from accelerate import init_empty_weights
+
+
+def replace_with_awq_linear(
+    model, modules_to_not_convert=None, current_key_name=None, quantization_config=None, has_been_replaced=False
+):
+    """
+    Private method that wraps the recursion for module replacement.
+
+    Returns the converted model and a boolean that indicates if the conversion has been successfull or not.
+    """
+    if modules_to_not_convert is None:
+        modules_to_not_convert = []
+
+    for name, module in model.named_children():
+        if current_key_name is None:
+            current_key_name = []
+        current_key_name.append(name)
+
+        if isinstance(module, nn.Linear) and name not in modules_to_not_convert:
+            # Check if the current key is not in the `modules_to_not_convert`
+            if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):
+                with init_empty_weights():
+                    in_features = module.in_features
+                    out_features = module.out_features
+
+                    model._modules[name] = WQLinear_GEMM(
+                        w_bit=quantization_config.w_bit,
+                        group_size=quantization_config.q_group_size,
+                        in_features=in_features,
+                        out_features=out_features,
+                        bias=module.bias is not None,
+                        dev=module.weight.device,
+                    )
+                    has_been_replaced = True
+
+                    # Store the module class in case we need to transpose the weight later
+                    model._modules[name].source_cls = type(module)
+                    # Force requires grad to False to avoid unexpected errors
+                    model._modules[name].requires_grad_(False)
+        if len(list(module.children())) > 0:
+            _, has_been_replaced = replace_with_awq_linear(
+                module,
+                modules_to_not_convert,
+                current_key_name,
+                quantization_config,
+                has_been_replaced=has_been_replaced,
+            )
+        # Remove the last key for recursion
+        current_key_name.pop(-1)
+    return model, has_been_replaced

src/transformers/modeling_utils.py

@@ -69,6 +69,7 @@
     extract_commit_hash,
     has_file,
     is_accelerate_available,
+    is_auto_awq_available,
     is_auto_gptq_available,
     is_bitsandbytes_available,
     is_flash_attn_2_available,
@@ -89,7 +90,7 @@
     is_torch_fx_proxy,
     is_torchdynamo_compiling,
 )
-from .utils.quantization_config import BitsAndBytesConfig, GPTQConfig, QuantizationMethod
+from .utils.quantization_config import AWQConfig, BitsAndBytesConfig, GPTQConfig, QuantizationMethod
 from .utils.versions import require_version_core
 
 
@@ -2787,6 +2788,12 @@ def from_pretrained(
                 logger.info("We suggest you to set `torch_dtype=torch.float16` for better efficiency with GPTQ.")
 
             quantizer = GPTQQuantizer.from_dict(quantization_config.to_dict())
+        elif quantization_method_from_config == QuantizationMethod.AWQ:
+            if not torch.cuda.is_available():
+                raise RuntimeError("GPU is required to run AWQ quantized model.")
+
+            if not is_auto_awq_available():
+                raise ImportError("Loading an AWQ quantized model requires auto-awq library (`pip install autoawq`)")
 
         if (
             is_8bit_serializable
@@ -3224,6 +3231,15 @@ def from_pretrained(
         if quantization_method_from_config == QuantizationMethod.GPTQ:
             model = quantizer.convert_model(model)
             model._is_quantized_training_enabled = True
+        elif quantization_method_from_config == QuantizationMethod.AWQ:
+            from .integrations import replace_with_awq_linear
+
+            if quantization_config is None:
+                quantization_config = AWQConfig.from_dict(config.quantization_config)
+
+            model, _ = replace_with_awq_linear(
+                model, quantization_config=quantization_config, modules_to_not_convert=["lm_head"]
+            )
 
         if quantization_method_from_config is not None:
             model.quantization_method = quantization_method_from_config

src/transformers/utils/__init__.py

@@ -104,6 +104,7 @@
     get_torch_version,
     is_accelerate_available,
     is_apex_available,
+    is_auto_awq_available,
     is_auto_gptq_available,
     is_bitsandbytes_available,
     is_bs4_available,

src/transformers/utils/import_utils.py

@@ -107,6 +107,8 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[
 _openai_available = _is_package_available("openai")
 _optimum_available = _is_package_available("optimum")
 _auto_gptq_available = _is_package_available("auto_gptq")
+# `importlib.metadata.version` doesn't work with `awq`
+_auto_awq_available = importlib.util.find_spec("awq") is not None
 _pandas_available = _is_package_available("pandas")
 _peft_available = _is_package_available("peft")
 _phonemizer_available = _is_package_available("phonemizer")
@@ -675,6 +677,10 @@ def is_optimum_available():
     return _optimum_available
 
 
+def is_auto_awq_available():
+    return _auto_awq_available
+
+
 def is_auto_gptq_available():
     return _auto_gptq_available
 

src/transformers/utils/quantization_config.py

@@ -37,6 +37,7 @@
 class QuantizationMethod(str, Enum):
     BITS_AND_BYTES = "bitsandbytes"
     GPTQ = "gptq"
+    AWQ = "awq"
 
 
 @dataclass
@@ -418,3 +419,36 @@ def post_init(self):
                     f"""dataset needs to be either a list of string or a value in
                     ['wikitext2','c4','c4-new','ptb','ptb-new'], but we found {self.dataset}"""
                 )
+
+
+@dataclass
+class AWQConfig(QuantizationConfigMixin):
+    """
+    This is a wrapper class about all possible attributes and features that you can play with a model that has been
+    loaded using `auto-awq` library awq quantization relying on auto_awq backend.
+
+    Args:
+        w_bit (`int`):
+            The number of bits to quantize to.
+        zero_point (`bool`, *optional*, defaults to `True`):
+            Whether to use zero point quantization.
+        q_group_size (`int`, *optional*, defaults to 128):
+            The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
+        version (`str`, *optional*, defaults to `GEMM`):
+            The version of the quantization algorithm to use.
+    """
+
+    def __init__(
+        self,
+        w_bit: int,
+        q_group_size: int = 128,
+        zero_point: bool = True,
+        version: str = "GEMM",
+        **kwargs,
+    ):
+        self.quant_method = QuantizationMethod.AWQ
+
+        self.w_bit = w_bit
+        self.q_group_size = q_group_size
+        self.zero_point = zero_point
+        self.version = version