Aluode/PerceptionLabPortable

1# Copyright 2024 The HuggingFace 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.14 15from collections import defaultdict16from typing import TYPE_CHECKING17 18from ..integrations import prepare_for_hqq_linear19from ..utils import is_hqq_available, is_torch_available, logging20from .base import HfQuantizer21from .quantizers_utils import get_module_from_name22 23 24if TYPE_CHECKING:25    from ..modeling_utils import PreTrainedModel26 27 28if is_torch_available():29    import torch30 31if is_hqq_available():32    from hqq.core.quantize import HQQLinear33 34    # This is a compatibility hack. HQQ-quantized linear layers do not have a `weight` attribute,35    # but some models attempt to access `weight.dtype` during the forward pass. To prevent runtime errors,36    # we patch HQQLinear with a dummy `weight` property that returns an empty tensor with the correct dtype and device.37    @property38    def weight(self):39        return torch.empty(0, dtype=self.compute_dtype, device=self.device)40 41    HQQLinear.weight = weight42 43logger = logging.get_logger(__name__)44 45 46class HqqHfQuantizer(HfQuantizer):47    """48    HQQ quantizer base HF class.49    nn.Linear modules are first tagged with quant_config in _process_model_before_weight_loading().50    """51 52    use_keep_in_fp32_modules = False53    requires_parameters_quantization = True54    requires_calibration = False55    required_packages = ["hqq"]56 57    def __init__(self, quantization_config, **kwargs):58        if not is_hqq_available():59            raise ImportError(60                "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`."61            )62        super().__init__(quantization_config, **kwargs)63        self.dtype = None64        self.using_multi_gpu = False65        # Keys that are serialized specifically by hqq66        self.hqq_keys = HQQLinear(None, None).state_dict_keys() - {"bias"}67 68        if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):69            raise ValueError(70                "Converting weights from tf/flax weights is currently not supported, please make"71                " sure the weights are in PyTorch format."72            )73 74        if self.dtype is None:75            if "dtype" in kwargs:76                self.dtype = kwargs["dtype"]77            else:78                self.dtype = torch.float3279                logger.info("Setting dtype to torch.float32 as the default value since it was not specified.")80 81        device_map = kwargs.get("device_map")82        if isinstance(device_map, dict):83            if "cpu" in device_map.values() or "disk" in device_map.values():84                raise ValueError(85                    "You are attempting to use an HQQ model with a device_map that contains a CPU or disk device."86                    " This is not supported. Please remove the CPU or disk device from the device_map."87                )88            else:89                self.using_multi_gpu = len(set(device_map.values())) > 190 91    def update_missing_keys(92        self, model: "PreTrainedModel", missing_keys: list[str], prefix: str, **kwargs93    ) -> list[str]:94        if self.pre_quantized:95            return [key for key in missing_keys if ("weight" not in key)]96        else:97            return missing_keys98 99    # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear100    def update_expected_keys(101        self, model: "PreTrainedModel", expected_keys: list[str], loaded_keys: list[str]102    ) -> list[str]:103        if not self.pre_quantized:104            return expected_keys105 106        # Collects all quantizable (linear) layers107        def _find_hqq_quantizable_layers(model, layers):108            for name, module in model.named_children():109                if isinstance(module, (torch.nn.Linear)):110                    layers.add(module.name)111                _find_hqq_quantizable_layers(module, layers)112 113        new_keys = set(expected_keys)114 115        # Name modules116        for name, module in model.named_modules():117            module.name = name118 119        # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params120        _valid_modules = set()121        _find_hqq_quantizable_layers(model, _valid_modules)122 123        # Remove skipped modules124        _skipped_modules = set()125        for _module in _valid_modules:126            for _skip_module in model.config.quantization_config["skip_modules"]:127                if _skip_module in _module:128                    _skipped_modules.add(_module)129        _valid_modules -= _skipped_modules130 131        # Append new expected layers based on _ref_keys132        _ref_keys = HQQLinear(133            linear_layer=None,134            quant_config=None,135            compute_dtype=torch.float16,136            device="cpu",137            del_orig=False,138        ).state_dict_keys() - {"bias"}139 140        # Clean-up141        _rm_keys = set()142        for key in new_keys:143            if any(_module in key for _module in _valid_modules):144                _rm_keys.add(key)145        new_keys -= _rm_keys146        # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear147 148        # Re-populate Linear/HQQLinear149        for _module in _valid_modules:150            if _module + ".weight" in loaded_keys:151                new_keys.add(_module + ".weight")152            else:153                new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})154            if _module + ".bias" in loaded_keys:155                new_keys.add(_module + ".bias")156 157        return list(new_keys)158 159    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:160        module, _ = get_module_from_name(model, param_name)161        # Since we do not prepare the modules in advance, we need every param of the Linear layer to go through162        # `create_quantized_param`, even when `self.is_quantized == True`163        return isinstance(module, torch.nn.Linear)164 165    def create_quantized_param(166        self,167        model: "PreTrainedModel",168        param_value: "torch.Tensor",169        param_name: str,170        target_device: "torch.device",171        **kwargs,172    ):173        module, tensor_name = get_module_from_name(model, param_name)174        module_name = param_name.rsplit(".", 1)[0]175        parent_module, node = get_module_from_name(model, module_name)176 177        quant_config = model.config.quantization_config["quant_config"]178        skip_modules = model.config.quantization_config["skip_modules"]179 180        # In this case we do not quantize this layer (it's explicitly skipped) -> simply load param181        if any(skip_module in module.name for skip_module in skip_modules):182            module.load_state_dict(183                {tensor_name: param_value.to(device=target_device, dtype=self.dtype)}, strict=False, assign=True184            )185            return186 187        # We need this hack as the model is not pre-prepared as an empty skeleton on meta device188        if self.pre_quantized:189            # Save them for later190            if not hasattr(self, "hqq_params"):191                self.hqq_params = defaultdict(dict)192            self.hqq_params[module_name].update({tensor_name: param_value})193            hqq_params = self.hqq_params[module_name]194 195            # If they are all present and saved, make it a HQQLinear layer! (we cannot do it param after param because196            # hqq does not support it...)197            if all(k in hqq_params for k in self.hqq_keys) and ("bias" in hqq_params or module.bias is None):198                hqq_layer = HQQLinear(199                    linear_layer=None,200                    quant_config=None,201                    compute_dtype=self.dtype,202                    device=target_device,203                    del_orig=False,204                )205                hqq_layer.load_state_dict(hqq_params)206 207                if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):208                    hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)209                if self.using_multi_gpu:210                    hqq_layer = self._patch_layer_for_multigpu(hqq_layer)211 212                setattr(parent_module, node, hqq_layer)213                del self.hqq_params[module_name], module214            return215 216        # Load param in the module (without caring about device or dtype, it will be changed later)217        module.load_state_dict({tensor_name: param_value}, strict=False, assign=True)218 219        # If both the weight and bias have already been loaded, time to quantize!220        module_is_ready = module.weight.device.type != "meta" and (221            module.bias is None or module.bias.device.type != "meta"222        )223 224        if module_is_ready:225            module_tag = ".".join(module.name.split(".")[-2:])226            if "weight_quant_params" in quant_config:227                module_quant_config = quant_config228            elif module_tag in quant_config:229                module_quant_config = quant_config[module_tag]230 231            hqq_layer = HQQLinear(232                module,233                quant_config=module_quant_config,234                compute_dtype=self.dtype,235                device=target_device,236                del_orig=True,237            )238 239            if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):240                hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)241 242            if self.using_multi_gpu:243                hqq_layer = self._patch_layer_for_multigpu(hqq_layer)244 245            setattr(parent_module, node, hqq_layer)246 247    def _patch_layer_for_multigpu(self, hqq_layer):248        def forward_with_device(self, x):249            out = torch.matmul(x.to(self.device), self.dequantize().t())250            if self.bias is not None:251                out += self.bias252            return out253 254        hqq_layer.forward = lambda x: forward_with_device(hqq_layer, x)255        return hqq_layer256 257    def _process_model_before_weight_loading(258        self,259        model: "PreTrainedModel",260        **kwargs,261    ):262        # Add the corresponding quant_config to each valid module. This allows us to do the actual nn.Linear -> HQQLinear conversion in create_quantized_param().263        # prepare_for_hqq_linear() also sets the right quantization config inside the model (model.config.quantization_config) and the layers (hqq_layer.quant_config)264        model = prepare_for_hqq_linear(model, quantization_config=self.quantization_config)265 266    def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):267        model.is_hqq_quantized = True268        model.is_hqq_serializable = self.is_serializable()269        return model270 271    def is_serializable(self, safe_serialization=None):272        return True273 274    @property275    def is_trainable(self) -> bool:276        return True277