Load large text file for LM pre-training resulting in OOM
@lhoestq Thanks for pointing to n1t0, just to clarify. That warning was doing fine-tuning, without collator:
```

from datasets import load_dataset, load_metric
import numpy as np

GLUE_TASKS = [
    "cola",
    "mnli",
    "mnli-mm",
    "mrpc",
    "qnli",
    "qqp",
    "rte",
    "sst2",
    "stsb",
    "wnli",
]
task = "mnli"
actual_task = "mnli" if task == "mnli-mm" else task
dataset = load_dataset("glue", actual_task)
metric = load_metric("glue", actual_task)
batch_size = 16
attention_type = "linear"

from transformers.models.mobilebert_mod import (
    MobileBertForSequenceClassification,
    MobileBertTokenizerFast,
)
from transformers.models.mobilebert_mod.configuration_mobilebert import (
    MobileBertConfigMod,
)
from transformers import TrainingArguments, Trainer

num_labels = 3 if task.startswith("mnli") else 1 if task == "stsb" else 2
tokenizer = MobileBertTokenizerFast.from_pretrained(
    "/media/ad/00b5422b-9d54-4449-8b5d-08eab5cdac8c/training_trfm/big_linear_layerdrop_shared/checkpoint-23000/",
    max_len=512,
)
model = MobileBertForSequenceClassification.from_pretrained(
    "/media/ad/00b5422b-9d54-4449-8b5d-08eab5cdac8c/training_trfm/big_linear_layerdrop_shared/checkpoint-23000/",
    num_labels=num_labels,
)
print(model.num_parameters())

task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mnli-mm": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

sentence1_key, sentence2_key = task_to_keys[task]
if sentence2_key is None:
    print(f"Sentence: {dataset['train'][0][sentence1_key]}")
else:
    print(f"Sentence 1: {dataset['train'][0][sentence1_key]}")
    print(f"Sentence 2: {dataset['train'][0][sentence2_key]}")


def preprocess_function(examples):
    if sentence2_key is None:
        return tokenizer(examples[sentence1_key], truncation=True)
    return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)


encoded_dataset = dataset.map(preprocess_function, batched=True)
metric_name = (
    "pearson"
    if task == "stsb"
    else "matthews_correlation"
    if task == "cola"
    else "accuracy"
)

args = TrainingArguments(
    f"test-glue/{task}_{attention_type}",
    evaluation_strategy="steps",
    learning_rate=1e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    logging_steps=200,
    num_train_epochs=5,
    gradient_accumulation_steps=1,
    warmup_steps=10000,
    fp16=True,
    dataloader_num_workers=10,
    weight_decay=0.1,
    load_best_model_at_end=True,
    metric_for_best_model=metric_name,
)


def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    if task != "stsb":
        predictions = np.argmax(predictions, axis=1)
    else:
        predictions = predictions[:, 0]
    return metric.compute(predictions=predictions, references=labels)


validation_key = (
    "validation_mismatched"
    if task == "mnli-mm"
    else "validation_matched"
    if task == "mnli"
    else "validation"
)

trainer = Trainer(
    model,
    args,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset[validation_key],
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

trainer.train()
```

Now, I have come back to pre-training. The changes that I think I have done are: not formatting the dataset to torch: ~~`big_dataset.set_format(type='torch', columns=["text", "input_ids", "attention_mask", "token_type_ids"])`~~ so maybe some column is dropped and not freezed in memory and now I have not setted any validation dataset in the trainer. 

My validation dataset before:
```
book_corpus_eval = load_dataset(
    "bookcorpus",
    "plain_text",
    cache_dir="/home/ad/Desktop/bookcorpus",
    split="train[98:99%]",
)
book_corpus_eval = book_corpus_eval.map(encode, batched=True)
book_corpus_eval.set_format(
    type="torch", columns=["text", "input_ids", "attention_mask", "token_type_ids"]
)
**book_corpus_eval = book_corpus_eval.select([i for i in range(1500)])**
```
Maybe _selecting_ or indexing the dataset before feeding it to the trainer, do something strange.

My trainer now:
```

big_dataset = load_from_disk("/home/ad/Desktop/35percent_data.arrow/")

from transformers import DataCollatorForWholeWordMask

data_collator = DataCollatorForWholeWordMask(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15)

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./big_linear_layerdrop_shared_silu_secondtry",
    overwrite_output_dir=True,
    per_device_train_batch_size=60,
    per_device_eval_batch_size=60,
    save_steps=500,
    save_total_limit=10,
    logging_first_step=True,
    logging_steps=100,
#     evaluation_strategy='steps',
#     eval_steps=250,
    gradient_accumulation_steps=8,
    fp16=True,
    dataloader_num_workers=10,
    warmup_steps=15000,
    learning_rate=6e-4,
    adam_epsilon=1e-6,
    adam_beta2=0.98,
    weight_decay=0.01,
    max_grad_norm=1.0,
    max_steps=500000, 
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=big_dataset,
#     eval_dataset=book_corpus_eval,
    tokenizer=tokenizer)

import wandb
wandb.login()

trainer.train()
```

And surprisingly, the ram now keeps going up and down. The training is up now for 12h without collapse the ram. I don't know what could cause the leakage. :mag: 

Edit: I didn't see the swap memory, that keeps increasing. So the problem persist. 
I tried to pretrain Longformer using transformers and datasets. But I got OOM issues with loading a large text file. My script is almost like this:

```python
from datasets import load_dataset

@dataclass
class DataCollatorForDatasetsLanguageModeling(DataCollatorForLanguageModeling):
    """
    Data collator used for language modeling based on DataCollatorForLazyLanguageModeling
    - collates batches of tensors, honoring their tokenizer's pad_token
    - preprocesses batches for masked language modeling
    """

    block_size: int = 512

    def __call__(self, examples: List[dict]) -> Dict[str, torch.Tensor]:
        examples = [example['text'] for example in examples]
        batch, attention_mask = self._tensorize_batch(examples)
        if self.mlm:
            inputs, labels = self.mask_tokens(batch)
            return {"input_ids": inputs, "labels": labels}
        else:
            labels = batch.clone().detach()
            if self.tokenizer.pad_token_id is not None:
                labels[labels == self.tokenizer.pad_token_id] = -100
            return {"input_ids": batch, "labels": labels}

    def _tensorize_batch(self, examples: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:

        if self.tokenizer._pad_token is None:
            raise ValueError(
                "You are attempting to pad samples but the tokenizer you are using"
                f" ({self.tokenizer.__class__.__name__}) does not have one."
            )

        tensor_examples = self.tokenizer.batch_encode_plus(
            [ex for ex in examples if ex],
            max_length=self.block_size,
            return_tensors="pt",
            pad_to_max_length=True,
            return_attention_mask=True,
            truncation=True,
        )

        input_ids, attention_mask = tensor_examples["input_ids"], tensor_examples["attention_mask"]
        return input_ids, attention_mask

dataset = load_dataset('text', data_files='train.txt',cache_dir="./", , split='train')
data_collator = DataCollatorForDatasetsLanguageModeling(tokenizer=tokenizer, mlm=True, 
                      mlm_probability=0.15, block_size=tokenizer.max_len)
trainer = Trainer(model=model, args=args, data_collator=data_collator,
                      train_dataset=train_dataset, prediction_loss_only=True, )
trainer.train(model_path=model_path)
```
This train.txt is about 1.1GB and has 90k lines where each line is a sequence of 4k words. 
During training, the memory usage increased fast as the following graph and resulted in OOM before the finish of training.

![image](https://user-images.githubusercontent.com/29704017/93292112-5576b280-f817-11ea-8da2-b2db9bf35665.png)

Could you please give me any suggestions on why this happened and how to fix it?
Thanks. 
Load large text file for LM pre-training resulting in OOM 
 I tried to pretrain Longformer using transformers and datasets. But I got OOM issues with loading a large text file. My script is almost like this:

```python
from datasets import load_dataset

@dataclass
class DataCollatorForDatasetsLanguageModeling(DataCollatorForLanguageModeling):
    """
    Data collator used for language modeling based on DataCollatorForLazyLanguageModeling
    - collates batches of tensors, honoring their tokenizer's pad_token
    - preprocesses batches for masked language modeling
    """

    block_size: int = 512

    def __call__(self, examples: List[dict]) -> Dict[str, torch.Tensor]:
        examples = [example['text'] for example in examples]
        batch, attention_mask = self._tensorize_batch(examples)
        if self.mlm:
            inputs, labels = self.mask_tokens(batch)
            return {"input_ids": inputs, "labels": labels}
        else:
            labels = batch.clone().detach()
            if self.tokenizer.pad_token_id is not None:
                labels[labels == self.tokenizer.pad_token_id] = -100
            return {"input_ids": batch, "labels": labels}

    def _tensorize_batch(self, examples: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:

        if self.tokenizer._pad_token is None:
            raise ValueError(
                "You are attempting to pad samples but the tokenizer you are using"
                f" ({self.tokenizer.__class__.__name__}) does not have one."
            )

        tensor_examples = self.tokenizer.batch_encode_plus(
            [ex for ex in examples if ex],
            max_length=self.block_size,
            return_tensors="pt",
            pad_to_max_length=True,
            return_attention_mask=True,
            truncation=True,
        )

        input_ids, attention_mask = tensor_examples["input_ids"], tensor_examples["attention_mask"]
        return input_ids, attention_mask

dataset = load_dataset('text', data_files='train.txt',cache_dir="./", , split='train')
data_collator = DataCollatorForDatasetsLanguageModeling(tokenizer=tokenizer, mlm=True, 
                      mlm_probability=0.15, block_size=tokenizer.max_len)
trainer = Trainer(model=model, args=args, data_collator=data_collator,
                      train_dataset=train_dataset, prediction_loss_only=True, )
trainer.train(model_path=model_path)
```
This train.txt is about 1.1GB and has 90k lines where each line is a sequence of 4k words. 
During training, the memory usage increased fast as the following graph and resulted in OOM before the finish of training.

![image](https://user-images.githubusercontent.com/29704017/93292112-5576b280-f817-11ea-8da2-b2db9bf35665.png)

Could you please give me any suggestions on why this happened and how to fix it?
Thanks.  
 @lhoestq Thanks for pointing to n1t0, just to clarify. That warning was doing fine-tuning, without collator:
```

from datasets import load_dataset, load_metric
import numpy as np

GLUE_TASKS = [
    "cola",
    "mnli",
    "mnli-mm",
    "mrpc",
    "qnli",
    "qqp",
    "rte",
    "sst2",
    "stsb",
    "wnli",
]
task = "mnli"
actual_task = "mnli" if task == "mnli-mm" else task
dataset = load_dataset("glue", actual_task)
metric = load_metric("glue", actual_task)
batch_size = 16
attention_type = "linear"

from transformers.models.mobilebert_mod import (
    MobileBertForSequenceClassification,
    MobileBertTokenizerFast,
)
from transformers.models.mobilebert_mod.configuration_mobilebert import (
    MobileBertConfigMod,
)
from transformers import TrainingArguments, Trainer

num_labels = 3 if task.startswith("mnli") else 1 if task == "stsb" else 2
tokenizer = MobileBertTokenizerFast.from_pretrained(
    "/media/ad/00b5422b-9d54-4449-8b5d-08eab5cdac8c/training_trfm/big_linear_layerdrop_shared/checkpoint-23000/",
    max_len=512,
)
model = MobileBertForSequenceClassification.from_pretrained(
    "/media/ad/00b5422b-9d54-4449-8b5d-08eab5cdac8c/training_trfm/big_linear_layerdrop_shared/checkpoint-23000/",
    num_labels=num_labels,
)
print(model.num_parameters())

task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mnli-mm": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

sentence1_key, sentence2_key = task_to_keys[task]
if sentence2_key is None:
    print(f"Sentence: {dataset['train'][0][sentence1_key]}")
else:
    print(f"Sentence 1: {dataset['train'][0][sentence1_key]}")
    print(f"Sentence 2: {dataset['train'][0][sentence2_key]}")


def preprocess_function(examples):
    if sentence2_key is None:
        return tokenizer(examples[sentence1_key], truncation=True)
    return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)


encoded_dataset = dataset.map(preprocess_function, batched=True)
metric_name = (
    "pearson"
    if task == "stsb"
    else "matthews_correlation"
    if task == "cola"
    else "accuracy"
)

args = TrainingArguments(
    f"test-glue/{task}_{attention_type}",
    evaluation_strategy="steps",
    learning_rate=1e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    logging_steps=200,
    num_train_epochs=5,
    gradient_accumulation_steps=1,
    warmup_steps=10000,
    fp16=True,
    dataloader_num_workers=10,
    weight_decay=0.1,
    load_best_model_at_end=True,
    metric_for_best_model=metric_name,
)


def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    if task != "stsb":
        predictions = np.argmax(predictions, axis=1)
    else:
        predictions = predictions[:, 0]
    return metric.compute(predictions=predictions, references=labels)


validation_key = (
    "validation_mismatched"
    if task == "mnli-mm"
    else "validation_matched"
    if task == "mnli"
    else "validation"
)

trainer = Trainer(
    model,
    args,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset[validation_key],
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

trainer.train()
```

Now, I have come back to pre-training. The changes that I think I have done are: not formatting the dataset to torch: ~~`big_dataset.set_format(type='torch', columns=["text", "input_ids", "attention_mask", "token_type_ids"])`~~ so maybe some column is dropped and not freezed in memory and now I have not setted any validation dataset in the trainer. 

My validation dataset before:
```
book_corpus_eval = load_dataset(
    "bookcorpus",
    "plain_text",
    cache_dir="/home/ad/Desktop/bookcorpus",
    split="train[98:99%]",
)
book_corpus_eval = book_corpus_eval.map(encode, batched=True)
book_corpus_eval.set_format(
    type="torch", columns=["text", "input_ids", "attention_mask", "token_type_ids"]
)
**book_corpus_eval = book_corpus_eval.select([i for i in range(1500)])**
```
Maybe _selecting_ or indexing the dataset before feeding it to the trainer, do something strange.

My trainer now:
```

big_dataset = load_from_disk("/home/ad/Desktop/35percent_data.arrow/")

from transformers import DataCollatorForWholeWordMask

data_collator = DataCollatorForWholeWordMask(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15)

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./big_linear_layerdrop_shared_silu_secondtry",
    overwrite_output_dir=True,
    per_device_train_batch_size=60,
    per_device_eval_batch_size=60,
    save_steps=500,
    save_total_limit=10,
    logging_first_step=True,
    logging_steps=100,
#     evaluation_strategy='steps',
#     eval_steps=250,
    gradient_accumulation_steps=8,
    fp16=True,
    dataloader_num_workers=10,
    warmup_steps=15000,
    learning_rate=6e-4,
    adam_epsilon=1e-6,
    adam_beta2=0.98,
    weight_decay=0.01,
    max_grad_norm=1.0,
    max_steps=500000, 
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=big_dataset,
#     eval_dataset=book_corpus_eval,
    tokenizer=tokenizer)

import wandb
wandb.login()

trainer.train()
```

And surprisingly, the ram now keeps going up and down. The training is up now for 12h without collapse the ram. I don't know what could cause the leakage. :mag: 

Edit: I didn't see the swap memory, that keeps increasing. So the problem persist. 