TPU Trainer memory leak and memory requirements
Author: misrasaurabh1
Pull request: no
Labels: name: wontfix

# 🐛 Bug
The TPU Trainer for Pytorch is extremely memory inefficient as it eats up lot of CPU memory while training. Moreover the training pipeline has a memory leak somewhere that leads to increased memory utilization as the training proceeds.

## Information
I am using a huge machine (n1-highmem-32) with 208 GB of RAM. I am trying to train seq-seq models like T5-base and BART-base with 880,000 examples. The problem is that with xla_spawn.py method of training models, multiprocessing starts 8 processes that all load up the features separately, resulting in 8x the memory usage than if they all shared the memory. Also this scaling of memory restricts what TPU one can use, I cannot imagine having to use a 1 TB RAM machine just to use v2-32 instance with 32 processes! I am trying to productionize a training pipeline with TPUs and the memory scaling problem makes it really hard and costly. TPUs are designed for training really large models and with large datasets but this memory issue is a large bottleneck against it. I could train with larger datasets but the memory usage prevents me.

In the images below one can see that for BART training, the memory utilization at the beginning is 80% but it rises to 100% in 14 hours, killing the training. I have to train the model 3 times to make progress.
![image](https://user-images.githubusercontent.com/1271289/86500284-e1169200-bd44-11ea-877f-337d5854c393.png)
![image](https://user-images.githubusercontent.com/1271289/86500267-c17f6980-bd44-11ea-8083-1949883883d9.png)
This is the memory usage with T5. It starts with 45% and ends with 75%.
![image](https://user-images.githubusercontent.com/1271289/86500290-f1c70800-bd44-11ea-8aac-8cc182f750a0.png)

Could you consider rearchitecting the TPU Trainer so that it does not takes humongous amounts of RAM? Maybe a streaming or lazy loading architecture that only loads the features as and when they are required? Also multiprocessing is inherently not scalable, does one launch 1024 processes to train on a 1024 Core TPU Pod? I have made a TPU Training pipeline with TensorFlow 1 previously and it worked spectacularly effectively with its architecture of TFRecords and Tensorflow handling all data streaming through Google storage. The pipeline performance scaled linearly with the number of TPU Cores without consuming large memory. (Image attached below).

![image](https://user-images.githubusercontent.com/1271289/86500867-12915c80-bd49-11ea-83ef-7cb9271c5022.png)

A scalable pipeline would make Huggingface a real contender for production level TPU PyTorch training pipelines.

Model I am using (Bert, XLNet ...): T5, BART

Language I am using the model on (English, Chinese ...): English

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:

Use the official example at https://github.com/huggingface/transformers/tree/master/examples#running-on-tpus

<!-- 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.-->

## Expected behavior

Memory usage contained and does not increase with time.

## 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: 2.11.0
- Platform: Linux-5.3.0-1026-gcp-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.7.0a0+4121d34 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: False
- Using distributed or parallel set-up in script?: Yes



Comments:
Hi @misrasaurabh1  I have also faced this issue. One immediate solution is to use lazy loading dataset, as xla by default loads the dataset in all processes. Something like `nlp` can really help here. as posted by @thomwolf  here https://gist.github.com/thomwolf/13ca2b2b172b2d17ac66685aa2eeba62

`nlp` can load full english wikipedia dataset (17 GB+) in just 9 MB of RAM. However I've not  been able to use it with TPU's yet.

Very interesting, if one could load data in a lazy-loading manner using `nlp` it would help a lot! Although I am not sure if we can use our custom datasets with `nlp`.

If there was a way to use a shared memory between different processes or each process reading independently from memory-mapped on drive, that could solve this problem.

It's possible to use custom datasets with `nlp`.  This section can help you https://github.com/huggingface/nlp/blob/master/CONTRIBUTING.md#how-to-add-a-dataset

We can create dataset as described in above steps, and then just pass the path of  created dataset to `nlp.load_dataset` 

Although I still want to highlight that this bug is still unfixed and a big problem. Drawing attention from the HuggingFace team!

Update on how I solved the large memory usage. Because I just wanted one copy of the features in memory, so I used a redis-server locally to cache the features in memory in a separate process. I am also using unix sockets to make connections faster. Items stored in Redis have keys equal to the array index and value equal to the pickled feature dictionary. In the __getitem__ I am unpickling the data on the fly.
This reduced memory consumption from 100s of GBs to 30GB. Surprisingly, the training also became 20% faster after using redis-server.

Thanks for sharing this @misrasaurabh1 

Thanks for letting us know @misrasaurabh1! We're very close to having TPU CI, only a few days away to a week, so we'll be able to monitor such metrics more closely.

Although the memory still keeps increasing as the training progresses. If it becomes too much of an issue for a reliable training pipeline I can dig deeper.

The memory spikes up everytime evaluation loop happens.
![image](https://user-images.githubusercontent.com/1271289/89491224-9df47800-d763-11ea-9e5a-692744df4663.png)


This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.