multi-lora setup with llama

[simeneide]

Environment info

• transformers version: 4.43.4
• Platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.35
• Python version: 3.12.4
• Huggingface_hub version: 0.23.4
• Safetensors version: 0.4.3
• Accelerate version: 0.34.2
• Accelerate config: not found
• PyTorch version (GPU?): 2.4.0+cu121 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: no
• GPU type: NVIDIA A100-SXM4-80GB
• adapters version: 1.0.0

Information

I am trying to set up a simple ac.Parallel on two lora modules. However, running an slightly modified example gives me RuntimeError: shape '[2, 6, 32, 128]' is invalid for input of size 98304

Adapter setup I am using (if any):

The problem arises when using:

• the official example scripts: (give details below)
• my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: (give the name)
• my own task or dataset: (give details below)

To reproduce

import adapters.composition as ac
from adapters import AutoAdapterModel
from transformers import AutoTokenizer
import adapters
import torch
model_name="meta-llama/Llama-3.2-1B"
model_name = "meta-llama/Llama-3.1-8B"
model = AutoAdapterModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

input_ids = tokenizer(["Adapters are great!", "Adapters are awesome!"], return_tensors="pt")

from adapters import LoRAConfig
config = LoRAConfig(r=8, alpha=16)
model.add_adapter("a", config=config)
model.add_adapter("b", config=config)

model.active_adapters = ac.Parallel("a", "b")

output1, output2 = model(**input_ids)

Error message:

[320](~/miniforge3/envs/schiblm/lib/python3.12/site-packages/adapters/models/llama/modeling_llama.py:320) key_states = self.k_proj(hidden_states)
    [321](~/miniforge3/envs/schiblm/lib/python3.12/site-packages/adapters/models/llama/modeling_llama.py:321) value_states = self.v_proj(hidden_states)
--> [323](~/miniforge3/envs/schiblm/lib/python3.12/site-packages/adapters/models/llama/modeling_llama.py:323) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    [324](~/miniforge3/envs/schiblm/lib/python3.12/site-packages/adapters/models/llama/modeling_llama.py:324) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
    [325](~/miniforge3/envs/schiblm/lib/python3.12/site-packages/adapters/models/llama/modeling_llama.py:325) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

RuntimeError: shape '[2, 6, 32, 128]' is invalid for input of size 98304

Expected behavior

It worked using the bert example from the documentation, so something with the lora or the autoregressive model here maybe?