RAM OOM with a large save_steps using trainer API for MLM training
Author: Sumsky21
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.17.0
- Platform: Linux-4.14.105-1-tlinux3-0013-x86_64-with-glibc2.17
- Python version: 3.9.2
- PyTorch version (GPU?): 1.10.2+cu113 (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?: yes
- Using distributed or parallel set-up in script?: yes (in fact the trainer API handles it automatically, with 8 gpus)

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- T5, BART, Marian, Pegasus, EncoderDecoder: @patrickvonplaten
- Blenderbot, MBART: @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet, BigBird: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder, UniSpeech, UniSpeechSAT, SEW, SEW-D, Speech2Text: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten @narsil
- Tokenizers: @SaulLu
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->
Electra model: @LysandreJik   
Trainer: @sgugger 

## Information

Model I am using (Bert, XLNet ...): Electra(ForMaskedLM)

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:
* [x] an official task: Masked LM (training from scratch)
* [ ] my own task or dataset: (give details below)

## Expected behavior

I don't know if it was a bug, so I'll describe the situation first. 

I'm using the trainer API to train a ElectraForMaskedLM model from scratch. I write a simple script (will be shown below) and set `save_step` argument as 10000. But after running the script, I observe the RAM(memory, NOT CUDA) usage continuously increases, finally casued an OOM error thus terminate the docker. At the same time, GPU, CUDA memory and CPU usage is quite normal.

![屏幕截图 2022-03-13 171751](https://user-images.githubusercontent.com/58537872/158053023-cda95dd5-c43b-43f3-adc0-3d9e9f3247cd.png)

To solve the problem, I've tried adjusting several training arguments, and then find that when I set a smaller `save_step` (like 2000), RAM usage  goes up and down periodically, and the cycle length is approximately equal to the interval between two save checkpoints, like the figure below: 

![屏幕截图 2022-03-13 172757](https://user-images.githubusercontent.com/58537872/158053488-e3956015-9986-4485-a6f7-f1b8107ef7cc.png)

So is there a bug about the checkpoint saving mechanism? Or it's normal so I should have set a small `save_steps`?

The training script is below: 

```python
import torch
import numpy as np
from transformers import ElectraConfig, ElectraTokenizerFast, ElectraForMaskedLM
import os
os.environ["TOKENIZERS_PARALLELISM"] = "true"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = "cpu"
print(device)

# Tokenizer
print("Importing tokenizer and model...")
tokenizer = ElectraTokenizerFast(tokenizer_file='./models/bert-custom.json')
# Model
config = ElectraConfig(vocab_size=10000)
model = ElectraForMaskedLM(config=config).to(device)

# dataset
print("Importing dataset...")
from transformers import LineByLineTextDataset
dataset = LineByLineTextDataset(
    tokenizer=tokenizer,
    file_path="./data/in_neg_1e7.csv",
    block_size=80,
)
eval_set = LineByLineTextDataset(
    tokenizer=tokenizer,
    file_path="./data/in_eval_1e6.csv",
    block_size=80,
)

from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

from transformers import Trainer, TrainingArguments, EarlyStoppingCallback
print("Start Training...")
training_args = TrainingArguments(
    output_dir="./electra_save/",
    overwrite_output_dir=True,
    num_train_epochs=50,
    per_device_train_batch_size=512,
    per_device_eval_batch_size=512,
    eval_accumulation_steps=300,
    save_steps=10000,                             # the key argument
    save_total_limit=2,
    dataloader_num_workers=80,
    evaluation_strategy='steps', 
    eval_steps=10000,
    logging_steps=1000,
    metric_for_best_model='eval_loss',
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    tokenizer=tokenizer,
    train_dataset=dataset,
    eval_dataset=eval_set,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
)

trainer.train()

trainer.save_model(output_dir='./electra_save/')
```

## To reproduce

Steps to reproduce the behavior:

1. run the training script above with a large dataset (I use a dataset with 10M lines of text) and a large `save_steps` argument
2. wait for a period of time and monitor on the memory usage
3. encounter with OOM

<!-- 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:
Addition: Previously I've used a similar script to train BertForMaskedLM, and no memory fluctuations were observed that time: 
![屏幕截图 2022-03-13 185104](https://user-images.githubusercontent.com/58537872/158056119-dd62f856-06fa-410b-8c7e-95d3ec5266fa.png)
That's why I suspect it's abnormal this time.

So the issue only appears with Electra, right? Did you encounter it for any other model?

Yes. I've only tried on Bert and Electra so far (and there's no problem during Bert training, as above figure has shown). 

If there's a need, I can try the similar script on RobertaforMaskedLM in next few days.

That would be helpful if you could!

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.