transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc1__py3-none-any.whl

transformers/integrations/awq.py

@@ -13,21 +13,14 @@
 # limitations under the License.
 "AWQ (Activation aware Weight Quantization) integration file"
 
-import importlib
+from typing import Optional, Union
 
-from packaging import version
+from ..quantizers.quantizers_utils import should_convert_module
+from ..utils import is_accelerate_available, is_torch_available, logging
 
-from ..activations import ACT2FN
-from ..modeling_rope_utils import ROPE_INIT_FUNCTIONS
-from ..modeling_utils import PreTrainedModel
-from ..utils import is_auto_awq_available, is_ipex_available, is_torch_available, logging
-from ..utils.quantization_config import (
-    AwqBackendPackingMethod,
-    AwqConfig,
-    AWQLinearVersion,
-    ExllamaVersion,
-)
 
+if is_accelerate_available():
+    from accelerate import init_empty_weights
 
 if is_torch_available():
     import torch
@@ -35,44 +28,6 @@ if is_torch_available():
 
 logger = logging.get_logger(__name__)
 
-AWQ_FUSED_MAPPINGS = {
-    "mistral": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-        "mlp": ["gate_proj", "up_proj", "down_proj"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-    "mixtral": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-        "mlp": ["w1", "w3", "w2"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-    "llama": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-        "mlp": ["gate_proj", "up_proj", "down_proj"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-    "llava": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-        "mlp": ["gate_proj", "up_proj", "down_proj"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-    "qwen2": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-        "mlp": ["gate_proj", "up_proj", "down_proj"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-    "qwen3": {
-        "attention": ["q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm"],
-        "mlp": ["gate_proj", "up_proj", "down_proj"],
-        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-        "use_alibi": False,
-    },
-}
 
 AWQ_SCALES_MAPPINGS = {
     "starcoder2": {"act": "act", "layer_before_act": "c_fc"},
@@ -86,55 +41,8 @@ AWQ_SCALES_MAPPINGS = {
 }
 
 
-if is_auto_awq_available():
-    from awq.modules.fused.attn import RoPE
-
-    class AWQRoPE(RoPE):
-        """
-        AWQRoPE module for hacking rope implementation in AWQ fused attention modules to support more models.
-
-        Args:
-            rope_type (`str`):
-                The rope type to use.
-            head_dim (`int`):
-                The head dimension.
-            max_seq_len (`int`):
-                The maximum sequence length.
-            config (`PreTrainedConfig`):
-                The model config object.
-            device (`torch.device`):
-                The device to put the module on.
-        """
-
-        def __init__(self, rope_type, head_dim, max_seq_len, config, device):
-            rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type]
-            self.inv_freq, self.attention_scaling = rope_init_fn(config, device)
-            # Use fake rope_theta to initialize the parent class
-            super().__init__(head_dim=head_dim, max_seq_len=max_seq_len, device=device, rope_theta=-1)
-
-        def precompute_freqs_cis(self, dim: int, end: int, theta=-1):
-            t = torch.arange(end, device=self.inv_freq.device)
-            freqs = torch.outer(t, self.inv_freq).float()
-            freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
-            del self.inv_freq  # free the memory
-            return freqs_cis
-
-        def forward(
-            self,
-            xq: torch.Tensor,
-            xk: torch.Tensor,
-            start_pos: int,
-            seqlen: int,
-            partial: bool = False,
-        ):
-            xq_out, xk_out = super().forward(xq, xk, start_pos, seqlen, partial)
-            xq_out = (xq_out * self.attention_scaling).type_as(xq)
-            xk_out = (xk_out * self.attention_scaling).type_as(xk)
-            return xq_out, xk_out
-
-
 def replace_quantization_scales(model, model_type):
-    from awq.modules.act import ScaledActivation
+    from gptqmodel.quantization.awq.modules.act import ScaledActivation
 
     if model_type not in AWQ_SCALES_MAPPINGS:
         return model
@@ -154,437 +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.
-    """
-    if modules_to_not_convert is None:
-        modules_to_not_convert = []
-
-    backend = quantization_config.backend
-
-    if not is_auto_awq_available():
-        raise ValueError(
-            "AWQ (either `autoawq` or `llmawq`) is not available. Please install it with `pip install autoawq` or check out the installation guide in https://github.com/mit-han-lab/llm-awq"
-        )
-
-    if backend == AwqBackendPackingMethod.AUTOAWQ:
-        if quantization_config.version == AWQLinearVersion.GEMM:
-            from awq.modules.linear.gemm import WQLinear_GEMM
-
-            target_cls = WQLinear_GEMM
-        elif quantization_config.version == AWQLinearVersion.GEMV:
-            from awq.modules.linear.gemv import WQLinear_GEMV
-
-            target_cls = WQLinear_GEMV
-        elif quantization_config.version == AWQLinearVersion.EXLLAMA:
-            if quantization_config.exllama_config["version"] == ExllamaVersion.ONE:
-                from awq.modules.linear.exllama import WQLinear_Exllama
-
-                target_cls = WQLinear_Exllama
-            elif quantization_config.exllama_config["version"] == ExllamaVersion.TWO:
-                from awq.modules.linear.exllamav2 import WQLinear_ExllamaV2
-
-                target_cls = WQLinear_ExllamaV2
-            else:
-                raise ValueError(f"Unrecognized Exllama version: {quantization_config.exllama_config['version']}")
-        elif quantization_config.version == AWQLinearVersion.IPEX:
-            from awq.modules.linear.gemm_ipex import WQLinear_IPEX
-
-            target_cls = WQLinear_IPEX
-        else:
-            raise ValueError(f"Unrecognized AWQ version: {quantization_config.version}")
-    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
-
-                model._modules[name] = target_cls(
-                    w_bit=quantization_config.bits,
+        device_map (`Union[str, dict]`, *optional*, defaults to `None`):
+            The device map that maps the parameters to the 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=in_features,
-                    out_features=out_features,
+                    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,
-            )
-        # Remove the last key for recursion
-        current_key_name.pop(-1)
-    return model, has_been_replaced
-
-
-def get_modules_to_fuse(model, quantization_config):
-    """
-    Returns the fusing mapping given the quantization config and the model
-
-    Args:
-        model (`~PreTrainedModel`):
-            The model to fuse - note this model should have been converted into AWQ format beforehand.
-        quantization_config (`~transformers.quantization_config.AWQConfig`):
-            The quantization configuration to use.
-    """
-    if not isinstance(model, PreTrainedModel):
-        raise TypeError(f"The model should be an instance of `PreTrainedModel`, got {model.__class__.__name__}")
-
-    # Always default to `quantization_config.modules_to_fuse`
-    if quantization_config.modules_to_fuse is not None:
-        current_fused_mapping = quantization_config.modules_to_fuse
-        current_fused_mapping["max_seq_len"] = quantization_config.fuse_max_seq_len
-    elif model.config.model_type in AWQ_FUSED_MAPPINGS:
-        current_fused_mapping = AWQ_FUSED_MAPPINGS[model.config.model_type]
-
-        # Properly deal with the case where we have a multi-modal model as well (e.g. Llava)
-        config = model.config.get_text_config(decoder=True)
-
-        # Handle hidden_size, num_attention_heads, num_key_value_heads, rope_parameters on our own.
-        hidden_size = config.hidden_size
-        num_attention_heads = config.num_attention_heads
-        num_key_value_heads = getattr(config, "num_key_value_heads", num_attention_heads)
-        rope_parameters = config.rope_parameters
-
-        # Fill `current_fused_mapping` with the expected values
-        current_fused_mapping["hidden_size"] = hidden_size
-        current_fused_mapping["num_attention_heads"] = num_attention_heads
-        current_fused_mapping["num_key_value_heads"] = num_key_value_heads
-        current_fused_mapping["rope_parameters"] = rope_parameters
-        current_fused_mapping["max_seq_len"] = quantization_config.fuse_max_seq_len
-    else:
-        raise ValueError(
-            "Fusing mapping not found either on the quantization config or the supported `AWQ_FUSED_MAPPINGS`. Please pass a `fused_mapping` argument"
-            " in the `quantization_config` or raise an issue on transformers https://github.com/huggingface/transformers to add its support."
-        )
-    return current_fused_mapping
-
-
-def fuse_awq_modules(model, quantization_config):
-    """
-    Optionally fuse some modules in the model to speedup inference.
-
-    Args:
-        model (`~PreTrainedModel`):
-            The model to fuse - note this model should have been converted into AWQ format beforehand.
-        quantization_config (`Union[AwqConfig, dict]`):
-            The quantization configuration to use.
-    """
-    # We need to convert it from dict in order to get an AwqConfig object
-    # otherwise the fields `backend` etc. will not be available
-    # https://github.com/huggingface/transformers/pull/27411#discussion_r1414044495
-    if isinstance(quantization_config, dict):
-        quantization_config = AwqConfig.from_dict(quantization_config)
-    backend = quantization_config.backend
-
-    modules_to_fuse = get_modules_to_fuse(model, quantization_config)
-    modules_to_not_convert = getattr(quantization_config, "modules_to_not_convert", None)
-
-    if backend == AwqBackendPackingMethod.AUTOAWQ:
-        from awq.modules.fused.attn import QuantAttentionFused
-        from awq.modules.fused.mlp import QuantFusedMLP
-        from awq.modules.fused.norm import FasterTransformerRMSNorm
-
-        # Hack QuantAttentionFused to modify the return value of forward function to avoid returning past_key_value
-        old_quant_attention_fused_forward = QuantAttentionFused.forward
-
-        def new_quant_attention_fused_forward(self, *args, **kwargs):
-            attn_output, attention_weight, _ = old_quant_attention_fused_forward(self, *args, **kwargs)
-            return attn_output, attention_weight
-
-        QuantAttentionFused.forward = new_quant_attention_fused_forward
-    else:
-        raise ValueError("Fusing is only supported for the AutoAWQ backend")
-
-    fused_attention_modules = []
-
-    for name, module in model.named_modules():
-        if modules_to_not_convert is not None:
-            if any(module_name_to_not_convert in name for module_name_to_not_convert in modules_to_not_convert):
-                continue
-
-        # Replace layer norms
-        _fuse_awq_layernorm(modules_to_fuse["layernorm"], module, FasterTransformerRMSNorm)
-
-        # Replace MLP layers if awq version is not ipex.
-        if quantization_config.version != "ipex":
-            _fuse_awq_mlp(model, name, modules_to_fuse["mlp"], module, QuantFusedMLP)
-        else:
-            logger.info("The IPEX version AWQ does not support fuse mlp for now.")
-
-        # Replace attention layers
-        attention_has_been_fused = _fuse_awq_attention_layers(
-            model, module, modules_to_fuse, name, QuantAttentionFused
-        )
-
-        if attention_has_been_fused:
-            fused_attention_modules.append(name.split(".")[0])
-
-    # For AWQ fused + Llama we need to set `config._attn_implementation` = "custom" to avoid unexpected behavior and pass
-    # `None` attention mask to the fused attention modules as now the attention mask is dropped by our models and dealt
-    # by the `AttentionMaskConverter` module.
-    if len(fused_attention_modules) > 0:
-        for module_name, module in model.named_modules():
-            if any(
-                module_name in fused_attention_modules for fused_attention_parent_module in fused_attention_modules
-            ):
-                if hasattr(module, "config") and hasattr(module.config, "_attn_implementation"):
-                    module.config._attn_implementation = "custom"
-    return model
-
-
-def _fuse_awq_layernorm(fuse_module_names, module, target_cls):
-    """
-    Fuse the LayerNorm layers into a target class using autoawq
-
-    Args:
-        fuse_module_names (`list[str]`):
-            The list of module names to fuse
-        module (`nn.Module`):
-            The pytorch parent module that has layernorm modules to fuse
-        target_cls (`~autoawq.FasterTransformerRMSNorm`):
-            The `FasterTransformerRMSNorm` class as it only supports that class
-            for now.
-    """
-    for module_name in fuse_module_names:
-        if hasattr(module, module_name):
-            old_module = getattr(module, module_name)
-            module._modules[module_name] = target_cls(
-                old_module.weight,
-                old_module.variance_epsilon,
-            ).to(old_module.weight.device)
-            del old_module
-
-
-def _fuse_awq_mlp(model, current_module_name, fuse_module_names, module, target_cls):
-    """
-    Fuse the MLP layers into a target class using autoawq
-
-    Args:
-        model (`~PreTrainedModel`):
-            The input pretrained model
-        current_module_name (`str`):
-            The current submodule name
-        fuse_module_names (`list[str]`):
-            The list of module names to fuse. For the MLP layers it has to be an array
-            of length 3 that consists of the 3 MLP layers in the order (gate (dense layer post-attention) / up / down layers)
-        module (`nn.Module`):
-            The pytorch parent module that has layernorm modules to fuse
-        target_cls (`~autoawq.QuantFusedMLP`):
-            The `QuantFusedMLP` class as it only supports that class
-            for now.
-    """
-    if len(fuse_module_names) == 0:
-        return
-
-    if hasattr(module, fuse_module_names[0]):
-        gate_proj = getattr(module, fuse_module_names[0])
-        up_proj = getattr(module, fuse_module_names[1])
-        down_proj = getattr(module, fuse_module_names[2])
-
-        previous_device = gate_proj.qweight.device
-
-        # Deal also with the case model has `text_config` attribute
-        config = model.config.get_text_config(decoder=True)
-        hidden_act = config.hidden_act
-        activation_fn = ACT2FN[hidden_act]
-        new_module = target_cls(gate_proj, down_proj, up_proj, activation_fn)
-
-        parent_name, child_name = current_module_name.rsplit(".", 1)
-        parent = model.get_submodule(parent_name)
-        setattr(parent, child_name, new_module.to(previous_device))
-
-        del gate_proj, up_proj, down_proj
-
-
-def _fuse_awq_attention_layers(model, module, modules_to_fuse, current_module_name, target_cls):
-    """
-    Fuse the Attention layers into a target class using autoawq
-
-    Args:
-        model (`~PreTrainedModel`):
-            The input pretrained model
-        module (`nn.Module`):
-            The pytorch parent module that has layernorm modules to fuse
-        modules_to_fuse (`list[str]`):
-            The module fusing mapping. The dictionary has to contain a field `attention` with attention module names
-            in the correct order: q, k, v, o layer, (q_norm, k_norm) optional
-        current_module_name (`str`):
-            The current submodule name
-        target_cls (`~autoawq.QuantAttentionFused`):
-            The `QuantAttentionFused` class as it only supports that class
-            for now.
-    """
-    from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
-
-    module_has_been_fused = False
-
-    if len(modules_to_fuse["attention"]) == 0:
-        return module_has_been_fused
-
-    if hasattr(module, modules_to_fuse["attention"][0]):
-        # First, we pack the QKV layers together
-        q_proj = getattr(module, modules_to_fuse["attention"][0])
-
-        if isinstance(q_proj, WQLinear_GEMV):
-            linear_target_cls = WQLinear_GEMV
-            cat_dim = 0
-        elif isinstance(q_proj, WQLinear_GEMM):
-            linear_target_cls = WQLinear_GEMM
-            cat_dim = 1
-        elif is_ipex_available() and version.parse(importlib.metadata.version("autoawq")) > version.parse("0.2.6"):
-            from awq.modules.linear import WQLinear_IPEX
-
-            if isinstance(q_proj, WQLinear_IPEX):
-                linear_target_cls = WQLinear_IPEX
-                cat_dim = 1
-        else:
-            raise ValueError("Unsupported q_proj type: {type(q_proj)}")
-
-        previous_device = q_proj.qweight.device
-
-        k_proj = getattr(module, modules_to_fuse["attention"][1])
-        v_proj = getattr(module, modules_to_fuse["attention"][2])
-        o_proj = getattr(module, modules_to_fuse["attention"][3])
-
-        # maybe there are q_norm and k_norm layers
-        if len(modules_to_fuse["attention"]) > 4:
-            q_norm = getattr(module, modules_to_fuse["attention"][4])
-            k_norm = getattr(module, modules_to_fuse["attention"][5])
-        else:
-            q_norm = None
-            k_norm = None
-
-        bias = torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0) if q_proj.bias is not None else None
-
-        qkv_layer = linear_target_cls(
-            q_proj.w_bit,
-            q_proj.group_size,
-            q_proj.in_features,
-            q_proj.out_features + k_proj.out_features + v_proj.out_features,
-            q_proj.bias is not None,
-            next(iter(module.state_dict().values())).device,
-        )
-
-        qkv_layer.qweight = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=cat_dim)
-        qkv_layer.qzeros = torch.cat([q_proj.qzeros, k_proj.qzeros, v_proj.qzeros], dim=cat_dim)
-        qkv_layer.scales = torch.cat([q_proj.scales, k_proj.scales, v_proj.scales], dim=cat_dim)
-
-        if isinstance(qkv_layer, WQLinear_GEMV):
-            qkv_layer.split_k_iters = q_proj.split_k_iters
-
-        qkv_layer.bias = bias
-
-        fused_attention_layer = target_cls(
-            modules_to_fuse["hidden_size"],
-            modules_to_fuse["num_attention_heads"],
-            modules_to_fuse["num_key_value_heads"],
-            qkv_layer,
-            o_proj,
-            previous_device,
-            modules_to_fuse["max_seq_len"],
-            use_alibi=modules_to_fuse["use_alibi"],
-            # The default value in autoawq is set to 10000.0
-            rope_theta=modules_to_fuse["rope_parameters"].get("rope_theta", 10000.0),
-            q_norm=q_norm,
-            k_norm=k_norm,
-        )
-
-        # Hack the rope module if not using alibi and rope_type is not default
-        # As the default rope implementation in autoawq only supports the "default" rope type
-        rope_type = modules_to_fuse["rope_parameters"].get("rope_type", "default")
-        if not modules_to_fuse["use_alibi"] and rope_type != "default":
-            fused_attention_layer.rope = AWQRoPE(
-                rope_type,
-                modules_to_fuse["hidden_size"] // modules_to_fuse["num_attention_heads"],
-                modules_to_fuse["max_seq_len"],
-                model.config.get_text_config(decoder=True),
-                previous_device,
-            )
-
-        fused_attention_layer.is_hf_transformers = True
-
-        parent_name, child_name = current_module_name.rsplit(".", 1)
-        parent = model.get_submodule(parent_name)
-        setattr(parent, child_name, fused_attention_layer.to(previous_device))
-
-        del q_proj, k_proj, v_proj, o_proj, q_norm, k_norm
-        module_has_been_fused = True
-
-    return module_has_been_fused
-
-
-def post_init_awq_exllama_modules(model, exllama_config):
-    """
-    Runs post init for Exllama layers which performs:
-        - Weights unpacking, reordering and repacking
-        - Devices scratch space allocation
-    """
-
-    if exllama_config["version"] == ExllamaVersion.ONE:
-        from awq.modules.linear.exllama import exllama_post_init
-
-        model = exllama_post_init(model)
-    elif exllama_config["version"] == ExllamaVersion.TWO:
-        from awq.modules.linear.exllamav2 import exllamav2_post_init
-
-        model = exllamav2_post_init(
-            model,
-            max_input_len=exllama_config["max_input_len"],
-            max_batch_size=exllama_config["max_batch_size"],
+    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."
         )
-    else:
-        raise ValueError(f"Unrecognized Exllama version: {exllama_config['version']}")
-
-    return model
-
-
-def post_init_awq_ipex_modules(model):
-    """
-    Runs post init for IPEX layers which performs:
-        - Weights packing, reordering and repacking
-    """
-
-    from awq.modules.linear.gemm_ipex import ipex_post_init
-
-    model = ipex_post_init(model)
 
     return model