Aluode/PerceptionLabPortable

1# Copyright 2024 The HuggingFace Inc. team. All rights reserved.2# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15import warnings16from typing import Optional, Union17 18from ..models.auto.configuration_auto import AutoConfig19from ..utils import logging20from ..utils.quantization_config import (21    AqlmConfig,22    AutoRoundConfig,23    AwqConfig,24    BitNetQuantConfig,25    BitsAndBytesConfig,26    CompressedTensorsConfig,27    EetqConfig,28    FbgemmFp8Config,29    FineGrainedFP8Config,30    FPQuantConfig,31    GPTQConfig,32    HiggsConfig,33    HqqConfig,34    Mxfp4Config,35    QuantizationConfigMixin,36    QuantizationMethod,37    QuantoConfig,38    QuarkConfig,39    SpQRConfig,40    TorchAoConfig,41    VptqConfig,42)43from .base import HfQuantizer44from .quantizer_aqlm import AqlmHfQuantizer45from .quantizer_auto_round import AutoRoundQuantizer46from .quantizer_awq import AwqQuantizer47from .quantizer_bitnet import BitNetHfQuantizer48from .quantizer_bnb_4bit import Bnb4BitHfQuantizer49from .quantizer_bnb_8bit import Bnb8BitHfQuantizer50from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer51from .quantizer_eetq import EetqHfQuantizer52from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer53from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer54from .quantizer_fp_quant import FPQuantHfQuantizer55from .quantizer_gptq import GptqHfQuantizer56from .quantizer_higgs import HiggsHfQuantizer57from .quantizer_hqq import HqqHfQuantizer58from .quantizer_mxfp4 import Mxfp4HfQuantizer59from .quantizer_quanto import QuantoHfQuantizer60from .quantizer_quark import QuarkHfQuantizer61from .quantizer_spqr import SpQRHfQuantizer62from .quantizer_torchao import TorchAoHfQuantizer63from .quantizer_vptq import VptqHfQuantizer64 65 66AUTO_QUANTIZER_MAPPING = {67    "awq": AwqQuantizer,68    "bitsandbytes_4bit": Bnb4BitHfQuantizer,69    "bitsandbytes_8bit": Bnb8BitHfQuantizer,70    "gptq": GptqHfQuantizer,71    "aqlm": AqlmHfQuantizer,72    "quanto": QuantoHfQuantizer,73    "quark": QuarkHfQuantizer,74    "fp_quant": FPQuantHfQuantizer,75    "eetq": EetqHfQuantizer,76    "higgs": HiggsHfQuantizer,77    "hqq": HqqHfQuantizer,78    "compressed-tensors": CompressedTensorsHfQuantizer,79    "fbgemm_fp8": FbgemmFp8HfQuantizer,80    "torchao": TorchAoHfQuantizer,81    "bitnet": BitNetHfQuantizer,82    "vptq": VptqHfQuantizer,83    "spqr": SpQRHfQuantizer,84    "fp8": FineGrainedFP8HfQuantizer,85    "auto-round": AutoRoundQuantizer,86    "mxfp4": Mxfp4HfQuantizer,87}88 89AUTO_QUANTIZATION_CONFIG_MAPPING = {90    "awq": AwqConfig,91    "bitsandbytes_4bit": BitsAndBytesConfig,92    "bitsandbytes_8bit": BitsAndBytesConfig,93    "eetq": EetqConfig,94    "gptq": GPTQConfig,95    "aqlm": AqlmConfig,96    "quanto": QuantoConfig,97    "quark": QuarkConfig,98    "fp_quant": FPQuantConfig,99    "hqq": HqqConfig,100    "compressed-tensors": CompressedTensorsConfig,101    "fbgemm_fp8": FbgemmFp8Config,102    "higgs": HiggsConfig,103    "torchao": TorchAoConfig,104    "bitnet": BitNetQuantConfig,105    "vptq": VptqConfig,106    "spqr": SpQRConfig,107    "fp8": FineGrainedFP8Config,108    "auto-round": AutoRoundConfig,109    "mxfp4": Mxfp4Config,110}111 112logger = logging.get_logger(__name__)113 114 115class AutoQuantizationConfig:116    """117    The Auto-HF quantization config class that takes care of automatically dispatching to the correct118    quantization config given a quantization config stored in a dictionary.119    """120 121    @classmethod122    def from_dict(cls, quantization_config_dict: dict):123        quant_method = quantization_config_dict.get("quant_method")124        # We need a special care for bnb models to make sure everything is BC ..125        if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):126            suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"127            quant_method = QuantizationMethod.BITS_AND_BYTES + suffix128        elif quant_method is None:129            raise ValueError(130                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"131            )132 133        if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING:134            raise ValueError(135                f"Unknown quantization type, got {quant_method} - supported types are:"136                f" {list(AUTO_QUANTIZER_MAPPING.keys())}"137            )138 139        target_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_method]140        return target_cls.from_dict(quantization_config_dict)141 142    @classmethod143    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):144        model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)145        if getattr(model_config, "quantization_config", None) is None:146            raise ValueError(147                f"Did not found a `quantization_config` in {pretrained_model_name_or_path}. Make sure that the model is correctly quantized."148            )149        quantization_config_dict = model_config.quantization_config150        quantization_config = cls.from_dict(quantization_config_dict)151        # Update with potential kwargs that are passed through from_pretrained.152        quantization_config.update(**kwargs)153        return quantization_config154 155 156class AutoHfQuantizer:157    """158     The Auto-HF quantizer class that takes care of automatically instantiating to the correct159    `HfQuantizer` given the `QuantizationConfig`.160    """161 162    @classmethod163    def from_config(cls, quantization_config: Union[QuantizationConfigMixin, dict], **kwargs):164        # Convert it to a QuantizationConfig if the q_config is a dict165        if isinstance(quantization_config, dict):166            quantization_config = AutoQuantizationConfig.from_dict(quantization_config)167 168        quant_method = quantization_config.quant_method169 170        # Again, we need a special care for bnb as we have a single quantization config171        # class for both 4-bit and 8-bit quantization172        if quant_method == QuantizationMethod.BITS_AND_BYTES:173            if quantization_config.load_in_8bit:174                quant_method += "_8bit"175            else:176                quant_method += "_4bit"177 178        if quant_method not in AUTO_QUANTIZER_MAPPING:179            raise ValueError(180                f"Unknown quantization type, got {quant_method} - supported types are:"181                f" {list(AUTO_QUANTIZER_MAPPING.keys())}"182            )183 184        target_cls = AUTO_QUANTIZER_MAPPING[quant_method]185        return target_cls(quantization_config, **kwargs)186 187    @classmethod188    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):189        quantization_config = AutoQuantizationConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)190        return cls.from_config(quantization_config)191 192    @classmethod193    def merge_quantization_configs(194        cls,195        quantization_config: Union[dict, QuantizationConfigMixin],196        quantization_config_from_args: Optional[QuantizationConfigMixin],197    ):198        """199        handles situations where both quantization_config from args and quantization_config from model config are present.200        """201        if quantization_config_from_args is not None:202            warning_msg = (203                "You passed `quantization_config` or equivalent parameters to `from_pretrained` but the model you're loading"204                " already has a `quantization_config` attribute. The `quantization_config` from the model will be used."205            )206        else:207            warning_msg = ""208 209        if isinstance(quantization_config, dict):210            # Convert the config based on the type of quantization_config_from_args (e.g., AutoRoundConfig), which takes priority before automatic configuration dispatch.211            if isinstance(quantization_config_from_args, AutoRoundConfig):212                quantization_config = AutoRoundConfig.from_dict(quantization_config)213            else:214                quantization_config = AutoQuantizationConfig.from_dict(quantization_config)215 216        if (217            quantization_config_from_args is not None218            and quantization_config.__class__.__name__ != quantization_config_from_args.__class__.__name__219        ):220            raise ValueError(221                f"The model is quantized with {quantization_config.__class__.__name__} but you are passing a {quantization_config_from_args.__class__.__name__} config. "222                "Please make sure to pass the same quantization config class to `from_pretrained` with different loading attributes."223            )224 225        if (226            isinstance(227                quantization_config,228                (GPTQConfig, AwqConfig, AutoRoundConfig, FbgemmFp8Config, CompressedTensorsConfig, Mxfp4Config),229            )230            and quantization_config_from_args is not None231        ):232            loading_attr_dict = quantization_config_from_args.get_loading_attributes()233            for attr, val in loading_attr_dict.items():234                setattr(quantization_config, attr, val)235 236            warning_msg += f"However, loading attributes (e.g. {list(loading_attr_dict.keys())}) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored."237 238        if warning_msg != "" and not isinstance(quantization_config, Mxfp4Config):239            warnings.warn(warning_msg)240        else:241            # in the case of mxfp4, we don't want to print the warning message, bit confusing for users242            logger.info(warning_msg)243        return quantization_config244 245    @staticmethod246    def supports_quant_method(quantization_config_dict):247        quant_method = quantization_config_dict.get("quant_method", None)248        if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):249            suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"250            quant_method = QuantizationMethod.BITS_AND_BYTES + suffix251        elif quant_method is None:252            raise ValueError(253                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"254            )255 256        if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING:257            logger.warning(258                f"Unknown quantization type, got {quant_method} - supported types are:"259                f" {list(AUTO_QUANTIZER_MAPPING.keys())}. Hence, we will skip the quantization. "260                "To remove the warning, you can delete the quantization_config attribute in config.json"261            )262            return False263        return True264 265 266def register_quantization_config(method: str):267    """Register a custom quantization configuration."""268 269    def register_config_fn(cls):270        if method in AUTO_QUANTIZATION_CONFIG_MAPPING:271            raise ValueError(f"Config '{method}' already registered")272 273        if not issubclass(cls, QuantizationConfigMixin):274            raise TypeError("Config must extend QuantizationConfigMixin")275 276        AUTO_QUANTIZATION_CONFIG_MAPPING[method] = cls277        return cls278 279    return register_config_fn280 281 282def register_quantizer(name: str):283    """Register a custom quantizer."""284 285    def register_quantizer_fn(cls):286        if name in AUTO_QUANTIZER_MAPPING:287            raise ValueError(f"Quantizer '{name}' already registered")288 289        if not issubclass(cls, HfQuantizer):290            raise ValueError("Quantizer must extend HfQuantizer")291 292        AUTO_QUANTIZER_MAPPING[name] = cls293        return cls294 295    return register_quantizer_fn296 297 298def get_hf_quantizer(config, quantization_config, dtype, from_tf, from_flax, device_map, weights_only, user_agent):299    pre_quantized = hasattr(config, "quantization_config")300    if pre_quantized and not AutoHfQuantizer.supports_quant_method(config.quantization_config):301        pre_quantized = False302 303    if pre_quantized or quantization_config is not None:304        if pre_quantized:305            config.quantization_config = AutoHfQuantizer.merge_quantization_configs(306                config.quantization_config, quantization_config307            )308        else:309            config.quantization_config = quantization_config310 311        hf_quantizer = AutoHfQuantizer.from_config(312            config.quantization_config,313            pre_quantized=pre_quantized,314        )315    else:316        hf_quantizer = None317 318    if hf_quantizer is not None:319        hf_quantizer.validate_environment(320            dtype=dtype,321            from_tf=from_tf,322            from_flax=from_flax,323            device_map=device_map,324            weights_only=weights_only,325        )326        dtype = hf_quantizer.update_dtype(dtype)327        device_map = hf_quantizer.update_device_map(device_map)328        config = hf_quantizer.update_tp_plan(config)329        config = hf_quantizer.update_ep_plan(config)330 331        # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry`332        if not getattr(hf_quantizer.quantization_config, "dequantize", False):333            quant_method = hf_quantizer.quantization_config.quant_method334            user_agent["quant"] = getattr(quant_method, "value", quant_method)335    return hf_quantizer, config, dtype, device_map336