optimum/habana/peft/peft_model.py

import warnings

import packaging.version
import torch
import transformers
from peft import PeftType


def gaudi_generate(self, *args, **kwargs):
    peft_config = self.active_peft_config
    self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation
    if hasattr(self.base_model, "model"):
        self.base_model.model.generation_config = self.generation_config
    else:
        self.base_model.generation_config = self.generation_config
    try:
        if not peft_config.is_prompt_learning:
            with self._enable_peft_forward_hooks(*args, **kwargs):
                kwargs = {k: v for k, v in kwargs.items() if k not in self.special_peft_forward_args}
                outputs = self.base_model.generate(*args, **kwargs)
        else:
            kwargs["num_virtual_tokens"] = peft_config.num_virtual_tokens
            outputs = self.base_model.generate(**kwargs)
    except:
        self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
        raise
    else:
        self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
        return outputs


def gaudi_prepare_inputs_for_generation(self, *args, task_ids: torch.Tensor = None, **kwargs):
    """
    Copied from PeftModelForCausalLM.prepare_inputs_for_generation: https://github.com/huggingface/peft/blob/v0.9.0/src/peft/peft_model.py#L1156
    The only differences are:
    - add token_idx disposal for prompt learning
    """
    peft_config = self.active_peft_config
    model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
    # https://github.com/huggingface/transformers/pull/26681/ introduced new cache format
    # for some architectures which requires a special fix for prompt tuning etc.
    # TODO: starting with transformers 4.38, all architectures should support caching.
    uses_transformers_4_38 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.38.0")
    uses_transformers_4_36 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.36.0")
    transformers_new_cache_archs = ["llama", "mistral", "persimmon", "phi"]
    uses_cache = uses_transformers_4_38 or (
        uses_transformers_4_36 and self.base_model.config.model_type in transformers_new_cache_archs
    )
    if peft_config.peft_type == PeftType.POLY:
        model_kwargs["task_ids"] = task_ids

    if peft_config.is_prompt_learning:
        if uses_cache and (model_kwargs["past_key_values"] is not None):
            # change in the logic of `prepare_inputs_for_generation` makes the below code necessary
            # In prompt learning methods, past key values are longer when compared to the `input_ids`.
            # As such only consider the last input ids in the autogressive generation phase.
            if model_kwargs.get("reuse_cache", False):
                if model_kwargs["past_key_values"][0][0][-2] >= model_kwargs["input_ids"].shape[1]:
                    model_kwargs["input_ids"] = model_kwargs["input_ids"][:, -1:]
            else:
                if model_kwargs["past_key_values"][0][0].shape[-2] >= model_kwargs["input_ids"].shape[1]:
                    model_kwargs["input_ids"] = model_kwargs["input_ids"][:, -1:]

        if model_kwargs.get("attention_mask", None) is not None:
            size = model_kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens
            prefix_attention_mask = torch.ones(size).to(model_kwargs["input_ids"].device)
            model_kwargs["attention_mask"] = torch.cat((prefix_attention_mask, model_kwargs["attention_mask"]), dim=1)

        token_idx = model_kwargs.get("token_idx", None)
        if token_idx is not None:
            token_idx = token_idx + peft_config.num_virtual_tokens

        token_idx_cpu = model_kwargs.get("token_idx_cpu", None)
        if token_idx_cpu is not None:
            token_idx_cpu = token_idx_cpu + peft_config.num_virtual_tokens
            model_kwargs["token_idx_cpu"] = token_idx_cpu

        if model_kwargs.get("position_ids", None) is not None and token_idx is None:
            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
            model_kwargs["position_ids"] = None

        if kwargs.get("token_type_ids", None) is not None:
            warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids")
            kwargs["token_type_ids"] = None

        if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING:
            past_key_values = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0])
            model_kwargs["past_key_values"] = past_key_values
        else:
            if model_kwargs["past_key_values"] is None:
                inputs_embeds = self.word_embeddings(model_kwargs["input_ids"])
                prompts = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0], task_ids=task_ids)
                prompts = prompts.to(inputs_embeds.dtype)
                model_kwargs["inputs_embeds"] = torch.cat((prompts, inputs_embeds), dim=1)
                model_kwargs["input_ids"] = None
        if token_idx is not None:
            attention_mask = model_kwargs["attention_mask"]
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if peft_config.peft_type == PeftType.PREFIX_TUNING and model_kwargs["input_ids"].shape[-1] != 1:
                position_ids = position_ids[:, -model_kwargs["input_ids"].shape[-1] :]
            if model_kwargs["past_key_values"] is not None and model_kwargs["input_ids"].shape[-1] == 1:
                position_ids = torch.index_select(position_ids, 1, token_idx - 1)
            model_kwargs["position_ids"] = position_ids
            model_kwargs["token_idx"] = token_idx
    # For transformers>=4.38.0 - for some architectures such as Llama, `cache_position` is
    # passed in the forward pass to keep track of the position ids of the cache. We have to
    # pop that from `model_kwargs` as `cache_position` is properly created by the model, using the passed
    # `inputs_embeds`: https://github.com/huggingface/transformers/blob/593230f0a1150ea9c0477b9d859f25daf73c8c33/src/transformers/models/llama/modeling_llama.py#L956
    _ = model_kwargs.pop("cache_position", None)

    return model_kwargs