Issue: Trainer error on `evaluate()` in multithreaded/distributed context (shape mismatch)

Environment info

• transformers version: 4.3.3
• Platform: Linux-5.4.83.1.fi-x86_64-with-centos-7.8.2003-Core
• Python version: 3.7.3
• PyTorch version (GPU?): 1.8.0 (True)
• Tensorflow version (GPU?): not installed (NA)
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: yes - multinode/multigpu and multigpu settings.

Who can help

@LysandreJik @sgugger

Information

Model I am using (GPT2):

The problem arises when using:

• the official example scripts: (give details below)
• my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: (give the name)
• my own task or dataset: (give details below)

To reproduce

Steps to reproduce the behavior: I have witnessed this error in two contexts Using a custom torch.utils.data.IterableDataset.

First:

1. specify dataloader_num_workers > 1 in TrainingArguments and run trainer.train() with an eval dataset

Second:

1. In distributed setting, fire up multiple training instances on separate nodes using the torch.distributed.launch command, run trainer.train() with an eval dataset

Error message:

  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer.py", line 1655, in prediction_loop
    eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer_pt_utils.py", line 338, in add_arrays
    slice_len = self._nested_set_tensors(self._storage, arrays)
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer_pt_utils.py", line 354, in _nested_set_tensors
    storage[self._offsets[i] : self._offsets[i] + slice_len] = arrays[i * slice_len : (i + 1) * slice_len]
ValueError: could not broadcast input array from shape (104,) into shape (96,)

The broadcast input array shape varies. In the first case, the broadcast shape will be dataloader_num_workers * expected_shape (in this case (96,)). Above exhibits the second case error message.

Expected behavior

The evaluate loop should run without error.

Dataset information

The dataset object is an IterableDataset that is abc.Sized.

Script information

The script is fairly generic, involving training and evaluating GPT2 via the Trainer object for next-token prediction.

[sgugger]

That’s correct, especially for distributed evaluation.

[djberenberg]

Ok, that makes sense. So just to conclude, transformers.Trainer won’t work in distributed setting with an torch.utils.data.IterableDataset, in principal due to the fact that IterableDatasets are not amenable to that use case, since it isn’t clear how describe a distributed sampling procedure for them. Is that correct? Thanks in advance