Aluode/PerceptionLabPortable

1# Copyright 2024 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import TYPE_CHECKING, Optional15 16from .base import HfQuantizer17 18 19if TYPE_CHECKING:20    from ..modeling_utils import PreTrainedModel21 22from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging23from .quantizers_utils import get_module_from_name24 25 26if is_torch_available():27    import torch28 29 30logger = logging.get_logger(__name__)31 32 33class EetqHfQuantizer(HfQuantizer):34    """35    8-bit quantization from EETQ quantization method:36        before loading: converts transformer layers into W8A16Linear during loading: load 16bit weight and pass to the37        layer object after: quantizes individual weights in Linear8bitLt into 8bit at first .cuda() call38    """39 40    requires_parameters_quantization = True41    requires_calibration = False42 43    required_packages = ["eetq", "accelerate"]44 45    def __init__(self, quantization_config, **kwargs):46        super().__init__(quantization_config, **kwargs)47        self.quantization_config = quantization_config48 49    def validate_environment(self, *args, **kwargs):50        if not is_eetq_available():51            raise ImportError(52                "Using `eetq` 8-bit quantization requires eetq."53                "Please install the latest version of eetq from : https://github.com/NetEase-FuXi/EETQ"54            )55 56        try:57            import eetq  # noqa: F40158        except ImportError as exc:59            if "shard_checkpoint" in str(exc):60                # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed61                # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34.62                # TODO: Update message once eetq releases a fix63                raise ImportError(64                    "You are using a version of EETQ that is incompatible with the current transformers version. "65                    "Either downgrade transformers to <= v4.46.3 or, if available, upgrade EETQ to > v1.0.0."66                ) from exc67            else:68                raise69 70        if not is_accelerate_available():71            raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)")72 73        if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):74            raise ValueError(75                "Converting into 8-bit weights from tf/flax weights is currently not supported, please make"76                " sure the weights are in PyTorch format."77            )78 79        if not torch.cuda.is_available():80            raise RuntimeError("No GPU found. A GPU is needed for quantization.")81 82        device_map = kwargs.get("device_map")83        if device_map is None:84            logger.warning_once(85                "You have loaded an EETQ model on CPU and have a CUDA device available, make sure to set "86                "your model on a GPU device in order to run your model."87            )88        elif device_map is not None:89            if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):90                raise ValueError(91                    "You are attempting to load an EETQ model with a device_map that contains a CPU or disk device."92                    " This is not supported. Please remove the CPU or disk device from the device_map."93                )94 95    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":96        if dtype is None:97            dtype = torch.float1698            logger.info(99                "Overriding dtype=%s with `dtype=torch.float16` due to "100                "requirements of `eetq` to enable model loading in 8-bit. "101                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"102                " dtype=torch.float16 to remove this warning.",103                dtype,104            )105        elif dtype != torch.float16:106            logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with EETQ.")107        return dtype108 109    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:110        from eetq import EetqLinear111 112        module, tensor_name = get_module_from_name(model, param_name)113 114        if isinstance(module, EetqLinear):115            if self.pre_quantized or tensor_name == "bias":116                return False117            else:118                return True119        return False120 121    def create_quantized_param(122        self,123        model: "PreTrainedModel",124        param_value: "torch.Tensor",125        param_name: str,126        target_device: "torch.device",127        **kwargs,128    ):129        from eetq import EetqLinear, quantize_and_preprocess_weights130 131        module, tensor_name = get_module_from_name(model, param_name)132        new_value, weight_scale = quantize_and_preprocess_weights(param_value)133 134        # Samity check135        if isinstance(module, EetqLinear):136            if self.pre_quantized or tensor_name == "bias":137                if tensor_name == "weight" and param_value.dtype != torch.int8:138                    raise ValueError("Expect quantized weights but got an unquantized weight")139            else:140                if tensor_name == "weight_scale":141                    raise ValueError("Expect unquantized weights but got a quantized weight_scale")142 143        module._buffers[tensor_name] = new_value.to(target_device)144        module.register("weight_scales", weight_scale.to(target_device))145 146    def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):147        return model148 149    def _process_model_before_weight_loading(150        self,151        model: "PreTrainedModel",152        keep_in_fp32_modules: Optional[list[str]] = None,153        **kwargs,154    ):155        from ..integrations import replace_with_eetq_linear156 157        self.modules_to_not_convert = self.get_modules_to_not_convert(158            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules159        )160 161        model = replace_with_eetq_linear(162            model,163            modules_to_not_convert=self.modules_to_not_convert,164            quantization_config=self.quantization_config,165            pre_quantized=self.pre_quantized,166        )167 168        model.config.quantization_config = self.quantization_config169 170    def is_serializable(self, safe_serialization=None):171        return True172 173    @property174    def is_trainable(self) -> bool:175        return True176