Repository: huggingface/transformers (Python) — Bug Fix

Issue:
Multiple validation datasets unsupported with `dataloader_persistent_workers=True` 
### System Info

- `transformers` version: 4.40.1
- Platform: Linux-6.8.0-76060800daily20240311-generic-x86_64-with-glibc2.35
- Python version: 3.11.8
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.3
- Accelerate version: 0.29.3
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: False

### Who can help?

@muellerzr @pacman100

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from transformers import Trainer, TrainingArguments

DIM = 2


class DummyDataset(Dataset):
    def __init__(self, size=10000, label=0):
        self.size = size
        self.data = torch.rand(size, DIM)
        self.labels = torch.full((size,), label)

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return {"input_ids": self.data[idx], "labels": self.labels[idx]}


class DummyModel(torch.nn.Module):
    def __init__(self):
        super(DummyModel, self).__init__()
        self.linear = torch.nn.Linear(DIM, 2)

    def forward(self, input_ids, labels=None):
        outputs = self.linear(input_ids)
        loss = F.cross_entropy(outputs, labels)
        return {"logits": outputs, "loss": loss}


if __name__ == "__main__":
    model = DummyModel()
    train_dataset = DummyDataset(label=0)
    good_validation_dataset = DummyDataset(label=0)
    bad_validation_dataset = DummyDataset(label=1)

    training_args = TrainingArguments(
        output_dir="./outputs",
        learning_rate=0.01,
        num_train_epochs=5,
        per_device_train_batch_size=128,
        per_device_eval_batch_size=128,
        dataloader_num_workers=2,
        dataloader_persistent_workers=True,
        evaluation_strategy="epoch",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset={"good": good_validation_dataset, "bad": bad_validation_dataset},
    )

    trainer.train()
```
With `dataloader_persistent_workers=True` :
```bash
{'eval_good_loss': 0.04770788177847862, 'eval_good_runtime': 0.0951, 'eval_good_samples_per_second': 105140.955, 'eval_good_steps_per_second': 830.614, 'epoch': 1.0}                                                                                                             
{'eval_bad_loss': 0.04770788177847862, 'eval_bad_runtime': 0.1225, 'eval_bad_samples_per_second': 81619.03, 'eval_bad_steps_per_second': 644.79, 'epoch': 1.0}                                                                                                                    
{'eval_good_loss': 0.024791115894913673, 'eval_good_runtime': 0.0995, 'eval_good_samples_per_second': 100488.125, 'eval_good_steps_per_second': 793.856, 'epoch': 2.0}                                                                                                            
{'eval_bad_loss': 0.024791115894913673, 'eval_bad_runtime': 0.1183, 'eval_bad_samples_per_second': 84530.882, 'eval_bad_steps_per_second': 667.794, 'epoch': 2.0}                                                                                                                 
{'eval_good_loss': 0.017540939152240753, 'eval_good_runtime': 0.095, 'eval_good_samples_per_second': 105282.943, 'eval_good_steps_per_second': 831.735, 'epoch': 3.0}                                                                                                             
{'eval_bad_loss': 0.017540939152240753, 'eval_bad_runtime': 0.0814, 'eval_bad_samples_per_second': 122839.094, 'eval_bad_steps_per_second': 970.429, 'epoch': 3.0}                                                                                                                
{'eval_good_loss': 0.014589476399123669, 'eval_good_runtime': 0.1745, 'eval_good_samples_per_second': 57297.904, 'eval_good_steps_per_second': 452.653, 'epoch': 4.0}                                                                                                             
{'eval_bad_loss': 0.014589476399123669, 'eval_bad_runtime': 0.1389, 'eval_bad_samples_per_second': 71998.668, 'eval_bad_steps_per_second': 568.789, 'epoch': 4.0}                                                                                                                 
{'eval_good_loss': 0.01373046450316906, 'eval_good_runtime': 0.0833, 'eval_good_samples_per_second': 120031.709, 'eval_good_steps_per_second': 948.25, 'epoch': 5.0}                                                                                                              
{'eval_bad_loss': 0.01373046450316906, 'eval_bad_runtime': 0.0865, 'eval_bad_samples_per_second': 115601.295, 'eval_bad_steps_per_second': 913.25, 'epoch': 5.0}                                                                                                                  
{'train_runtime': 1.8571, 'train_samples_per_second': 26923.771, 'train_steps_per_second': 212.698, 'train_loss': 0.03968705527390106, 'epoch': 5.0}
```
With `dataloader_persistent_workers=False` :
```bash
{'eval_good_loss': 0.10046054422855377, 'eval_good_runtime': 0.1053, 'eval_good_samples_per_second': 95006.818, 'eval_good_steps_per_second': 750.554, 'epoch': 1.0}                                                                                                              
{'eval_bad_loss': 2.533043622970581, 'eval_bad_runtime': 0.0946, 'eval_bad_samples_per_second': 105667.808, 'eval_bad_steps_per_second': 834.776, 'epoch': 1.0}                                                                                                                   
{'eval_good_loss': 0.05101846158504486, 'eval_good_runtime': 0.161, 'eval_good_samples_per_second': 62102.692, 'eval_good_steps_per_second': 490.611, 'epoch': 2.0}                                                                                                               
{'eval_bad_loss': 3.2872579097747803, 'eval_bad_runtime': 0.1805, 'eval_bad_samples_per_second': 55403.336, 'eval_bad_steps_per_second': 437.686, 'epoch': 2.0}                                                                                                                   
{'eval_good_loss': 0.03576516732573509, 'eval_good_runtime': 0.1225, 'eval_good_samples_per_second': 81623.001, 'eval_good_steps_per_second': 644.822, 'epoch': 3.0}                                                                                                              
{'eval_bad_loss': 3.694115161895752, 'eval_bad_runtime': 0.1046, 'eval_bad_samples_per_second': 95635.471, 'eval_bad_steps_per_second': 755.52, 'epoch': 3.0}                                                                                                                     
{'eval_good_loss': 0.029605071991682053, 'eval_good_runtime': 0.0998, 'eval_good_samples_per_second': 100165.593, 'eval_good_steps_per_second': 791.308, 'epoch': 4.0}                                                                                                            
{'eval_bad_loss': 3.9129879474639893, 'eval_bad_runtime': 0.0825, 'eval_bad_samples_per_second': 121274.534, 'eval_bad_steps_per_second': 958.069, 'epoch': 4.0}                                                                                                                  
{'eval_good_loss': 0.027824044227600098, 'eval_good_runtime': 0.0903, 'eval_good_samples_per_second': 110771.994, 'eval_good_steps_per_second': 875.099, 'epoch': 5.0}                                                                                                            
{'eval_bad_loss': 3.9852359294891357, 'eval_bad_runtime': 0.1141, 'eval_bad_samples_per_second': 87625.956, 'eval_bad_steps_per_second': 692.245, 'epoch': 5.0}                                                                                                                   
{'train_runtime': 2.0821, 'train_samples_per_second': 24014.737, 'train_steps_per_second': 189.716, 'train_loss': 0.08233800960492484, 'epoch': 5.0}
```

### Expected behavior

Hi there,

When using multiple validation datasets with `transformers.Trainer` and setting `dataloader_persistent_workers=True` in the `transformers.TrainingArguments`,  all evaluations are done using the first validation dataset.

In the example above, the model only learns to predict the class `0`, so we should have a big loss for the "bad" validation dataset and a small one for the "good" one.

This seems related to #28469 and #29538; which does not support passing a dictionary of evaluation datasets :
```python
# def get_eval_dataloader in src/transformers/trainer.py
if hasattr(self, "_eval_dataloader") and self.args.dataloader_persistent_workers:
            return self.accelerator.prepare(self._eval_dataloader)
```
The evaluation dataloaders should probably also be stored in a dictionary, or the `_eval_dataloader` attribute should be suffixed with the `eval_dataset_name`.

I can look into opening a PR for this.


Discussion:
@bastienlc feel free to open a PR to support this!
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.
This issue is still relevant and is being addressed in #30627.

@muellerzr

Resolving patch:
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -919,25 +919,36 @@ def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.
         else:
             return None
 
-    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
+    def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader:
         """
         Returns the evaluation [`~torch.utils.data.DataLoader`].
 
         Subclass and override this method if you want to inject some custom behavior.
 
         Args:
-            eval_dataset (`torch.utils.data.Dataset`, *optional*):
-                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted
-                by the `model.forward()` method are automatically removed. It must implement `__len__`.
+            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
+                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
         """
         if eval_dataset is None and self.eval_dataset is None:
             raise ValueError("Trainer: evaluation requires an eval_dataset.")
 
         # If we have persistent workers, don't do a fork bomb especially as eval datasets
         # don't change during training
-        if hasattr(self, "_eval_dataloader") and self.args.dataloader_persistent_workers:
-            return self.accelerator.prepare(self._eval_dataloader)
-        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
+        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
+        if (
+            hasattr(self, "_eval_dataloaders")
+            and dataloader_key in self._eval_dataloaders
+            and self.args.dataloader_persistent_workers
+        ):
+            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])
+
+        eval_dataset = (
+            self.eval_dataset[eval_dataset]
+            if isinstance(eval_dataset, str)
+            else eval_dataset
+            if eval_dataset is not None
+            else self.eval_dataset
+        )
         data_collator = self.data_collator
 
         if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
@@ -962,7 +973,10 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
         # we need to store the non-prepared version
         eval_dataloader = DataLoader(eval_dataset, **dataloader_params)
         if self.args.dataloader_persistent_workers:
-            self._eval_dataloader = eval_dataloader
+            if hasattr(self, "_eval_dataloaders"):
+                self._eval_dataloaders[dataloader_key] = eval_dataloader
+            else:
+                self._eval_dataloaders = {dataloader_key: eval_dataloader}
 
         return self.accelerator.prepare(eval_dataloader)
 
@@ -3584,12 +3598,13 @@ def evaluate(
             dictionary also contains the epoch number which comes from the training state.
         """
         # handle multipe eval datasets
-        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
+        override = eval_dataset is not None
+        eval_dataset = eval_dataset if override else self.eval_dataset
         if isinstance(eval_dataset, dict):
             metrics = {}
             for eval_dataset_name, _eval_dataset in eval_dataset.items():
                 dataset_metrics = self.evaluate(
-                    eval_dataset=_eval_dataset,
+                    eval_dataset=_eval_dataset if override else eval_dataset_name,
                     ignore_keys=ignore_keys,
                     metric_key_prefix=f"{metric_key_prefix}_{eval_dataset_name}",
                 )


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -1231,6 +1231,97 @@ def test_dataloader_without_dataset(self):
             trainer.train()
             trainer.evaluate()
 
+    def test_get_eval_dataloader_without_persistent_workers(self):
+        train_dataset = RegressionDataset()
+        config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4)
+        tiny_gpt2 = GPT2LMHeadModel(config)
+        args = TrainingArguments("./test", report_to="none", dataloader_persistent_workers=False)
+
+        # Single evaluation dataset
+        eval_dataset = RegressionDataset()
+        trainer = Trainer(tiny_gpt2, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
+        # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader
+        trainer.accelerator.prepare = lambda x: x
+
+        default_dataloader = trainer.get_eval_dataloader()
+        dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset)
+
+        self.assertEqual(default_dataloader.dataset, eval_dataset)
+        self.assertEqual(dataloader_with_dataset.dataset, eval_dataset)
+        self.assertNotEqual(default_dataloader, dataloader_with_dataset)
+
+        # Multiple evaluation datasets
+        first_dataset = RegressionDataset()
+        second_dataset = RegressionDataset()
+        trainer = Trainer(
+            tiny_gpt2,
+            args,
+            train_dataset=train_dataset,
+            eval_dataset={"first": first_dataset, "second": second_dataset},
+        )
+        # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader
+        trainer.accelerator.prepare = lambda x: x
+
+        first_dataloader = trainer.get_eval_dataloader("first")
+        first_dataloader_repeated = trainer.get_eval_dataloader("first")
+        second_dataloader = trainer.get_eval_dataloader("second")
+        second_dataloader_repeated = trainer.get_eval_dataloader("second")
+
+        self.assertEqual(first_dataset, first_dataloader.dataset)
+        self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset)
+        self.assertEqual(second_dataset, second_dataloader.dataset)
+        self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset)
+        self.assertNotEqual(first_dataloader, first_dataloader_repeated)
+        self.assertNotEqual(second_dataloader, second_dataloader_repeated)
+
+    def test_get_eval_dataloader_with_persistent_workers(self):
+        train_dataset = RegressionDataset()
+        config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4)
+        tiny_gpt2 = GPT2LMHeadModel(config)
+        args = TrainingArguments(
+            "./test",
+            report_to="none",
+            dataloader_persistent_workers=True,
+            dataloader_num_workers=2,
+        )
+
+        # Single evaluation dataset
+        eval_dataset = RegressionDataset()
+        trainer = Trainer(tiny_gpt2, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
+        # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader
+        trainer.accelerator.prepare = lambda x: x
+
+        default_dataloader = trainer.get_eval_dataloader()
+        dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset)
+
+        self.assertEqual(default_dataloader.dataset, eval_dataset)
+        self.assertEqual(dataloader_with_dataset.dataset, eval_dataset)
+        self.assertEqual(default_dataloader, dataloader_with_dataset)
+
+        # Multiple evaluation datasets
+        first_dataset = RegressionDataset()
+        second_dataset = RegressionDataset()
+        trainer = Trainer(
+            tiny_gpt2,
+            args,
+            train_dataset=train_dataset,
+            eval_dataset={"first": first_dataset, "second": second_dataset},
+        )
+        # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader
+        trainer.accelerator.prepare = lambda x: x
+
+        first_dataloader = trainer.get_eval_dataloader("first")
+        first_dataloader_repeated = trainer.get_eval_dataloader("first")
+        second_dataloader = trainer.get_eval_dataloader("second")
+        second_dataloader_repeated = trainer.get_eval_dataloader("second")
+
+        self.assertEqual(first_dataset, first_dataloader.dataset)
+        self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset)
+        self.assertEqual(second_dataset, second_dataloader.dataset)
+        self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset)
+        self.assertEqual(first_dataloader, first_dataloader_repeated)
+        self.assertEqual(second_dataloader, second_dataloader_repeated)
+
     @require_lomo
     @require_torch_gpu
     def test_lomo(self):