Aluode/PerceptionLabPortable

1# Copyright 2025 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 HfQuantizer17from .quantizers_utils import get_module_from_name18 19 20if TYPE_CHECKING:21    from ..modeling_utils import PreTrainedModel22 23from ..utils import is_fp_quant_available, is_qutlass_available, is_torch_available, logging24from ..utils.quantization_config import QuantizationConfigMixin25 26 27if is_torch_available():28    import torch29 30logger = logging.get_logger(__name__)31 32 33class FPQuantHfQuantizer(HfQuantizer):34    """35    Quantizer for the FP-Quant method. Enables the loading of prequantized models and in-flight quantization of full-precision models.36    """37 38    requires_calibration = False39    requires_parameters_quantization = True40    is_qat_trainable = True41    required_packages = ["fp_quant"]42 43    def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):44        super().__init__(quantization_config, **kwargs)45        self.quantization_config = quantization_config46 47    def validate_environment(self, device_map, **kwargs):48        if not torch.cuda.is_available():49            raise NotImplementedError(50                "FPQuant quantization is only supported on GPU. Please use a different quantizer."51            )52 53        if not is_qutlass_available() and not self.quantization_config.pseudoquantization:54            raise ImportError(55                "Using `fp_quant` with real quantization requires a **Blackwell GPU** and qutlass: `git clone https://github.com/IST-DASLab/qutlass.git && cd qutlass && pip install --no-build-isolation .`. You can use `FPQuantConfig(pseudoquantization=True, ...)` to use Triton-based pseudo-quantization. It doesn't provide any speedups but emulates the quantization behavior of the real quantization."56            )57 58        if self.quantization_config.pseudoquantization:59            logger.warning(60                "Using pseudo-quantization for FP-Quant. This doesn't provide any speedups but emulates the quantization behavior of the real quantization."61            )62 63        if not is_fp_quant_available():64            raise ImportError("Using `fp_quant` quantization requires fp_quant: `pip install fp_quant`")65 66        if device_map is None and not self.quantization_config.pseudoquantization:67            raise ValueError(68                "You are attempting to load a FPQuant model without setting device_map."69                " Please set device_map comprised of 'cuda' devices."70            )71        elif (72            isinstance(device_map, dict)73            and ("cpu" in device_map.values() or "disk" in device_map.values())74            and not self.quantization_config.pseudoquantization75        ):76            raise ValueError(77                "You are attempting to load a FPQuant model with a device_map that contains a CPU or disk device."78                " This is not supported. Please remove the CPU or disk device from the device_map."79            )80 81    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":82        if dtype is None:83            logger.info("`dtype` is None. Setting `dtype=torch.bfloat16` for qutlass compatibility.")84            dtype = torch.bfloat1685        elif dtype != torch.bfloat16:86            raise ValueError(f"Invalid `dtype` {dtype}. fp_quant quantization only supports `dtype=torch.bfloat16`.")87 88        return dtype89 90    def create_quantized_param(91        self,92        model: "PreTrainedModel",93        param_value: "torch.Tensor",94        param_name: str,95        target_device: "torch.device",96        **kwargs,97    ):98        module, _ = get_module_from_name(model, param_name)99 100        # The module holds either:101        #  * `weight` when `store_master_weights=True`102        #  * `qweight` and `scales` when `store_master_weights=False` and `pseudoquantization=False`103        #  * `dqweight` when `store_master_weights=False` and `pseudoquantization=True`104 105        if param_name.endswith(".qweight"):106            # Loading a real quantized checkpoint without master weights107            module.qweight = torch.nn.Parameter(108                param_value.to(target_device),109                requires_grad=False,110            )111            module.weight = None112            module.dqweight = None113            return114 115        if param_name.endswith(".dqweight"):116            # Loading a pseudo-quantized checkpoint without master weights117            module.dqweight = torch.nn.Parameter(param_value.to(target_device))118            module.weight = None119            module.qweight = None120            module.scales = None121            return122 123        # Loading master weights or an unquantized checkpoint124        module.weight = torch.nn.Parameter(param_value.to(target_device))125        # Let pre-forward handle the quantization and set None where necessary126        module.pre_forward()127 128    def _process_model_before_weight_loading(129        self,130        model: "PreTrainedModel",131        **kwargs,132    ):133        from fp_quant import replace_with_fp_quant_linear134 135        from ..integrations.fp_quant import adapt_fp_quant_config136 137        replace_with_fp_quant_linear(138            model,139            fp_quant_linear_config=adapt_fp_quant_config(self.quantization_config),140        )141        model.config.quantization_config = self.quantization_config142 143    def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):144        return model145 146    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:147        from fp_quant import FPQuantLinear148 149        fp_quant_names = {name for name, module in model.named_modules() if isinstance(module, FPQuantLinear)}150 151        def should_exclude(key: str) -> bool:152            if key.endswith(".weight") or key.endswith(".bias"):153                return False154            full_key = f"{prefix}.{key}"155            return any(name in key or name in full_key for name in fp_quant_names)156 157        return [key for key in missing_keys if not should_exclude(key)]158 159    @property160    def is_trainable(self, model: Optional["PreTrainedModel"] = None):161        trainable = self.quantization_config.store_master_weights162        if not trainable:163            logger.warning(164                "You are attempting to train a model with FPQuant quantization. This is only supported when `store_master_weights=True`. Please set `store_master_weights=True` to train the model."165            )166        return trainable167 168    def is_serializable(self, safe_serialization=None):169        return True170 171    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:172        from fp_quant import FPQuantLinear173 174        module, tensor_name = get_module_from_name(model, param_name)175        if isinstance(module, FPQuantLinear) and tensor_name in ["weight", "qweight", "dqweight"]:176            # Only quantize weights of FPQuantLinear modules that are not already quantized177            return True178        else:179            return False180