Prefix Tuning does not work with T5-3B

[saparina]

Environment info

• adapter-transformers version: 3.1.0

Information

Model I am using: T5-3B

Adapter setup I am using (if any): Prefix Tuning

To reproduce

I believe this bug occurs every time T5-3B and Prefix Tuning are used, for example, it can be reproduced by running the summarization example:

python examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path t5-small \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --train_adapter \
    --adapter_config prefix_tuning

Error is raised when prefix embeddings should be concatenated with keys and values:

  File "/home/hpcsapa1/.conda/envs/adapters-torch1.9/lib/python3.9/site-packages/transformers/adapters/prefix_tuning.py", line 329, in forward
    key_states = torch.cat([prefix_keys, key_states], dim=2)
RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 32 but got size 128 for tensor number 1 in the list.

Expected behavior

There is a mismatch in the sizes of prefix embeddings and keys and values. I think this is because the prefix size is based on hidden_size which is defined as d_model for T5 models while it should be num_heads * d_kv. For T5-Small, T5-Base and T5-Large d_model == num_heads * d_kv (and these models work fine) but this is not true for T5-3B where d_model=1024 and num_heads=32, d_kv=128 .

[calpt]

Hey @saparina, thanks for bringing this up.

Your explanation of the issue makes a lot of sense, we'll look into a posible fix for prefix tuning. It's a bit odd however that the constraint d_model == num_heads * d_kv is not fulfilled by all models as the HuggingFace documentation for T5 explicitly states that this is required to be the case: https://huggingface.co/docs/transformers/main/en/model_doc/t5#transformers.T5Config.d_kv

[adapter-hub-bert]

This issue has been automatically marked as stale because it has been without activity for 90 days. This issue will be closed in 14 days unless you comment or remove the stale label.

[vijetadeshpande]

Is there any update on this issue?

[calpt]

Sorry for the delay on this, it should be fixed once #621 is merged.