Repository: huggingface/peft

Issue:
LLaMA Adapter Fails during the inference due to mixed precision weights.
Running the HF LLaMA model with LLaMa Adapter, it works fine during the training. However to load back the PEFT checkpoints and run inference it fails and complains about 

``` bash
File "/home/ubuntu/peft/src/peft/tuners/adaption_prompt.py", line 347, in forward
    scores = torch.matmul(query_states, adapter_k.transpose(2, 3)) / math.sqrt(self.model.head_dim)
RuntimeError: expected scalar type Float but found Half
```

[Repro code ](https://gist.github.com/HamidShojanazeri/d6a8069cbc40cf52a54f6352d0d305c2)

[Error Logs](https://gist.github.com/HamidShojanazeri/7eaae28db4ec2a8a6799243a7334ac9f)

The model weights are all in fp16 and LLaMa Adapter layers in fp32.

I would appreciate any work around/fix.


Discussion:
cc: @younesbelkada , @pacman100 
thanks @younesbelkada for the pointer, seems like change the following in [LLaMa Adapter](https://github.com/huggingface/peft/blob/main/src/peft/tuners/adaption_prompt.py#L347) can fix the issue.

```
previous_dtype = query_states.dtype
query_states = query_states.to(torch.float32)
```
and finally before  the return
```
output.to(previous_dtype)

```


Resolving patch:
diff --git a/src/peft/tuners/adaption_prompt.py b/src/peft/tuners/adaption_prompt.py
--- a/src/peft/tuners/adaption_prompt.py
+++ b/src/peft/tuners/adaption_prompt.py
@@ -294,10 +294,12 @@ def __init__(self, model_type: str, adapter_len: int, model):
         # https://github.com/ZrrSkywalker/LLaMA-Adapter/blob/41c3546fe1997ab8a65809dc8d8f9252b19d9faf/llama/model.py#L234
         # (bsz, adapter_len, hidden_size)
         self.adaption_prompt = nn.Parameter(
-            torch.empty(1, adapter_len, self.model.hidden_size, device=device).normal_()
+            torch.empty(
+                1, adapter_len, self.model.hidden_size, device=device, dtype=model.q_proj.weight.dtype
+            ).normal_()
         )
         # Initialize the gate to 0 as this is "zero-init".
-        self.adaption_gate = nn.Parameter(torch.zeros(1, device=device))
+        self.adaption_gate = nn.Parameter(torch.zeros(1, device=device, dtype=model.q_proj.weight.dtype))
 
     def forward(self, **kwargs):
         """
@@ -343,11 +345,12 @@ def forward(self, **kwargs):
         # (bsz, num_heads, q_len, head_dim)
         query_states = compute_query_states(model=self.model, **kwargs)
 
+        previous_dtype = query_states.dtype
         # (bsz, num_heads, q_len, adapter_len)
         scores = torch.matmul(query_states, adapter_k.transpose(2, 3)) / math.sqrt(self.model.head_dim)
         # Upcast attention to fp32
         # (bsz, num_heads, q_len, adapter_len)
-        scores = self.adaption_gate * F.softmax(scores, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        scores = self.adaption_gate * F.softmax(scores, dim=-1, dtype=torch.float32).to(previous_dtype)
         # (bsz, q_len, num_heads * head_dim)
         adapter_output = torch.matmul(scores, adapter_v).transpose(1, 2).reshape(bsz, q_len, -1)
         # (bsz, q_len, hidden_size)
@@ -356,4 +359,7 @@ def forward(self, **kwargs):
 
         # Add adaption prompt output to original output.
         output = output + adapter_output
+
+        # Restore original dtype.
+        output = output.to(previous_dtype)
         return output, None, past_key_value


Test patch:
diff --git a/tests/test_adaption_prompt.py b/tests/test_adaption_prompt.py
--- a/tests/test_adaption_prompt.py
+++ b/tests/test_adaption_prompt.py
@@ -312,3 +312,15 @@ def test_use_cache(self) -> None:
         adapted.base_model.config.use_cache = True
         actual = adapted.generate(input_ids=input_ids, max_length=8)
         assert_close(expected, actual, rtol=0, atol=0)
+
+    def test_bf16_inference(self) -> None:
+        """Test that AdaptionPrompt works when Llama using a half-precision model."""
+        input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
+        original = LlamaForCausalLM.from_pretrained(
+            "trl-internal-testing/tiny-random-LlamaForCausalLM", torch_dtype=torch.bfloat16
+        )
+        adapted = get_peft_model(
+            original, AdaptionPromptConfig(adapter_layers=2, adapter_len=4, task_type="CAUSAL_LM")
+        )
+        adapted = adapted.to(self.torch_device)
+        _ = adapted.generate(input_ids=input_ids)