Fix DeBERTa prefix tuning w. enabled relative attention

.github/workflows/tests_torch.yml

@@ -56,7 +56,7 @@ jobs:
             ${{ runner.os }}-pip-
       - name: Install
         run: |
-          pip install torch
+          pip install torch==1.12.1
           pip install .[sklearn,testing,sentencepiece,vision]
           pip install datasets
       - name: Test

src/transformers/models/deberta/modeling_deberta.py

@@ -655,6 +655,7 @@ def linear(w, b, x):
         query_layer = query_layer + self.transpose_for_scores(self.q_bias[None, None, :])
         value_layer = value_layer + self.transpose_for_scores(self.v_bias[None, None, :])
 
+        orig_key_layer = key_layer  # save this for relative attention
         key_layer, value_layer, attention_mask = self.prefix_tuning(
             key_layer, value_layer, hidden_states, attention_mask, False
         )
@@ -667,10 +668,14 @@ def linear(w, b, x):
         attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
         if self.relative_attention:
             rel_embeddings = self.pos_dropout(rel_embeddings)
-            rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
+            rel_att = self.disentangled_att_bias(
+                query_layer, orig_key_layer, relative_pos, rel_embeddings, scale_factor
+            )
 
         if rel_att is not None:
-            attention_scores = attention_scores + rel_att
+            rel_att_padded = torch.zeros_like(attention_scores)
+            rel_att_padded[:, :, :, -rel_att.size(-1) :] = rel_att
+            attention_scores = attention_scores + rel_att_padded
 
         # bxhxlxd
         if self.talking_head:

src/transformers/models/deberta_v2/modeling_deberta_v2.py

@@ -713,6 +713,9 @@ def forward(
         key_layer = self.transpose_for_scores_extended(self.key_proj(hidden_states), self.num_attention_heads)
         value_layer = self.transpose_for_scores_extended(self.value_proj(hidden_states), self.num_attention_heads)
 
+        orig_key_layer = key_layer.contiguous().view(
+            -1, key_layer.size(2), key_layer.size(-1)
+        )  # save this for relative attention
         key_layer, value_layer, attention_mask = self.prefix_tuning(
             key_layer, value_layer, hidden_states, attention_mask, False
         )  # [:, 0, :, 0])
@@ -732,11 +735,13 @@ def forward(
         if self.relative_attention:
             rel_embeddings = self.pos_dropout(rel_embeddings)
             rel_att = self.disentangled_attention_bias(
-                query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+                query_layer, orig_key_layer, relative_pos, rel_embeddings, scale_factor
             )
 
         if rel_att is not None:
-            attention_scores = attention_scores + rel_att
+            rel_att_padded = torch.zeros_like(attention_scores)
+            rel_att_padded[:, :, -rel_att.size(2) :] = rel_att
+            attention_scores = attention_scores + rel_att_padded
         attention_scores = attention_scores
         attention_scores = attention_scores.view(
             -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)

tests_adapters/test_deberta.py

@@ -4,7 +4,14 @@
 from transformers import DebertaAdapterModel
 from transformers.testing_utils import require_torch
 
-from .methods import BottleneckAdapterTestMixin, UniPELTTestMixin, CompacterTestMixin, IA3TestMixin, LoRATestMixin, PrefixTuningTestMixin
+from .methods import (
+    BottleneckAdapterTestMixin,
+    CompacterTestMixin,
+    IA3TestMixin,
+    LoRATestMixin,
+    PrefixTuningTestMixin,
+    UniPELTTestMixin,
+)
 from .test_adapter import AdapterTestBase, make_config
 from .test_adapter_backward_compability import CompabilityTestMixin
 from .test_adapter_composition import ParallelAdapterInferenceTestMixin, ParallelTrainingMixin
@@ -32,6 +39,8 @@ class DebertaAdapterTestBase(AdapterTestBase):
         num_attention_heads=4,
         intermediate_size=37,
         hidden_act="gelu",
+        relative_attention=True,
+        pos_att_type="p2c|c2p",
     )
     tokenizer_name = "microsoft/deberta-base"
 

tests_adapters/test_debertaV2.py

@@ -4,7 +4,14 @@
 from transformers import DebertaV2AdapterModel
 from transformers.testing_utils import require_torch
 
-from .methods import BottleneckAdapterTestMixin, UniPELTTestMixin, CompacterTestMixin, IA3TestMixin, LoRATestMixin, PrefixTuningTestMixin
+from .methods import (
+    BottleneckAdapterTestMixin,
+    CompacterTestMixin,
+    IA3TestMixin,
+    LoRATestMixin,
+    PrefixTuningTestMixin,
+    UniPELTTestMixin,
+)
 from .test_adapter import AdapterTestBase, make_config
 from .test_adapter_backward_compability import CompabilityTestMixin
 from .test_adapter_composition import ParallelAdapterInferenceTestMixin, ParallelTrainingMixin
@@ -32,6 +39,8 @@ class DebertaV2AdapterTestBase(AdapterTestBase):
         num_attention_heads=4,
         intermediate_size=37,
         hidden_act="gelu",
+        relative_attention=True,
+        pos_att_type="p2c|c2p",
     )
     tokenizer_name = "microsoft/deberta-v3-base"