CPU OOM when using IterableDataset with dataloader_num_workers > 0
Author: nghuyong
Pull request: no

## Environment info
- `transformers` version: 4.15.0
- Platform: Linux-3.10.107-1-tlinux2_kvm_guest-0049-x86_64-with-glibc2.10
- Python version: 3.8.8
- PyTorch version (GPU?): 1.8.1 (False)
- 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?:  Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
@sgugger 

## Information
I want to pre-train a large language model. So the train-dataset is huge (~100G), I need to use IterableDataset. and the dev dataset is small (~10M) so I use Dataset. 
In order to speed up the training process, the `dataloader_num_workers` of `TrainingArguments` is set to 10. But when I train the model, the CPU memory keeps increasing and gets OOM finally.

![image](https://user-images.githubusercontent.com/16462374/151107817-075ea140-5112-4a0b-9d82-ea718c867100.png)

`dataloader_num_workers` of `TrainingArguments` is set to 0 and the CPU memory is stable but the training is too slow!

May be It is the same problem of this https://pytorch.org/docs/stable/data.html#multi-process-data-loading 

I want to know how to correctly use IterableDataset with huggingface's trainer (train-dataset is IterableDataset, and dev-dataset is Dataset, and dataloader_num_workers > 0)? Thanks a lot!

## To reproduce

I write a sample code to reproduce:

demo_train.py
```Python
import argparse
from itertools import cycle
from typing import Optional
from dataclasses import dataclass
from torch.nn import CrossEntropyLoss
from torch.utils.data import Dataset, IterableDataset
from transformers import BertTokenizer, Trainer, TrainingArguments, BertConfig, PreTrainedModel
import torch
import os
import numpy as np
from transformers.file_utils import ModelOutput


class DemoEvalDataset(Dataset):

    def __init__(self, data_file):
        self.items = []
        with open(data_file, 'rt', encoding='utf-8', buffering=10000) as f:
            for line in f:
                self.items.append(line.strip())

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        return self.items[index]


class DemoTrainDataset(IterableDataset):

    def __init__(self, data_file):
        super().__init__()
        self.data_file = data_file

    def get_stream(self):
        return cycle(self.parse_file_by_line())

    def __iter__(self):
        return self.get_stream()

    def parse_file_by_line(self):
        with open(self.data_file, 'rt', encoding='utf-8', buffering=10000) as file_obj:
            for line in file_obj:
                yield line.strip()

@dataclass
class DemoModelOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None


class DemoModel(PreTrainedModel):
    """
    token level binary classification
    """

    def __init__(self, config):
        super().__init__(config)
        config.hidden_size = 100
        config.label_num = 2
        self.config = config
        self.embedding = torch.nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
        self.mlp = torch.nn.Linear(config.hidden_size, config.label_num)
        self.loss_fct = CrossEntropyLoss()

    def forward(self, input_ids, labels=None):
        logits = self.mlp(self.embedding(input_ids))

        loss = None
        if labels is not None:
            loss = self.loss_fct(logits.view(-1, self.config.label_num), labels.view(-1))

        return DemoModelOutput(
            loss=loss,
            logits=logits
        )


def data_collator(tokenizer, features):
    texts = features
    texts_tokenized = tokenizer(texts, padding=True, return_tensors='np')
    input_ids = texts_tokenized['input_ids']
    batch_size, max_seq_len = input_ids.shape
    labels = np.random.randint(2, size=(batch_size, max_seq_len))
    return {
        'input_ids': torch.tensor(input_ids, dtype=torch.long),
        'labels': torch.tensor(labels, dtype=torch.long)
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument('--pretrained_model_path', required=True)
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--train_data_file', required=True)
    parser.add_argument('--eval_data_file', required=True)
    parser.add_argument('--train_batch_size', type=int, default=128)
    parser.add_argument('--eval_batch_size', type=int, default=128)
    parser.add_argument('--logging_steps', type=int, default=100)
    parser.add_argument('--eval_steps', type=int, default=1000)
    parser.add_argument('--save_steps', type=int, default=1000)
    parser.add_argument('--learning_rate', type=float, default=5e-5)
    parser.add_argument('--max_steps', type=int, default=100000)
    args = parser.parse_args()
    model_config = BertConfig.from_pretrained(args.pretrained_model_path)
    model = DemoModel(model_config)
    train_dataset = DemoTrainDataset(args.train_data_file)
    eval_dataset = DemoEvalDataset(args.eval_data_file)
    train_args = TrainingArguments(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.train_batch_size,
        per_device_eval_batch_size=args.eval_batch_size,
        do_train=True,
        do_eval=True,
        evaluation_strategy="steps",
        logging_steps=args.logging_steps,
        eval_steps=args.eval_steps,
        save_steps=args.save_steps,
        overwrite_output_dir=True,
        save_total_limit=1,
        local_rank=int(os.environ.get('LOCAL_RANK', -1)),
        learning_rate=args.learning_rate,
        metric_for_best_model='eval_loss',
        fp16=True,
        max_steps=args.max_steps,
        dataloader_num_workers=10 # here !!!
    )
    bert_tokenizer = BertTokenizer.from_pretrained(args.pretrained_model_path)
    trainer = Trainer(
        model=model,
        args=train_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=lambda features: data_collator(bert_tokenizer, features),
    )
    trainer.train()
```

run.sh
```bash
python -m torch.distributed.launch \
  --nproc_per_node=8 \
  demo_train.py \
  --pretrained_model_path $PRETRAINED_MODELS_DIR/roberta \
  --output_dir $CHECKPOINT_DIR \
  --train_data_file $DATA_DIR/large_file.txt \
  --eval_data_file $DATA_DIR/dev_file.txt \
  --learning_rate 5e-5 \
  --max_step 300000
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->


Comments:
DO NOT USE `from itertools import cycle`

@Ethan-yt  It works👍

```
class DemoTrainDataset(IterableDataset):

    def __init__(self, data_file):
        super().__init__()
        self.data_file = data_file

    def __iter__(self):
        while True:
            with open(self.data_file, 'rt', encoding='utf-8', buffering=10000) as file_obj:
                for line in file_obj:
                    yield line.strip()

```

![image](https://user-images.githubusercontent.com/16462374/151121188-db7d6a90-b5e4-41df-a787-6e4a7a82439a.png)


@sgugger 
When I use IterableDataset and dataloader_num_workers > 0, should I manually handle multiprocessing in the __iter__ follow the https://pytorch.org/docs/stable/data.html#multi-process-data-loading

![image](https://user-images.githubusercontent.com/16462374/151128420-b8372fcf-f4ee-4239-a426-05272c5f4754.png)


Yes, that's your responsibility, as the PyTorch doc highlights.

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/master/CONTRIBUTING.md) are likely to be ignored.