Memory accumulation when using hybrid clip script
Author: galuhsahid
Pull request: no

## Environment info
- `transformers` version: 4.9.1
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cpu (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.18
- JaxLib version: 0.1.69
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: yes, TPU v3-8

### Who can help
@patil-suraj 

## Information

Model I am using (Bert, XLNet ...): BERT + ViT

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)

I'm currently working on pretraining CLIP for Indonesian using scripts that are based on the [Hybrid CLIP](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects/hybrid_clip) example. I've run the code using the COCO dataset and it ran with no problem and managed to produce a working model, however I didn't track this run with wandb so I'm not sure how the system looked like with COCO.

Right now I'm trying to train a model on a larger dataset (~12M image-text pairs). However it seems that the memory keeps accumulating as seen on the graph below. During one of my runs it (not tracked) it eventually crashed after ~7 hours which is why I noticed this. After rerunning and tracking it (terminated - [here](https://wandb.ai/galuh/clip-indonesian/runs/2qb4zp6v?workspace=user-galuh) is the wandb run) turns out it's probably because of the memory:

![image](https://user-images.githubusercontent.com/10180442/129117152-cd9cb507-532a-482f-bdfb-a99178f2ccb8.png)

Looks like the sharp jumps happened during evaluation.

Changes on the script compared to the hybrid clip example:
- add wandb logging
- save training, evaluation metrics, and checkpoints in steps instead of epochs -> at first I thought the memory accumulation was due to the training and evaluation metrics, however the issue still persists despite having logged and cleared my training metrics with a 200 step interval <note: if it's OK to add to the existing script, maybe I can make a separate issue to add logging and saving by steps and make a PR for that?>
- use adafactor instead of adamw

## To reproduce

Steps to reproduce the behavior:

1. Run the code in [this folder](https://github.com/galuhsahid/clip-indonesian/tree/main/hybrid_clip) by running `run_training.sh`, with any large image-text dataset that is prepared as instructed in the [readme of the examples folder](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects/hybrid_clip) (jsonl file).

## Expected behavior

Memory stays roughly constant throughout


Comments:
Thanks a lot for posting the detailed issue.

I'm not exactly sure about this, but could you try disabling `persistent_workers` in `DataLoader`, it was causing some issues for another team. 

> if it's OK to add to the existing script, maybe I can make a separate issue to add logging and saving by steps and make a PR for that?

yes, that would be great! Feel free to open a PR for that :)

Along with what @patil-suraj suggested, you should also bring down your `num_workers` to 16 or 32 instead of 96. That should keep your memory in check if you're planning to train your model for 2-3 days.

@bhavitvyamalik Yup, for the run I linked in the original thread I've used `num_workers` = 16. I have used 96 earlier which crashed the run almost immediately with ~12M images. Decreasing the `num_workers` helped a lot indeed, although the memory is still seeing an increasing trend.

@patil-suraj Thank you for the suggestion - disabling `persistent_workers` seems to work, this is how the graph looks like after I disabled it ([wandb run](https://wandb.ai/galuh/clip-indonesian/runs/33dqdxtd/system?workspace=user-galuh)):

![image](https://user-images.githubusercontent.com/10180442/129300348-6b7a7ae0-d9bb-4c85-9c15-ec7ef5c330ed.png)

It's still increasing but at a much slower pace than when `persistent_workers` was enabled and somehow drops again much later. Not sure if this is expected though. 

(If it's not expected - this might be somehow related to this issue https://github.com/pytorch/pytorch/issues/13246. I've tried one of the solutions (converting the `examples` list in the DataLoader into np.array) but I'm still seeing the same increasing trend)

Also sure would be happy to open a PR later! Thank you

Looking at the PyTorch issue it does seem related to the dataloader.

In the `ImageTextDataset` in `run_hybrid_clip.py`, all examples, captions, image_paths are stored in python lists

https://github.com/huggingface/transformers/blob/bda1cb02360fc9d290636cfcb6bcbeb4a18484ce/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py#L220-L228

As suggested in that issue, could try storing the examples in a zero-copy object?

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.