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 abc import ABC, abstractmethod15from typing import TYPE_CHECKING, Any, Optional, Union16 17from ..utils import is_torch_available, logging18from ..utils.quantization_config import QuantizationConfigMixin, QuantizationMethod19from .quantizers_utils import get_module_from_name20 21 22if TYPE_CHECKING:23    from ..modeling_utils import PreTrainedModel24 25if is_torch_available():26    import torch27    from torch.nn import ModuleList28else:29    ModuleList = str30 31logger = logging.get_logger(__file__)32 33 34class HfQuantizer(ABC):35    """36    Abstract class of the HuggingFace quantizer. Supports for now quantizing HF transformers models for inference and/or quantization.37    This class is used only for transformers.PreTrainedModel.from_pretrained and cannot be easily used outside the scope of that method38    yet.39 40    Attributes41        quantization_config (`transformers.utils.quantization_config.QuantizationConfigMixin`):42            The quantization config that defines the quantization parameters of your model that you want to quantize.43        modules_to_not_convert (`list[str]`, *optional*):44            The list of module names to not convert when quantizing the model.45        required_packages (`list[str]`, *optional*):46            The list of required pip packages to install prior to using the quantizer47        requires_calibration (`bool`):48            Whether the quantization method requires to calibrate the model before using it.49        requires_parameters_quantization (`bool`):50            Whether the quantization method requires to create a new Parameter. For example, for bitsandbytes, it is51            required to create a new xxxParameter in order to properly quantize the model.52    """53 54    requires_calibration = False55    required_packages = None56    requires_parameters_quantization = False57 58    def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):59        self.quantization_config = quantization_config60 61        # -- Handle extra kwargs below --62        self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", [])63        self.pre_quantized = kwargs.pop("pre_quantized", True)64 65        if not self.pre_quantized and self.requires_calibration:66            raise ValueError(67                f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized."68                f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to "69                f"pass `pre_quantized=True` while knowing what you are doing."70            )71 72    def update_torch_dtype(self, dtype: "torch.dtype") -> "torch.dtype":73        """74        Deprecared in favor of `update_dtype`!75 76        Args:77            dtype (`torch.dtype`):78                The input dtype that is passed in `from_pretrained`79        """80        logger.warning_once(81            "`update_torch_dtype` is deprecated in favor of `update_dtype`! It will be removed in version v4.57"82        )83        return self.update_dtype(dtype)84 85    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":86        """87        Some quantization methods require to explicitly set the dtype of the model to a88        target dtype. You need to override this method in case you want to make sure that behavior is89        preserved90 91        Args:92            dtype (`torch.dtype`):93                The input dtype that is passed in `from_pretrained`94        """95        return dtype96 97    def update_device_map(self, device_map: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:98        """99        Override this method if you want to pass a override the existing device map with a new100        one. E.g. for bitsandbytes, since `accelerate` is a hard requirement, if no device_map is101        passed, the device_map is set to `"auto"``102 103        Args:104            device_map (`Union[dict, str]`, *optional*):105                The device_map that is passed through the `from_pretrained` method.106        """107        return device_map108 109    def adjust_target_dtype(self, dtype: "torch.dtype") -> "torch.dtype":110        """111        Override this method if you want to adjust the `target_dtype` variable used in `from_pretrained`112        to compute the device_map in case the device_map is a `str`. E.g. for bitsandbytes we force-set `target_dtype`113        to `torch.int8` and for 4-bit we pass a custom enum `accelerate.CustomDtype.int4`.114 115        Args:116            dtype (`torch.dtype`, *optional*):117                The dtype that is used to compute the device_map.118        """119        return dtype120 121    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:122        """123        Override this method if you want to adjust the `missing_keys`.124 125        Args:126            missing_keys (`list[str]`, *optional*):127                The list of missing keys in the checkpoint compared to the state dict of the model128        """129        return missing_keys130 131    def update_expected_keys(self, model, expected_keys: list[str], loaded_keys: list[str]) -> list[str]:132        """133        Override this method if you want to adjust the `update_expected_keys`.134 135        Args:136            expected_keys (`list[str]`, *optional*):137                The list of the expected keys in the initialized model.138            loaded_keys (`list[str]`, *optional*):139                The list of the loaded keys in the checkpoint.140        """141        return expected_keys142 143    def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:144        return unexpected_keys145 146    def get_special_dtypes_update(self, model, dtype: "torch.dtype") -> dict[str, "torch.dtype"]:147        """148        returns dtypes for modules that are not quantized - used for the computation of the device_map in case149        one passes a str as a device_map. The method will use the `modules_to_not_convert` that is modified150        in `_process_model_before_weight_loading`.151 152        Args:153            model (`~transformers.PreTrainedModel`):154                The model to quantize155            dtype (`torch.dtype`):156                The dtype passed in `from_pretrained` method.157        """158 159        return {160            name: dtype for name, _ in model.named_parameters() if any(m in name for m in self.modules_to_not_convert)161        }162 163    def adjust_max_memory(self, max_memory: dict[str, Union[int, str]]) -> dict[str, Union[int, str]]:164        """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization"""165        return max_memory166 167    def check_quantized_param(self, *args, **kwargs) -> bool:168        """DEPRECATED -> remove in v5"""169        logger.warning_once(170            "`check_quantized_param` is deprecated in favor of `param_needs_quantization`, which is a much "171            "more self.explanatory name for what the method achieves. It will be removed in v5"172        )173        return self.param_needs_quantization(*args, **kwargs)174 175    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:176        """177        Check whether a given param needs quantization as defined by `create_quantized_param`.178        """179        return False180 181    def create_quantized_param(self, *args, **kwargs):182        """183        Take needed components from state_dict (those from which `param_needs_quantization` is True) and create184        quantized param.185        It usually also load the new param directly in the `model`.186        Note: only applicable if requires_parameters_quantization == True.187        """188        if not self.requires_parameters_quantization:189            raise AttributeError(190                f"`.create_quantized_param()` method is not supported by quantizer class {self.__class__.__name__}."191            )192 193    def validate_environment(self, *args, **kwargs):194        """195        This method is used to potentially check for potential conflicts with arguments that are196        passed in `from_pretrained`. You need to define it for all future quantizers that are integrated with transformers.197        If no explicit check are needed, simply return nothing.198        """199        return200 201    def update_tp_plan(self, config):202        "updates the tp plan for the scales"203        return config204 205    def update_ep_plan(self, config):206        "updates the tp plan for the scales"207        return config208 209    def preprocess_model(self, model: "PreTrainedModel", **kwargs):210        """211        Setting model attributes and/or converting model before weights loading. At this point212        the model should be initialized on the meta device so you can freely manipulate the skeleton213        of the model in order to replace modules in-place. Make sure to override the abstract method `_process_model_before_weight_loading`.214 215        Args:216            model (`~transformers.PreTrainedModel`):217                The model to quantize218            kwargs (`dict`, *optional*):219                The keyword arguments that are passed along `_process_model_before_weight_loading`.220        """221        model.is_quantized = True222        model.quantization_method = self.quantization_config.quant_method223        if self.pre_quantized:224            self._convert_model_for_quantization(model)225        return self._process_model_before_weight_loading(model, **kwargs)226 227    def postprocess_model(self, model: "PreTrainedModel", **kwargs):228        """229        Post-process the model post weights loading.230        Make sure to override the abstract method `_process_model_after_weight_loading`.231 232        Args:233            model (`~transformers.PreTrainedModel`):234                The model to quantize235            kwargs (`dict`, *optional*):236                The keyword arguments that are passed along `_process_model_after_weight_loading`.237        """238        return self._process_model_after_weight_loading(model, **kwargs)239 240    def remove_quantization_config(self, model):241        """242        Remove the quantization config from the model.243        """244        if hasattr(model, "hf_quantizer"):245            del model.hf_quantizer246        if hasattr(model.config, "quantization_config"):247            del model.config.quantization_config248        if hasattr(model.config, "_pre_quantization_dtype"):249            del model.config._pre_quantization_dtype250        if hasattr(model, "quantization_method"):251            del model.quantization_method252        model.is_quantized = False253 254    def dequantize(self, model):255        """256        Potentially dequantize the model to retrieve the original model, with some loss in accuracy / performance.257        Note not all quantization schemes support this.258        """259        model = self._dequantize(model)260 261        # Delete quantizer and quantization config262        del model.hf_quantizer263        del model.config.quantization_config264        del model.config._pre_quantization_dtype265        del model.quantization_method266        model.is_quantized = False267 268        return model269 270    def get_accelerator_warm_up_factor(self):271        """272        The factor to be used in `caching_allocator_warmup` to get the number of bytes to pre-allocate to warm up accelerator.273        A factor of 2 means we allocate all bytes in the empty model (since we allocate in fp16), a factor of 4 means274        we allocate half the memory of the weights residing in the empty model, etc...275        """276        # By default we return 4, i.e. half the model size (this corresponds to the case where the model is not277        # really pre-processed, i.e. we do not have the info that weights are going to be 8 bits before actual278        # weight loading)279        return 4280 281    def _dequantize(self, model):282        raise NotImplementedError(283            f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub."284        )285 286    def get_param_name(self, param_name: str) -> str:287        """288        Override this method if you want to adjust the `param_name`.289        """290        return param_name291 292    @staticmethod293    def get_modules_to_not_convert(294        model: "PreTrainedModel",295        skip_modules: Optional[list[str]] = None,296        keep_in_fp32_modules: Optional[list[str]] = None,297        add_default_skips: bool = False,298    ):299        from ..integrations import get_keys_to_not_convert300 301        if skip_modules is None or add_default_skips:302            modules_to_not_convert = get_keys_to_not_convert(model)303        else:304            modules_to_not_convert = []305 306        if skip_modules is not None:307            modules_to_not_convert.extend(skip_modules)308 309        if keep_in_fp32_modules is not None:310            modules_to_not_convert.extend(keep_in_fp32_modules)311 312        return modules_to_not_convert313 314    @property315    def is_qat_trainable(self) -> bool:316        """Flag indicating whether the quantized model can carry out quantization aware training"""317        return False318 319    @property320    def is_compileable(self) -> bool:321        """Flag indicating whether the quantized model can be compiled"""322        return False323 324    def get_state_dict_and_metadata(self, model, safe_serialization=False):325        """Get state dict and metadata. Useful when we need to modify a bit the state dict due to quantization"""326        return None, {}327 328    def update_state_dict_with_metadata(self, state_dict, metadata):329        """Update state dict with metadata. Default behaviour returns state_dict"""330        return state_dict331 332    @abstractmethod333    def _process_model_before_weight_loading(self, model, **kwargs): ...334 335    @abstractmethod336    def _process_model_after_weight_loading(self, model, **kwargs): ...337 338    @abstractmethod339    def is_serializable(self, safe_serialization=None): ...340 341    @property342    @abstractmethod343    def is_trainable(self): ...344 345    def _convert_model_for_quantization(self, model):346        from accelerate import init_empty_weights347 348        for name, module in model.named_modules():349            module_class_name = module.__class__.__name__350            if module_class_name in MODULES_TO_PATCH_FOR_QUANTIZATION and (351                self.quantization_config.quant_method352                in MODULES_TO_PATCH_FOR_QUANTIZATION[module_class_name]["quantization_methods"]353            ):354                with init_empty_weights():355                    parent_module, name = get_module_from_name(model, name)356                    parent_module._modules[name] = MODULES_TO_PATCH_FOR_QUANTIZATION[module_class_name]["module_name"](357                        model.config.get_text_config()358                    )359 360 361class SequentialLlama4TextExperts(ModuleList):362    """363    A module that implements a compressed version of a list of expert modules.364    This is specifically designed to work with Llama4TextExperts in MoE layers.365    """366 367    def __init__(self, config):368        from transformers.models.llama4.modeling_llama4 import Llama4TextMLP369 370        super().__init__([Llama4TextMLP(config) for _ in range(config.num_local_experts)])371        self.num_experts = config.num_local_experts372 373    def forward(374        self,375        hidden_states: "torch.Tensor",376    ) -> "torch.Tensor":377        hidden_states = hidden_states.reshape(self.num_experts, -1, hidden_states.shape[-1])378        routed_out = torch.zeros_like(hidden_states)379        for expert_idx in range(self.num_experts):380            routed_out[expert_idx] = self[expert_idx](hidden_states[expert_idx])381        return routed_out382 383 384MODULES_TO_PATCH_FOR_QUANTIZATION = {385    "Llama4TextExperts": {386        "module_name": SequentialLlama4TextExperts,387        "quantization_methods": [388            QuantizationMethod.COMPRESSED_TENSORS,389            QuantizationMethod.BITS_AND_BYTES,390        ],391    }392}393