Fix embedding Training Bug (#655)

Fixes a bug that was prevent embeddings from training properly and
changes the test to catch that in the future.

src/adapters/model_mixin.py

@@ -467,6 +467,7 @@ def train_adapter(self, adapter_setup: Union[list, AdapterCompositionBlock], tra
         self.set_active_adapters(adapter_setup)
         if train_embeddings:
             self.get_input_embeddings().train()
+            self.get_input_embeddings().weight.requires_grad = True
 
     def train_fusion(self, adapter_setup: Union[list, AdapterCompositionBlock], unfreeze_adapters=False):
         """Sets the model into mode for training of adapter fusion determined by a list of adapter names."""
@@ -1356,7 +1357,8 @@ def train_adapter(self, adapter_setup: Union[list, AdapterCompositionBlock], tra
             super().train_adapter(adapter_setup, train_embeddings)
         else:
             self.base_model.train_adapter(adapter_setup, train_embeddings)
-        self.freeze_embeddings()
+        if not train_embeddings:
+            self.freeze_embeddings()
 
     def train_adapter_fusion(self, adapter_setup: Union[list, AdapterCompositionBlock], unfreeze_adapters=False):
         """

tests/test_adapter_embeddings.py

@@ -96,8 +96,10 @@ def test_training_embedding(self):
             self.assertTrue(v.requires_grad, k)
 
         self.assertTrue(model.get_input_embeddings().train)
+        self.assertTrue(model.get_input_embeddings().weight.requires_grad)
 
         state_dict_pre = copy.deepcopy(model.state_dict())
+        initial_embedding = model.get_input_embeddings().weight.clone()
 
         train_dataset = self.dataset()
         training_args = TrainingArguments(
@@ -118,6 +120,10 @@ def test_training_embedding(self):
         )
         trainer.train()
 
+        trained_embedding = model.get_input_embeddings().weight.clone()
+
+        self.assertFalse(torch.equal(initial_embedding, trained_embedding))
+
         self.assertFalse(
             all(
                 torch.equal(v1, v2)
@@ -130,6 +136,10 @@ def test_training_embedding(self):
                 torch.equal(v1, v2)
                 for ((k1, v1), (k2, v2)) in zip(state_dict_pre.items(), model.state_dict().items())
                 if "test" not in k1
+                and "embedding" not in k1
+                and "embed_tokens" not in k1
+                and "shared" not in k1
+                and "wte" not in k1
             )
         )