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_torch_available, is_vptq_available, logging23from ..utils.quantization_config import QuantizationConfigMixin24 25 26if is_torch_available():27    import torch28 29logger = logging.get_logger(__name__)30 31 32class VptqHfQuantizer(HfQuantizer):33    """34    Quantizer of the VPTQ method. Enables the loading of prequantized models.35    """36 37    requires_calibration = True38    required_packages = ["vptq"]39 40    def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):41        super().__init__(quantization_config, **kwargs)42        self.quantization_config = quantization_config43 44    def validate_environment(self, *args, **kwargs):45        if not is_accelerate_available():46            raise ImportError("Using `vptq` quantization requires Accelerate: `pip install accelerate`")47 48        if not is_vptq_available():49            raise ImportError("Using `vptq` quantization requires VPTQ>=0.0.4: `pip install -U vptq`")50 51    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":52        if dtype is None:53            if torch.cuda.is_available():54                dtype = torch.float1655                logger.info(56                    "CUDA available. Assuming VPTQ inference on GPU and loading the model in `torch.float16`. To overwrite it, set `dtype` manually."57                )58            else:59                import vptq60 61                device_availability = getattr(vptq, "device_availability", lambda device: False)62                if device_availability("cpu") is True:63                    raise RuntimeError("No GPU found. Please wait for the next release of VPTQ to use CPU inference")64                dtype = torch.float3265                logger.info("No GPU found. Assuming VPTQ inference on CPU and loading the model in `torch.float32`.")66        return dtype67 68    def _process_model_before_weight_loading(69        self,70        model: "PreTrainedModel",71        keep_in_fp32_modules: Optional[list[str]] = None,72        **kwargs,73    ):74        """75        we don't have param like modules_to_not_convert to indicate which layers should not be quantized76        because `quantization_config` include the layers that should be quantized77        """78        from ..integrations import replace_with_vptq_linear79 80        self.modules_to_not_convert = self.get_modules_to_not_convert(81            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules82        )83 84        replace_with_vptq_linear(85            model,86            quantization_config=self.quantization_config,87            modules_to_not_convert=self.modules_to_not_convert,88        )89        model.config.quantization_config = self.quantization_config90 91    def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):92        return model93 94    @property95    def is_trainable(self) -> bool:96        return False97 98    def is_serializable(self, safe_serialization=None):99        return True100