Repository: huggingface/transformers
Pull request: fix awq 

# What does this PR do?

This PR simplifies a bit the AWQ integration. cc @Qubitium 
Also, I fixed some issues with gptq tests as the hub tests are failiing due to them

Changed files:
filename: src/transformers/integrations/awq.py
patch: @@ -15,12 +15,13 @@
 
 from typing import Optional, Union
 
-from ..utils import is_gptqmodel_available, is_llm_awq_available, is_torch_available, logging
-from ..utils.quantization_config import (
-    AwqBackend,
-)
+from ..quantizers.quantizers_utils import should_convert_module
+from ..utils import is_accelerate_available, is_torch_available, logging
 
 
+if is_accelerate_available():
+    from accelerate import init_empty_weights
+
 if is_torch_available():
     import torch
     import torch.nn as nn
@@ -61,120 +62,63 @@ def replace_with_awq_linear(
     model,
     modules_to_not_convert=None,
     quantization_config=None,
-    current_key_name=None,
-    has_been_replaced=False,
     device_map: Optional[Union[str, dict]] = None,
 ) -> bool:
     """
-    Public method that recursively replaces the Linear layers of the given model with AWQ quantized layers.
-    `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the
-    conversion has been successful or not.
-
-    During the module replacement, we also infer the backend to use through the `quantization_config` object.
+    Public method that replaces the linear layers of the given model with awq quantized layers.
 
     Args:
         model (`torch.nn.Module`):
             The model to convert, can be any `torch.nn.Module` instance.
         quantization_config (`AwqConfig`):
             The quantization config object that contains the quantization parameters.
-        modules_to_not_convert (`list`, *optional*):
-            A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be
+        modules_to_not_convert (`list[str]`, *optional*, defaults to `None`):
+            A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be
             converted.
-        current_key_name (`list`, *optional*):
-            A list that contains the current key name. This is used for recursion and should not be passed by the user.
-        has_been_replaced (`bool`, *optional*):
-            A boolean that indicates if the conversion has been successful or not. This is used for recursion and
-            should not be passed by the user.
+        device_map (`Union[str, dict]`, *optional*, defaults to `None`):
+            The device map that maps the parameters to the device
     """
-    if modules_to_not_convert is None:
-        modules_to_not_convert = []
-
-    backend = quantization_config.backend
-
-    if not is_gptqmodel_available() and not is_llm_awq_available():
-        raise ValueError(
-            "AWQ (either `llmawq`) is not available. Please install it with `pip install gptqmodel` or check out the installation guide in https://github.com/mit-han-lab/llm-awq"
-        )
-
-    if backend != AwqBackend.LLMAWQ:
-        from gptqmodel.quantization import METHOD
-        from gptqmodel.utils.importer import hf_select_quant_linear_v2
-
-        target_cls = hf_select_quant_linear_v2(
-            bits=quantization_config.bits,
-            group_size=quantization_config.group_size,
-            desc_act=False,
-            sym=False,
-            format=quantization_config.format,
-            backend=quantization_config.backend,
-            device_map=device_map,
-            quant_method=METHOD.AWQ,
-            zero_point=quantization_config.zero_point,
-            pack=False,
-        )
-    else:
-        from awq.quantize.qmodule import WQLinear
-
-        target_cls = WQLinear
-
-    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):
-                in_features = module.in_features
-                out_features = module.out_features
-
-                if backend != AwqBackend.LLMAWQ:
-                    model._modules[name] = target_cls(
-                        bits=quantization_config.bits,
-                        sym=quantization_config.sym,
-                        desc_act=quantization_config.desc_act,
-                        group_size=quantization_config.group_size,
-                        in_features=in_features,
-                        out_features=out_features,
-                        bias=module.bias is not None,
-                        dev=module.weight.device,
-                        register_buffers=True,
-                    )
-                else:
-                    model._modules[name] = target_cls(
-                        w_bit=quantization_config.bits,
-                        group_size=quantization_config.group_size,
-                        in_features=in_features,
-                        out_features=out_features,
-                        bias=module.bias is not None,
-                        dev=module.weight.device,
-                    )
+    from gptqmodel.quantization import METHOD
+    from gptqmodel.utils.importer import hf_select_quant_linear_v2
+
+    target_cls = hf_select_quant_linear_v2(
+        bits=quantization_config.bits,
+        group_size=quantization_config.group_size,
+        desc_act=False,
+        sym=False,
+        format=quantization_config.format,
+        backend=quantization_config.backend,
+        device_map=device_map,
+        quant_method=METHOD.AWQ,
+        zero_point=quantization_config.zero_point,
+        pack=False,
+    )
+
+    for module_name, module in model.named_modules():
+        if not should_convert_module(module_name, modules_to_not_convert):
+            continue
+        with init_empty_weights():
+            if isinstance(module, nn.Linear):
+                new_module = target_cls(
+                    bits=quantization_config.bits,
+                    sym=quantization_config.sym,
+                    desc_act=quantization_config.desc_act,
+                    group_size=quantization_config.group_size,
+                    in_features=module.in_features,
+                    out_features=module.out_features,
+                    bias=module.bias is not None,
+                    dev=module.weight.device,
+                    register_buffers=True,
+                )
+                new_module.requires_grad_(False)
+                model.set_submodule(module_name, new_module)
                 has_been_replaced = True
 
-                # 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=modules_to_not_convert,
-                current_key_name=current_key_name,
-                quantization_config=quantization_config,
-                has_been_replaced=has_been_replaced,
-                device_map=device_map,
-            )
-        # Remove the last key for recursion
-        current_key_name.pop(-1)
-    return model, has_been_replaced
-
-
-def post_init_awq_ipex_modules(model):
-    """
-    Runs post init for IPEX layers which performs:
-        - Weights packing, reordering and repacking
-    """
-
-    from gptqmodel.quantization.awq.modules.linear.gemm_ipex import ipex_post_init
-
-    model = ipex_post_init(model)
+    if not has_been_replaced:
+        logger.warning(
+            "You are loading your model using eetq but no linear modules were found in your model."
+            " Please double check your model architecture, or submit an issue on github if you think this is"
+            " a bug."
+        )
 
     return model

filename: src/transformers/quantizers/quantizer_awq.py
patch: @@ -43,9 +43,11 @@ class AwqQuantizer(HfQuantizer):
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
 
-    def validate_environment(self, device_map, **kwargs):
+    def validate_environment(self, **kwargs):
         if not is_gptqmodel_available():
-            raise ImportError("Loading an AWQ quantized model requires gptqmodel library (`pip install gptqmodel`)")
+            raise ImportError(
+                "Loading an AWQ quantized model requires gptqmodel. Please install it with `pip install gptqmodel`"
+            )
 
         if not is_accelerate_available():
             raise ImportError("Loading an AWQ quantized model requires accelerate (`pip install accelerate`)")
@@ -72,7 +74,7 @@ def _process_model_before_weight_loading(
             model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules, add_default_skips=True
         )
 
-        model, has_been_replaced = replace_with_awq_linear(
+        model = replace_with_awq_linear(
             model,
             quantization_config=self.quantization_config,
             modules_to_not_convert=self.modules_to_not_convert,
@@ -81,18 +83,12 @@ def _process_model_before_weight_loading(
 
         model = replace_quantization_scales(model, model.config.model_type)
 
-        if not has_been_replaced:
-            logger.warning(
-                "You are loading an AWQ model but no linear modules were found in your model."
-                " Please double check your model architecture, or submit an issue on github if you think this is a bug."
-            )
-
     def _process_model_after_weight_loading(self, model, **kwargs):
         from gptqmodel.utils.model import hf_gptqmodel_post_init
 
         hf_gptqmodel_post_init(model, use_act_order=self.quantization_config.desc_act)
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         if self.quantization_config.backend in [AwqBackend.EXLLAMA_V1, AwqBackend.EXLLAMA_V2]:
             logger.warning("You cannot save an AWQ model that uses Exllama backend!")
             return False

filename: src/transformers/utils/quantization_config.py
patch: @@ -85,7 +85,6 @@ class AwqBackend(str, Enum):
     GEMV_FAST = "gemv_fast"
     TORCH_AWQ = "torch_awq"
     TORCH_FUSED_AWQ = "torch_fused_awq"
-    LLMAWQ = "llm-awq"
 
 
 @dataclass
@@ -806,8 +805,7 @@ class AwqConfig(GPTQConfig):
         zero_point (`bool`, *optional*, defaults to `True`):
             Whether to use zero point quantization.
         backend (`AwqBackend`, *optional*, defaults to `AwqBackend.AUTO`):
-            The quantization backend. Some models might be quantized using `llm-awq` backend. This is useful for users
-            that quantize their own models using `llm-awq` library.
+            The quantization backend.
         modules_to_not_convert (`list`, *optional*, default to `None`):
             The list of modules to not quantize, useful for quantizing models that explicitly require to have
             some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
@@ -852,16 +850,6 @@ def post_init(self):
         if self.backend not in AwqBackend.__members__.values():
             raise ValueError(f"Invalid backend '{self.backend}'. Must be one of: {[b.value for b in AwqBackend]}")
 
-        if self.backend == AwqBackend.LLMAWQ:
-            # Only cuda device can run this function
-            if not (torch.cuda.is_available() or torch.xpu.is_available()):
-                raise ValueError("LLM-AWQ backend is only supported on CUDA and XPU")
-            if torch.cuda.is_available():
-                compute_capability = torch.cuda.get_device_capability()
-                major, minor = compute_capability
-                if major < 8:
-                    raise ValueError("LLM-AWQ backend is only supported on CUDA GPUs with compute capability >= 8.0")
-
     def to_dict(self) -> dict[str, Any]:
         config_dict = super().to_dict()
         config_dict.pop("checkpoint_format")

filename: tests/quantization/autoawq/test_awq.py
patch: @@ -19,7 +19,6 @@
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AwqConfig, OPTForCausalLM
 from transformers.testing_utils import (
     backend_empty_cache,
-    get_device_properties,
     require_accelerate,
     require_gptqmodel,
     require_torch_accelerator,
@@ -58,22 +57,6 @@ def test_wrong_backend(self):
         with self.assertRaises(ValueError):
             AwqConfig(bits=4, backend="unexisting-backend")
 
-        # Only cuda and xpu devices can run this function
-        support_llm_awq = False
-        device_type, major, _ = get_device_properties()
-        if device_type == "cuda" and major >= 8:
-            support_llm_awq = True
-        elif device_type == "xpu":
-            support_llm_awq = True
-
-        if support_llm_awq:
-            # LLMAWQ should work on an A100
-            AwqConfig(bits=4, backend="llm-awq")
-        else:
-            # LLMAWQ does not work on a T4
-            with self.assertRaises(ValueError):
-                AwqConfig(bits=4, backend="llm-awq")
-
     def test_to_dict(self):
         """
         Simple test that checks if one uses a config and converts it to a dict, the dict is the same as the config object

filename: tests/quantization/gptq/test_gptq.py
patch: @@ -113,8 +113,6 @@ class GPTQTest(unittest.TestCase):
     group_size = 128
     desc_act = False
     act_group_aware = True
-    quant_backend = BACKEND.AUTO
-    load_backend = BACKEND.AUTO
     dataset = [
         "gptqmodel is an easy-to-use model quantization library with user-friendly APIs, based on the GPTQ algorithm."
     ]
@@ -143,7 +141,7 @@ def setUpClass(cls):
             desc_act=cls.desc_act,
             act_group_aware=cls.act_group_aware,
             sym=cls.sym,
-            backend=cls.quant_backend,
+            backend=BACKEND.AUTO,
         )
 
         cls.quantized_model = AutoModelForCausalLM.from_pretrained(
@@ -309,7 +307,6 @@ class GPTQTestActOrderExllamaV2(unittest.TestCase):
     # `act_group_aware` == `True` requires `desc_act` == `False` when both are explicitly set
     desc_act = True
     act_group_aware = False
-    load_backend = BACKEND.EXLLAMA_V2
 
     EXPECTED_OUTPUTS = set()
     # flaky test: gptqmodel kernels are not always bitwise deterministic even between transformer/torch versions
@@ -328,7 +325,7 @@ def setUpClass(cls):
             max_input_length=4028,
             desc_act=cls.desc_act,
             act_group_aware=cls.act_group_aware,
-            backend=cls.load_backend,
+            backend=BACKEND.EXLLAMA_V2,
         )
         cls.quantized_model = AutoModelForCausalLM.from_pretrained(
             cls.model_name,
@@ -376,7 +373,6 @@ class GPTQTestExllamaV2(unittest.TestCase):
     https://huggingface.co/docs/transformers/main_classes/quantization#transformers.GPTQConfig
     """
 
-    load_backend = BACKEND.EXLLAMA_V2
     EXPECTED_OUTPUTS = set()
     # flaky test: gptqmodel kernels are not always bitwise deterministic even between transformer/torch versions
     EXPECTED_OUTPUTS.add("Hello, how are you ? I'm doing good, thanks for asking.")
@@ -389,7 +385,7 @@ def setUpClass(cls):
         """
         Setup quantized model
         """
-        cls.quantization_config = GPTQConfig(bits=4, backend=cls.load_backend)
+        cls.quantization_config = GPTQConfig(bits=4, backend=BACKEND.EXLLAMA_V2)
         cls.quantized_model = AutoModelForCausalLM.from_pretrained(
             cls.model_name,
             dtype=torch.float16,