Issue: Trainer error on `evaluate()` in multithreaded/distributed context (shape mismatch)
Author: djberenberg
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `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)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] 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.


Comments:
I am unsure on what you want us to do: all the example scripts have been tested with evaluation and work in a distributed setup. Unless you share your script, there is little we can do to fix an issue we have not encountered.

Also make sure you have the latest version of Transformers installed as this bug might have been fixed already.

@sgugger Hello and thank you for the reply! Understood it's not clear how to help here. Unfortunately the error persists in version `4.5.0`. For a minimal example, will you need data or is code good enough? It's a very nonstandard dataset, composed of DNA strings ( and only 4 coding tokens plus a pad), but the only nonstandard way I am interacting with `transformers` is by feeding the custom datasets to a `Trainer`. 

training code:
```python
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import logging
from typing import Optional

import shutil
import argparse
import itertools
from pathlib import Path
from functools import partial

import torch
from tokenizers import ByteLevelBPETokenizer
from tokenizers.pre_tokenizers import Whitespace

import transformers

from min_data import DNADataset, read_list, chunk_fasta

def arguments():
    parser = argparse.ArgumentParser(description="Train GPT-2 model on DNA data.")
    parser.add_argument("--partitions", type=Path, nargs=2,
            help="Train, validation partition files")

    parser.add_argument("--session", type=Path,
            help="Training session directory; models are saved here along with other important metadata")

    parser.add_argument("--log-to", type=Path,
            help="Tensorboard logging root directory; logs will write to log_dir/session_name",
            default="tensorboard_logs",
            dest='log_dir')

    parser.add_argument("--tokenizer", 
                type=Path,
                help="Specify pre-trained tokenizer if it exists",
                required=True)

    # architecture specs
    parser.add_argument("--n-layer", type=int, default=4, help="# of layers")
    parser.add_argument("--n-embed", type=int, default=16, help="embedding dim.")
    parser.add_argument("--n-inner", type=int, default=1024, help="hidden dim.")
    parser.add_argument("--chunk-size", type=int, default=2000, help="max base pair width", dest="chunksize")
    parser.add_argument("--lr", default=1e-4, type=float)

    # training/logging specs
    parser.add_argument("--train-epochs", type=int, default=1)
    parser.add_argument("--save-steps", type=int, default=250)
    parser.add_argument("--save-up-to", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--lens", nargs=2, type=int, default=(10_000, 2000))
    parser.add_argument("--progress-bar", action='store_true', default=False, dest='tqdm')
    parser.add_argument("--local_rank", type=int, default=-1)
    return parser

def get_training_args(args: argparse.Namespace):
    output_dir = args.session / "outputs"
    save_steps = args.save_steps
    save_limit = args.save_up_to
    batch_size = args.batch_size
    log_dir = args.log_dir / args.session.name
    lr = args.lr
    
    max_steps = args.train_epochs * args.lens[0]
    return transformers.TrainingArguments(
            output_dir=str(output_dir),
            overwrite_output_dir=True,
            ddp_find_unused_parameters=False,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            evaluation_strategy="steps",
            learning_rate=lr,
            local_rank=args.local_rank,
            disable_tqdm=not args.tqdm,
            max_steps=max_steps,
            eval_steps=save_steps,
            save_steps=save_steps,
            prediction_loss_only=True,
            #logging_dir=str(log_dir),
            save_total_limit=save_limit
            )

def construct_dataset_from(fasta_list: Path, tokenizer,  **kwargs):
    dataset = DNADataset( read_list(fasta_list), tokenizer, **kwargs )
    return dataset

class DNATrainer(transformers.Trainer): # overwritten because error occur wrt sampler for  IterableDataset, still seems necessary in 4.5.0
   def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
        if isinstance(eval_dataset, torch.utils.data.IterableDataset):
            return None

if __name__ == '__main__':
    args: argparse.Namespace  = arguments()
        
    # get tokenizer, gpt2 config, training arguments 
    tokenizer: tokenizers.ByteLevelBPETokenizer   = LOAD_TOKENIZER(args)  # standard loading of tokenizer
    config: transformers.GPT2Config               = GET_GPT2_CONFIG(args) # generates a GPT2Config 
    training_args: transformers.TrainingArguments = get_training_args(args)

    model  = transformers.GPT2LMHeadModel(config=config)
    model.resize_token_embeddings(len(tokenizer))
    model.train()

    kwargs = {'pad_token': '<pad>', 'chunksize': args.chunksize}
    train_part, val_part = args.partitions # list of files for each dataset to stream from
    
    train_data = construct_dataset_from(train_part, tokenizer, asserted_len=args.lens[0], **kwargs)
    val_data   = construct_dataset_from(val_part, tokenizer, asserted_len=args.lens[1], **kwargs)

    collator  = transformers.DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    trainer = DNATrainer(model=model,
            args=training_args,
            data_collator=collator,
            train_dataset=train_data,
            eval_dataset=val_data,
            )

    trainer.train()
```

Oh, but you're overriding the sampler part of the `Trainer` code. There is no way distributed evaluation can work then, as it relies on this.

Aha! Thank you! I'm sure this is the right track, but now I am back to an error in how the `Trainer` chooses a sampler and constructs the `DataLoader`: (this is `transformers` version `4.5.0`)

```
ValueError: DataLoader with IterableDataset: expected unspecified sampler option, but got sampler=<transformers.trainer_pt_utils.SequentialDistributedSampler object at 0x1554e8f1e2b0>
self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer.py", line 1265, in _maybe_log_save_evaluate
    eval_dataloader = self.get_eval_dataloader(eval_dataset)
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer.py", line 612, in get_eval_dataloader
    pin_memory=self.args.dataloader_pin_memory,
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 231, in __init__
metrics = self.evaluate()
      File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer.py", line 1754, in evaluate
    "sampler option, but got sampler={}".format(sampler))
ValueError: DataLoader with IterableDataset: expected unspecified sampler option, but got sampler=<transformers.trainer_pt_utils.SequentialDistributedSampler object at 0x1554e8039dd8>
eval_dataloader = self.get_eval_dataloader(eval_dataset)
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/transformers/trainer.py", line 612, in get_eval_dataloader
    pin_memory=self.args.dataloader_pin_memory,
  File "/mnt/home/dberenberg/projects/metagenomics/huggingface_meta/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 231, in __init__
    "sampler option, but got sampler={}".format(sampler))
ValueError: DataLoader with IterableDataset: expected unspecified sampler option, but got sampler=<transformers.trainer_pt_utils.SequentialDistributedSampler object at 0x1554f2898c88>
```

Yes you need to forego the inheritance to `IterableDataset` as PyTorch does not let you take a sampler for those, so you will nee to implement a `__getitem__` instead of the `__iter__` for the evaluation.

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 `IterableDataset`s 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

That's correct, especially for distributed evaluation.