FIX Issues with transformers 4.36 (#1252)

Adjust for different type of past_key_values when using caching.

Also: Fix some seeds for flaky tests.

---------

Co-authored-by: Younes Belkada <[email protected]>
Co-authored-by: Joao Gante <[email protected]>

Author: 3 people

src/peft/peft_model.py

@@ -23,7 +23,9 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Union
 
+import packaging.version
 import torch
+import transformers
 from accelerate import dispatch_model, infer_auto_device_map
 from accelerate.hooks import AlignDevicesHook, add_hook_to_module, remove_hook_from_submodules
 from accelerate.utils import get_balanced_memory
@@ -1136,11 +1138,26 @@ def generate(self, **kwargs):
     def prepare_inputs_for_generation(self, *args, task_ids: torch.Tensor = None, **kwargs):
         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.37, all architectures should support caching.
+        uses_transformers_4_37 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.37.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_37 or (
+            uses_transformers_4_36 and self.base_model.config.model_type in transformers_new_cache_archs
+        )
+
         if peft_config.is_prompt_learning:
             if model_kwargs.get("attention_mask", None) is not None:
-                prefix_attention_mask = torch.ones(
-                    model_kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens
-                ).to(model_kwargs["input_ids"].device)
+                if uses_cache and (model_kwargs["past_key_values"] is not None):
+                    # TODO figure out why this workaround is necessary, see #1252 for context
+                    size = model_kwargs["input_ids"].shape[0], model_kwargs["past_key_values"][0][0].shape[-2]
+                else:
+                    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
                 )

src/peft/tuners/adaption_prompt/utils.py

@@ -73,7 +73,12 @@ def llama_compute_query_states(model: nn.Module, **kwargs) -> torch.Tensor:
 
     seq_len = q_len
     if past_key_value is not None:
-        seq_len += past_key_value[0].shape[-2]
+        if isinstance(past_key_value, tuple):
+            # for transformers <= 4.35
+            seq_len += past_key_value[0].shape[-2]
+        else:
+            # since transformers 4.36, this is a DynamicCache instance
+            seq_len += past_key_value.get_seq_length(model.layer_idx)
     cos, sin = model.rotary_emb(value_states, seq_len=seq_len)
 
     return llama_apply_rotary_pos_emb(query_states, cos, sin, position_ids)

tests/test_adaption_prompt.py

@@ -387,6 +387,7 @@ def test_add_and_set_while_disabled(self):
 
     def test_use_cache(self) -> None:
         """Test that AdaptionPrompt works when Llama config use_cache=True."""
+        torch.manual_seed(0)
         input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
         original = LlamaForCausalLM(
             LlamaConfig(
@@ -397,7 +398,7 @@ def test_use_cache(self) -> None:
                 num_attention_heads=4,
                 use_cache=False,
             )
-        )
+        ).eval()
         adapted = get_peft_model(
             original, AdaptionPromptConfig(adapter_layers=2, adapter_len=4, task_type="CAUSAL_LM")
         )

tests/test_multitask_prompt_tuning.py

@@ -220,10 +220,11 @@ def test_generate(self) -> None:
 
     def test_use_cache(self) -> None:
         """Test that MultiTaskPromptTuning works when Llama config use_cache=True."""
+        torch.manual_seed(0)
         input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
         task_ids = torch.LongTensor([1, 2]).to(self.torch_device)
 
-        original = LlamaForCausalLM(self._create_test_llama_config())
+        original = LlamaForCausalLM(self._create_test_llama_config()).eval()
         mpt = get_peft_model(original, self._create_multitask_prompt_tuning_config())
         mpt = mpt.to(self.torch_device)