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 ..utils import is_optimum_quanto_available, is_torch_available, logging16 17 18if is_torch_available():19    import torch20 21logger = logging.get_logger(__name__)22 23 24def replace_with_quanto_layers(25    model,26    quantization_config=None,27    modules_to_not_convert=None,28    current_key_name=None,29    has_been_replaced=False,30):31    """32    Public method that recursively replaces the Linear layers of the given model with Quanto quantized layers.33    Returns the converted model and a boolean that indicates if the conversion has been successful or not.34 35    Args:36        model (`torch.nn.Module`):37            The model to convert, can be any `torch.nn.Module` instance.38        quantization_config (`AqlmConfig`, defaults to `None`):39            The quantization config object that contains the quantization parameters.40        modules_to_not_convert (`list`, *optional*, defaults to `None`):41            A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be42            converted.43        current_key_name (`list`, *optional*, defaults to `None`):44            A list that contains the current key name. This is used for recursion and should not be passed by the user.45        has_been_replaced (`bool`, *optional*, defaults to `None`):46            A boolean that indicates if the conversion has been successful or not. This is used for recursion and47            should not be passed by the user.48    """49    from accelerate import init_empty_weights50 51    if is_optimum_quanto_available():52        from optimum.quanto import QLayerNorm, QLinear, qfloat8, qint2, qint4, qint853 54    w_mapping = {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2}55    a_mapping = {None: None, "float8": qfloat8, "int8": qint8}56 57    if modules_to_not_convert is None:58        modules_to_not_convert = []59 60    for name, module in model.named_children():61        if current_key_name is None:62            current_key_name = []63        current_key_name.append(name)64 65        if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):66            with init_empty_weights():67                if isinstance(module, torch.nn.Linear):68                    model._modules[name] = QLinear(69                        in_features=module.in_features,70                        out_features=module.out_features,71                        bias=module.bias is not None,72                        dtype=module.weight.dtype,73                        weights=w_mapping[quantization_config.weights],74                        activations=a_mapping[quantization_config.activations],75                    )76                    model._modules[name].requires_grad_(False)77                    has_been_replaced = True78                elif isinstance(module, torch.nn.LayerNorm):79                    if quantization_config.activations is not None:80                        model._modules[name] = QLayerNorm(81                            module.normalized_shape,82                            module.eps,83                            module.elementwise_affine,84                            module.bias is not None,85                            activations=a_mapping[quantization_config.activations],86                        )87                        has_been_replaced = True88        if len(list(module.children())) > 0:89            _, has_been_replaced = replace_with_quanto_layers(90                module,91                quantization_config=quantization_config,92                modules_to_not_convert=modules_to_not_convert,93                current_key_name=current_key_name,94                has_been_replaced=has_been_replaced,95            )96        # Remove the last key for recursion97        current_key_name.pop(-1)98    return model, has_been_replaced99