2023 the huggingface team 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 orglicenseslicense2 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 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 successfull or not during the module replacement we also infer the backend to use through the quantizationconfig object args model torch nn module the model to convert can be any torch nn module instance quantizationconfig awqconfig the quantization config object that contains the quantization parameters modulestonotconvert list optional a list of modules to not convert if a module name is in the list e g lmhead it will not be converted currentkeyname list optional a list that contains the current key name this is used for recursion and should not be passed by the user hasbeenreplaced 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 check if the current key is not in the modulestonotconvert force requires grad to false to avoid unexpected errors remove the last key for recursion 2023 the huggingface team 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 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 successfull or not during the module replacement we also infer the backend to use through the quantization_config object 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 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 check if the current key is not in the modules_to_not_convert force requires grad to false to avoid unexpected errors remove the last key for recursion
"AWQ (Activation aware Weight Quantization) integration file"
from ..utils import is_auto_awq_available, is_torch_available
from ..utils.quantization_config import AwqBackendPackingMethod, AWQLinearVersion


if is_torch_available():
    import torch.nn as nn


def replace_with_awq_linear(
    model,
    modules_to_not_convert=None,
    quantization_config=None,
    current_key_name=None,
    has_been_replaced=False,
) -> bool:
    
    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:
        from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
    elif backend == AwqBackendPackingMethod.LLMAWQ:
        from awq.quantize.qmodule import WQLinear

    if backend == AwqBackendPackingMethod.AUTOAWQ:
        target_cls = WQLinear_GEMM if quantization_config.version == AWQLinearVersion.GEMM else WQLinear_GEMV
    else:
        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:
            
            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,
                    group_size=quantization_config.group_size,
                    in_features=in_features,
                    out_features=out_features,
                    bias=module.bias is not None,
                    dev=module.weight.device,
                )
                has_been_replaced = True

                
                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,
            )
        
        current_key_name.pop(-1)
    return model, has_been_replaced