Repository: huggingface/transformers
Pull request: Share embedding modules in BART, not only weights

Embedding modules are now shared between encoder, decoder and shared - it is the same module, like in the T5 implementation.

This has the benefit that it does not matter which module is returned in `get_input_embeddings`, the caller of the latter can be sure that modifications done to that (e.g., hooks) apply to the embeddings.

Background: While revamping the gradient checkpointing tests in PEFT via peft#2860 we found that the gradient enable step
(`modeling_utils.enable_input_require_grads`) does not work for BART. This leads to gradient checkpointing with `use_reentrant=True` to fail as it will not detect any gradients. The reason for this is that the returned value by `get_input_embeddings` (`self.shared`) is not the module that is called in the encoder, therefore any hooks added to `self.shared` are not run - in this case the hook set by `enable_input_require_grads`.

Since the background is a missing hook I've added a test that tests directly the ability to define hooks and their ability to being called.

Changed files:
filename: src/transformers/models/bart/modeling_bart.py
patch: @@ -538,12 +538,12 @@ def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = No
         self.max_source_positions = config.max_position_embeddings
         embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
 
-        self.embed_tokens = BartScaledWordEmbedding(
-            config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
-        )
-
         if embed_tokens is not None:
-            self.embed_tokens.weight = embed_tokens.weight
+            self.embed_tokens = embed_tokens
+        else:
+            self.embed_tokens = BartScaledWordEmbedding(
+                config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
+            )
 
         self.embed_positions = BartLearnedPositionalEmbedding(
             config.max_position_embeddings,
@@ -682,12 +682,12 @@ def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = No
         self.max_target_positions = config.max_position_embeddings
         embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
 
-        self.embed_tokens = BartScaledWordEmbedding(
-            config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
-        )
-
         if embed_tokens is not None:
-            self.embed_tokens.weight = embed_tokens.weight
+            self.embed_tokens = embed_tokens
+        else:
+            self.embed_tokens = BartScaledWordEmbedding(
+                config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+            )
 
         self.embed_positions = BartLearnedPositionalEmbedding(
             config.max_position_embeddings,

filename: src/transformers/models/plbart/modeling_plbart.py
patch: @@ -343,12 +343,12 @@ def __init__(self, config: PLBartConfig, embed_tokens: Optional[nn.Embedding] =
         self.max_source_positions = config.max_position_embeddings
         embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
 
-        self.embed_tokens = PLBartScaledWordEmbedding(
-            config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
-        )
-
         if embed_tokens is not None:
-            self.embed_tokens.weight = embed_tokens.weight
+            self.embed_tokens = embed_tokens
+        else:
+            self.embed_tokens = PLBartScaledWordEmbedding(
+                config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
+            )
 
         self.embed_positions = PLBartLearnedPositionalEmbedding(
             config.max_position_embeddings,
@@ -595,12 +595,12 @@ def __init__(self, config: PLBartConfig, embed_tokens: Optional[nn.Embedding] =
         self.max_target_positions = config.max_position_embeddings
         embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
 
-        self.embed_tokens = PLBartScaledWordEmbedding(
-            config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
-        )
-
         if embed_tokens is not None:
-            self.embed_tokens.weight = embed_tokens.weight
+            self.embed_tokens = embed_tokens
+        else:
+            self.embed_tokens = PLBartScaledWordEmbedding(
+                config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+            )
 
         self.embed_positions = PLBartLearnedPositionalEmbedding(
             config.max_position_embeddings,

filename: tests/models/bart/test_modeling_bart.py
patch: @@ -16,6 +16,7 @@
 import copy
 import tempfile
 import unittest
+import unittest.mock
 from functools import cached_property
 
 import timeout_decorator  # noqa
@@ -477,6 +478,23 @@ def test_inputs_embeds(self):
             with torch.no_grad():
                 model(**inputs)[0]
 
+    def test_input_embeddings_support_forward_hook(self):
+        # Make sure that registering hooks on the input embeddings are indeed called
+        # in forward. This is necessary for gradient checkpointing in PEFT, see also #41821.
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+
+            hook = unittest.mock.MagicMock(return_value=None)
+            model.get_input_embeddings().register_forward_hook(hook)
+
+            inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))
+            model(**inputs)
+
+            self.assertGreater(hook.call_count, 0)
+
     @require_torch_fp16
     def test_generate_fp16(self):
         config, input_dict = self.model_tester.prepare_config_and_inputs()