Parallel training

src/transformers/adapters/heads/base.py

@@ -1,4 +1,5 @@
 import logging
+from dataclasses import dataclass
 from typing import List, Optional, Union
 
 import torch
@@ -23,6 +24,35 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class MultiHeadOutput(ModelOutput):
+ head_outputs: List[ModelOutput] = None
+ loss: Optional[torch.FloatTensor] = None
+
+ def __getitem__(self, k):
+ # with number indices the head output at that position is accessed
+ # e.g output[1] is equivalent to output.head_outputs[1]
+ if isinstance(k, int):
+ return self.head_outputs[k]
+ # with strings the attribute in the underlying dict can be adressed
+ # e.g output["loss"] is equivalent to output.loss
+ else:
+ return super().__getitem__(k)
+
+ def __setitem__(self, k, v):
+ if isinstance(k, int):
+ self.head_outputs[k] = v
+ else:
+ return super().__setitem__(k, v)
+
+ def __iter__(self):
+ # iterates over the head outputs
+ return iter(self.head_outputs)
+
+ def __len__(self):
+ return len(self.head_outputs)
+
+
 # Let this class inherit from nn.Sequential to provide iterable access as before
 class PredictionHead(nn.Sequential):
  def __init__(self, name):
@@ -662,14 +692,21 @@ def _get_head_input(outputs, cls_out, batch):
  )
  )
  head_outputs = []
+ labels = kwargs.pop("labels", None)
  for i, head in enumerate(self.active_head):
  head_module = self.heads[head]
  batch_idx = range(sum(self.active_head.batch_sizes[:i]), sum(self.active_head.batch_sizes[: i + 1]))
+ kwargs["labels"] = labels[batch_idx] if labels is not None else None
  head_inputs, head_cls_input = _get_head_input(all_outputs, cls_output, batch_idx)
  # head_attention = attention_mask[batch_idx] if attention_mask is not None else None
  head_output = head_module(head_inputs, head_cls_input, attention_mask, return_dict, **kwargs)
  head_outputs.append(head_output)
- return head_outputs
+ combined_loss = (
+ sum([out["loss"] for out in head_outputs])
+ if all("loss" in out and out["loss"] is not None for out in head_outputs)
+ else None
+ )
+ return MultiHeadOutput(head_outputs=head_outputs, loss=combined_loss)
  elif self.has_parallel_adapters or isinstance(self.active_head, Parallel):
  if len(self.active_head) != self.config.adapters.active_setup.parallel_channels:
  raise ValueError("The number of parallel adapters and the number of active heads must match.")
@@ -681,7 +718,12 @@ def _get_head_input(outputs, cls_out, batch):
  head_inputs, head_cls_input = _get_head_input(all_outputs, cls_output, batch_idx)
  head_output = head_module(head_inputs, head_cls_input, attention_mask, return_dict, **kwargs)
  head_outputs.append(head_output)
- return head_outputs
+ combined_loss = (
+ torch.sum(torch.stack([out["loss"] for out in head_outputs]))
+ if all("loss" in out and out["loss"] is not None for out in head_outputs)
+ else None
+ )
+ return MultiHeadOutput(head_outputs=head_outputs, loss=combined_loss)
  elif len(used_heads) > 1:
  head_outputs = []
  for head in used_heads:

tests/test_adapter.py

@@ -18,7 +18,7 @@
 from transformers.testing_utils import require_torch, torch_device
 
 from .test_adapter_common import AdapterModelTestMixin
-from .test_adapter_composition import ParallelAdapterInferenceTestMixin
+from .test_adapter_composition import ParallelAdapterInferenceTestMixin, ParallelTrainingMixin
 from .test_adapter_conversion import ModelClassConversionTestMixin
 from .test_adapter_fusion_common import AdapterFusionModelTestMixin
 from .test_adapter_heads import PredictionHeadModelTestMixin
@@ -78,6 +78,7 @@ class BertAdapterTest(
  PredictionHeadModelTestMixin,
  AdapterTrainingTestMixin,
  ParallelAdapterInferenceTestMixin,
+ ParallelTrainingMixin,
  BertAdapterTestBase,
  unittest.TestCase,
 ):
@@ -144,6 +145,7 @@ class DistilBertAdapterTest(
  PredictionHeadModelTestMixin,
  AdapterTrainingTestMixin,
  ParallelAdapterInferenceTestMixin,
+ ParallelTrainingMixin,
  DistilBertAdapterTestBase,
  unittest.TestCase,
 ):
@@ -181,6 +183,7 @@ class BartAdapterTest(
  PredictionHeadModelTestMixin,
  AdapterTrainingTestMixin,
  ParallelAdapterInferenceTestMixin,
+ ParallelTrainingMixin,
  BartAdapterTestBase,
  unittest.TestCase,
 ):
@@ -251,6 +254,7 @@ class GPT2AdapterTest(
  PredictionHeadModelTestMixin,
  AdapterTrainingTestMixin,
  ParallelAdapterInferenceTestMixin,
+ ParallelTrainingMixin,
  GPT2AdapterTestBase,
  unittest.TestCase,
 ):

tests/test_adapter_composition.py

@@ -1,8 +1,20 @@
+import copy
+import random
 import unittest
 
 import torch
 
-from transformers import AutoModelWithHeads, BertConfig, BertForSequenceClassification
+from tests.test_adapter_training import filter_parameters
+from transformers import (
+ AutoModelWithHeads,
+ AutoTokenizer,
+ BertConfig,
+ BertForSequenceClassification,
+ GlueDataset,
+ GlueDataTrainingArguments,
+ Trainer,
+ TrainingArguments,
+)
 from transformers.adapters.composition import BatchSplit, Fuse, Parallel, Split, Stack, parse_composition
 from transformers.testing_utils import require_torch, torch_device
 
@@ -207,3 +219,156 @@ def test_batch_split_with_heads(self):
  self.assertEqual(2, len(output))
  self.assertTrue(torch.allclose(output[0]["logits"], outputs_a["logits"]))
  self.assertTrue(torch.allclose(output[1]["logits"], outputs_b["logits"]))
+
+
+def create_twin_adapters(model, name):
+ # create adapter
+ adapter1, adapter2 = name + "_1", name + "_2"
+ model.add_adapter(adapter1)
+ model.add_classification_head(adapter1)
+ # create a twin initialized with the same random weights
+ model.add_adapter(adapter2)
+ model.add_classification_head(adapter2)
+
+ state_dict = model.state_dict()
+ for k, v in state_dict.items():
+ if adapter1 in k:
+ state_dict[k.replace(adapter1, adapter2)] = v
+ model.load_state_dict(state_dict)
+
+ return adapter1, adapter2
+
+
+class ParallelTrainingMixin:
+ def train_model(self, model, dataset):
+ # trains model in eval mode for 2 epochs
+ random.seed(42)
+ torch.manual_seed(42)
+ # Depending on the used optimizer the adapters are not exactly the same
+ optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+ for epoch in range(2):
+ for data_input in dataset:
+ optimizer.zero_grad()
+ output = model(**data_input)
+ loss = output["loss"]
+ loss.backward()
+ optimizer.step()
+ return model
+
+ def test_parallel_training(self):
+ tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, use_fast=False)
+ if tokenizer.pad_token is None:
+ tokenizer.pad_token = tokenizer.eos_token
+ model = AutoModelWithHeads.from_config(self.config())
+
+ model.add_adapter("mrpc1")
+ model.add_adapter("mrpc2")
+ model.add_classification_head("mrpc1", num_labels=2)
+ model.add_classification_head("mrpc2", num_labels=3)
+ model.eval()
+ model.active_adapters = Parallel("mrpc1", "mrpc2")
+ model.train_adapter(Parallel("mrpc1", "mrpc2"))
+
+ # all weights of the adapter should be activated
+ for k, v in filter_parameters(model, "adapters.mrpc1.").items():
+ self.assertTrue(v.requires_grad, k)
+ # all weights of the adapter not used for training should be freezed
+ for k, v in filter_parameters(model, "adapters.mrpc2.").items():
+ self.assertTrue(v.requires_grad, k)
+ # weights of the model should be freezed (check on some examples)
+ for k, v in filter_parameters(model, "encoder.layer.0.attention").items():
+ self.assertFalse(v.requires_grad, k)
+
+ state_dict_pre = copy.deepcopy(model.state_dict())
+
+ # setup dataset
+ data_args = GlueDataTrainingArguments(
+ task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True
+ )
+ train_dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="train")
+ training_args = TrainingArguments(
+ output_dir="./examples", do_train=True, learning_rate=0.1, max_steps=7, no_cuda=True
+ )
+
+ # evaluate
+ trainer = Trainer(
+ model=model,
+ args=training_args,
+ train_dataset=train_dataset,
+ )
+ trainer.train()
+
+ for ((k1, v1), (k2, v2)) in zip(state_dict_pre.items(), model.state_dict().items()):
+ if "mrpc" in k1:
+ self.assertFalse(torch.equal(v1, v2))
+ else:
+ self.assertTrue(torch.equal(v1, v2))
+
+ def test_parallel_training_equivalent_to_single_adapters(self):
+ model = AutoModelWithHeads.from_config(self.config())
+ model.eval()
+
+ a1, a2 = create_twin_adapters(model, "a")
+ b1, b2 = create_twin_adapters(model, "b")
+
+ dataset = []
+ for i in range(3):
+ input_data = self.get_input_samples((3, 128), config=model.config)
+ input_data["labels"] = torch.randint(0, 2, (3, 1))
+ dataset.append(input_data)
+
+ for adapter in [a1, b1]:
+ model.active_head = adapter
+ model.set_active_adapters(adapter)
+ model.train_adapter(adapter)
+ model.eval()
+
+ model = self.train_model(model, dataset)
+
+ model.set_active_adapters(Parallel(a2, b2))
+ model.train_adapter((Parallel(a2, b2)))
+ model.eval()
+
+ model = self.train_model(model, dataset)
+
+ state_dict = model.state_dict()
+ for k, v in state_dict.items():
+ if a1 in k:
+ self.assertTrue(torch.allclose(v, state_dict[k.replace(a1, a2)]))
+ if b1 in k:
+ self.assertTrue(torch.allclose(v, state_dict[k.replace(b1, b2)]))
+
+ def test_parallel_training_single_forward_pass(self):
+ model = AutoModelWithHeads.from_config(self.config())
+ model.eval()
+
+ a1, a2 = create_twin_adapters(model, "a")
+ b1, b2 = create_twin_adapters(model, "b")
+
+ state_dict = model.state_dict()
+ for k, v in state_dict.items():
+ if a1 in k:
+ self.assertTrue(torch.equal(v, state_dict[k.replace(a1, a2)]))
+ if b1 in k:
+ self.assertTrue(torch.equal(v, state_dict[k.replace(b1, b2)]))
+
+ input_data = self.get_input_samples((3, 128), config=model.config)
+ input_data["labels"] = torch.randint(0, 2, (3, 1))
+
+ outputs = []
+ for adapter in [a1, b1]:
+ model.active_head = adapter
+ model.set_active_adapters(adapter)
+ model.train_adapter(adapter)
+ model.eval()
+ outputs.append(model(**input_data))
+
+ model.set_active_adapters(Parallel(a2, b2))
+ model.train_adapter((Parallel(a2, b2)))
+ model.eval()
+
+ parallel_outputs = model(**input_data)
+
+ for out1, out2 in zip(outputs, parallel_outputs.head_outputs):
+ self.assertTrue(torch.equal(out1["loss"], out2["loss"]))
+ self.assertTrue(torch.allclose(out1["logits"], out2["logits"]))