nl_tasks/rpeft/mapping.py

# coding=utf-8
# Original License:
# Copyright 2023-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .peft_model import (
    PeftModel,
    PeftModelForCausalLM,
    PeftModelForSeq2SeqLM,
    PeftModelForSequenceClassification,
    PeftModelForTokenClassification,
)
from .rotation import RotationConfig, RotationTuner
from .utils import PromptLearningConfig

from transformers import PreTrainedModel

MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
    "SEQ_CLS": PeftModelForSequenceClassification,
    "SEQ_2_SEQ_LM": PeftModelForSeq2SeqLM,
    "CAUSAL_LM": PeftModelForCausalLM,
    "TOKEN_CLS": PeftModelForTokenClassification,
}

PEFT_TYPE_TO_CONFIG_MAPPING: dict = {
    "ROTATION": RotationConfig,
}

PEFT_TYPE_TO_TUNER_MAPPING: dict = {
    "ROTATION": RotationTuner
}

TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
    "t5": ["q", "v"],
    "mt5": ["q", "v"],
    "bart": ["q_proj", "v_proj"],
    "gpt2": ["c_attn"],
    "bloom": ["query_key_value"],
    "blip-2": ["q", "v", "q_proj", "v_proj"],
    "opt": ["q_proj", "v_proj"],
    "gptj": ["q_proj", "v_proj"],
    "gpt_neox": ["query_key_value"],
    "gpt_neo": ["q_proj", "v_proj"],
    "bert": ["query", "value"],
    "roberta": ["query", "value"],
    "xlm-roberta": ["query", "value"],
    "electra": ["query", "value"],
    "deberta-v2": ["query_proj", "value_proj"],
    "deberta": ["in_proj"],
    "layoutlm": ["query", "value"],
    "llama": ["q_proj", "v_proj"],
    "chatglm": ["query_key_value"],
    "gpt_bigcode": ["c_attn"],
    "mpt": ["Wqkv"],
    "RefinedWebModel": ["query_key_value"],
    "RefinedWeb": ["query_key_value"],
    "falcon": ["query_key_value"],
    "btlm": ["c_proj", "c_attn"],
    "codegen": ["qkv_proj"],
    "mistral": ["q_proj", "v_proj"],
    "mixtral": ["q_proj", "v_proj"],
    "stablelm": ["q_proj", "v_proj"],
    "phi": ["q_proj", "v_proj", "fc1", "fc2"],
    "gemma": ["q_proj", "v_proj"],
}



def get_peft_config(config_dict):
    """
    Returns a Peft config object from a dictionary.

    Args:
        config_dict (`Dict[str, Any]`): Dictionary containing the configuration parameters.
    """

    return PEFT_TYPE_TO_CONFIG_MAPPING[config_dict["peft_type"]](**config_dict)


def _prepare_prompt_learning_config(peft_config, model_config):
    if peft_config.num_layers is None:
        if "num_hidden_layers" in model_config:
            num_layers = model_config["num_hidden_layers"]
        elif "num_layers" in model_config:
            num_layers = model_config["num_layers"]
        elif "n_layer" in model_config:
            num_layers = model_config["n_layer"]
        else:
            raise ValueError("Please specify `num_layers` in `peft_config`")
        peft_config.num_layers = num_layers

    if peft_config.token_dim is None:
        if "hidden_size" in model_config:
            token_dim = model_config["hidden_size"]
        elif "n_embd" in model_config:
            token_dim = model_config["n_embd"]
        elif "d_model" in model_config:
            token_dim = model_config["d_model"]
        else:
            raise ValueError("Please specify `token_dim` in `peft_config`")
        peft_config.token_dim = token_dim

    if peft_config.num_attention_heads is None:
        if "num_attention_heads" in model_config:
            num_attention_heads = model_config["num_attention_heads"]
        elif "n_head" in model_config:
            num_attention_heads = model_config["n_head"]
        elif "num_heads" in model_config:
            num_attention_heads = model_config["num_heads"]
        elif "encoder_attention_heads" in model_config:
            num_attention_heads = model_config["encoder_attention_heads"]
        else:
            raise ValueError("Please specify `num_attention_heads` in `peft_config`")
        peft_config.num_attention_heads = num_attention_heads

    if getattr(peft_config, "encoder_hidden_size", None) is None:
        setattr(peft_config, "encoder_hidden_size", token_dim)

    return peft_config


def _prepare_lora_config(peft_config, model_config):
    if peft_config.target_modules is None:
        if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
            raise ValueError("Please specify `target_modules` in `peft_config`")
        peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]]
    if len(peft_config.target_modules) == 1:
        peft_config.fan_in_fan_out = True
        peft_config.enable_lora = [True, False, True]
    if peft_config.inference_mode:
        peft_config.merge_weights = True
    return peft_config
    


def get_peft_model(model, peft_config,
                    adapter_name: str = "default"):
    """
    Returns a Peft model object from a model and a config.

    Args:
        model ([`transformers.PreTrainedModel`]): Model to be wrapped.
        peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
    """
    model_config = model.config.to_dict()
    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys():
        if peft_config.peft_type == "LORA" or "QUANTA":
            peft_config = _prepare_lora_config(peft_config, model_config)
            return PeftModel(model, peft_config)
    if not isinstance(peft_config, PromptLearningConfig):
        if peft_config.peft_type == "LORA" or "QUANTA":
            peft_config = _prepare_lora_config(peft_config, model_config)
    else:
        peft_config = _prepare_prompt_learning_config(peft_config, model_config)
    # assert False
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](
        model,
        peft_config,
        adapter_name=adapter_name,
    )


# def get_peft_model(
#     model: PreTrainedModel,
#     peft_config,
#     adapter_name: str = "default",
#     mixed: bool = False,
#     autocast_adapter_dtype: bool = True,
#     revision: Optional[str] = None,
#     low_cpu_mem_usage: bool = False,
# ) -> PeftModel | PeftMixedModel:
#     """
#     Returns a Peft model object from a model and a config, where the model will be modified in-place.

#     Args:
#         model ([`transformers.PreTrainedModel`]):
#             Model to be wrapped.
#         peft_config ([`PeftConfig`]):
#             Configuration object containing the parameters of the Peft model.
#         adapter_name (`str`, `optional`, defaults to `"default"`):
#             The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
#         mixed (`bool`, `optional`, defaults to `False`):
#             Whether to allow mixing different (compatible) adapter types.
#         autocast_adapter_dtype (`bool`, *optional*):
#             Whether to autocast the adapter dtype. Defaults to `True`. Right now, this will only cast adapter weights
#             using float16 or bfloat16 to float32, as this is typically required for stable training, and only affect
#             select PEFT tuners.
#         revision (`str`, `optional`, defaults to `main`):
#             The revision of the base model. If this isn't set, the saved peft model will load the `main` revision for
#             the base model
#         low_cpu_mem_usage (`bool`, `optional`, defaults to `False`):
#             Create empty adapter weights on meta device. Useful to speed up the loading process. Leave this setting as
#             False if you intend on training the model, unless the adapter weights will be replaced by different weights
#             before training starts.
#     """
#     model_config = BaseTuner.get_model_config(model)
#     old_name = peft_config.base_model_name_or_path
#     new_name = model.__dict__.get("name_or_path", None)
#     peft_config.base_model_name_or_path = new_name

#     # Especially in notebook environments there could be a case that a user wants to experiment with different
#     # configuration values. However, it is likely that there won't be any changes for new configs on an already
#     # initialized PEFT model. The best we can do is warn the user about it.
#     if any(isinstance(module, BaseTunerLayer) for module in model.modules()):
#         warnings.warn(
#             "You are trying to modify a model with PEFT for a second time. If you want to reload the model with a "
#             "different config, make sure to call `.unload()` before."
#         )

#     if (old_name is not None) and (old_name != new_name):
#         warnings.warn(
#             f"The PEFT config's `base_model_name_or_path` was renamed from '{old_name}' to '{new_name}'. "
#             "Please ensure that the correct base model is loaded when loading this checkpoint."
#         )

#     if revision is not None:
#         if peft_config.revision is not None and peft_config.revision != revision:
#             warnings.warn(
#                 f"peft config has already set base model revision to {peft_config.revision}, overwriting with revision {revision}"
#             )
#         peft_config.revision = revision

#     if (
#         (isinstance(peft_config, PEFT_TYPE_TO_CONFIG_MAPPING["LORA"]))
#         and (peft_config.init_lora_weights == "eva")
#         and not low_cpu_mem_usage
#     ):
#         warnings.warn(
#             "lora with eva initialization used with low_cpu_mem_usage=False. "
#             "Setting low_cpu_mem_usage=True can improve the maximum batch size possible for eva initialization."
#         )

#     prefix = PEFT_TYPE_TO_PREFIX_MAPPING.get(peft_config.peft_type)
#     if prefix and adapter_name in prefix:
#         warnings.warn(
#             f"Adapter name '{adapter_name}' should not be contained in the prefix '{prefix}'. "
#             "This may lead to reinitialization of the adapter weights during loading."
#         )

#     if mixed:
#         # note: PeftMixedModel does not support autocast_adapter_dtype, so don't pass it
#         return PeftMixedModel(model, peft_config, adapter_name=adapter_name)

#     # We explicitly exclude prompt learning here since prompt learning is specific to the task and needs special
#     # handling in the PEFT model's forward method.
#     if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
#         return PeftModel(
#             model,
#             peft_config,
#             adapter_name=adapter_name,
#             autocast_adapter_dtype=autocast_adapter_dtype,
#             low_cpu_mem_usage=low_cpu_mem_usage,
#         )

#     return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](
#         model,
#         peft_config,
#         adapter_name=adapter_name,
#         autocast_adapter_dtype=autocast_adapter_dtype,
#         low_cpu_mem_usage=low_cpu_mem_usage,
#     )