Add past_key_values to forward pass for (M)BartModelWithHeads

src/transformers/models/bart/modeling_bart.py

@@ -1291,6 +1291,7 @@ def forward(
         output_hidden_states=None,
         return_dict=None,
         head=None,
+        past_key_values=None,
         **kwargs
     ):
         r"""
@@ -1318,6 +1319,7 @@ def forward(
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            past_key_values=past_key_values,
         )
         # sequence classification based on last token in sequence
         x = outputs[0]  # last hidden state

src/transformers/models/mbart/modeling_mbart.py

@@ -1291,6 +1291,7 @@ def forward(
         output_hidden_states=None,
         return_dict=None,
         head=None,
+        past_key_values=None,
         **kwargs
     ):
         r"""
@@ -1318,6 +1319,7 @@ def forward(
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            past_key_values=past_key_values,
         )
         # sequence classification based on last token in sequence
         x = outputs[0]  # last hidden state

tests/test_adapter_common.py

@@ -5,6 +5,7 @@
 
 from transformers import (
     ADAPTER_CONFIG_MAP,
+    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
     MODEL_WITH_HEADS_MAPPING,
     AdapterSetup,
     AutoModelWithHeads,
@@ -314,3 +315,37 @@ def test_loading_adapter_weights_without_prefix(self):
         output2 = model_base(**input_data)
         self.assertEqual(len(output1), len(output2))
         self.assertTrue(torch.equal(output1[0], output2[0]))
+
+    def test_forward_with_past(self):
+        if self.config_class not in MODEL_WITH_HEADS_MAPPING:
+            self.skipTest("Does not support flex heads.")
+        if self.config_class not in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING:
+            self.skipTest("No causal lm class.")
+
+        static_model = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[self.config_class](self.config())
+        flex_model = AutoModelWithHeads.from_pretrained(
+            None, config=self.config(), state_dict=static_model.state_dict()
+        )
+        static_model.add_adapter("dummy")
+        static_model.set_active_adapters("dummy")
+        static_model.eval()
+        flex_model.eval()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            static_model.save_adapter(temp_dir, "dummy")
+
+            loading_info = {}
+            flex_model.load_adapter(temp_dir, loading_info=loading_info)
+            flex_model.set_active_adapters("dummy")
+
+        input_data = self.get_input_samples((1, 128), config=static_model.config)
+        static_model.eval()
+        flex_model.eval()
+        static_model.to(torch_device)
+        flex_model.to(torch_device)
+        output = static_model(**input_data)
+
+        input_data["past_key_values"] = output["past_key_values"]
+        output_base = static_model(**input_data)
+        output_with_head = flex_model(**input_data)
+        self.assertTrue(torch.allclose(output_base["logits"], output_with_head["logits"]))

tests/test_adapter_conversion.py

@@ -141,3 +141,32 @@ def test_conversion_multiple_choice_model(self):
         label_dict = {}
         label_dict["labels"] = torch.ones(self.batch_size, dtype=torch.long, device=torch_device)
         self.run_test(model, input_shape=(self.batch_size, 2, self.seq_length), label_dict=label_dict)
+
+    def test_equivalent_language_generation(self):
+        if self.config_class not in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING:
+            self.skipTest("no causal lm class.")
+
+        static_model = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[self.config_class](self.config())
+        flex_model = AutoModelWithHeads.from_pretrained(
+            None, config=self.config(), state_dict=static_model.state_dict()
+        )
+        static_model.add_adapter("dummy")
+        static_model.set_active_adapters("dummy")
+        static_model.eval()
+        flex_model.eval()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            static_model.save_adapter(temp_dir, "dummy")
+
+            loading_info = {}
+            flex_model.load_adapter(temp_dir, loading_info=loading_info)
+            flex_model.set_active_adapters("dummy")
+
+        input_shape = (self.batch_size, 5)
+        input_samples = self.get_input_samples(input_shape, config=flex_model.config)
+
+        model_gen = static_model.generate(**input_samples)
+        flex_model_gen = flex_model.generate(**input_samples)
+
+        self.assertEquals(model_gen.shape, flex_model_gen.shape)
+        self.assertTrue(torch.equal(model_gen, flex_model_gen))