CUDA out of memory for bart-large while using deepspeed with Zero stage 3
Author: xpact
Pull request: no

### System Info

- `transformers` version: 4.25.1
- Platform: Linux-5.10.147+-x86_64-with-glibc2.27
- Python version: 3.8.16
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.0+cu116 (True)
- Tensorflow version (GPU?): 2.9.2 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: GPU 0: Tesla T4
- Using distributed or parallel set-up in script?: no

### Who can help?

@stas00

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Trying to run the following in colab [link](https://colab.research.google.com/drive/1zHogV6VnqGPV5LoQoDzTy9D2OafYpWkz?usp=sharing)

It is pretty stock example of running summarization training with bart-large-cnn with deepspeed turned. I was hoping to train t5-3b on A100 but running into issues with smaller model on T4.

### Expected behavior

Running this on a Tesla T4 with 16 GB of GPU RAM one would expect that bart-large-cnn would work with Zero stage 3 enabled. Instead I get OutOfMemoryError: CUDA out of memory.
Zero stage 3 seems to be enabled per following output in the log file:
...
[2023-01-14 18:34:57,363] [INFO] [config.py:1024:print]   zero_enabled ................. True
[2023-01-14 18:34:57,364] [INFO] [config.py:1024:print]   zero_optimization_stage ...... 3

Comments:
Really, this is a question for the Deepspeed Issues since HF Trainer only provides the integration, but why do you expect that this is supposed to fit into 16GB? Colab barely has any cpu ram so very often there isn't really any free memory to offload to. 

But let's figure out the math first and then it'd be easy to manage expectations.

How many params does this model have? e.g. if the checkpoint is saved in half precision - from `1.63 GB` model size it'd mean you'd need `1.63*9=14.67` i.e. 15GB just for the weights. If it were in fp32, than half of that and then it should fit w/o using zero at all.

You typically need `n_params * 18` just for the weights in mixed precision training. And then more for activations.

Also I don't see where you override the batch size - you should set it to `1` to start with and only increase that if you don't OOM.

p.s. to count the unique params:
```
sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
```

A bit more data. 
When using the standard GPU and high RAM environment before any training I see:
Gen RAM Free: 53.5 GB  | Proc size: 93.4 MB
GPU RAM Free: 15109MB | Used: 0MB | Util   0% | Total 15109MB

When explicitly setting the batch size to 1 it works even on the smaller GPU (T4). Once the training starts it keeps the GPU RAM utilization at about 7GB from what I can see:
GPU RAM Free: 8067MB | Used: 7042MB | Util  47% | Total 15109MB

For batch size 2:
GPU RAM Free: 6407MB | Used: 8702MB | Util  58% | Total 15109MB

For batch size 4:
GPU RAM Free: 1975MB | Used: 13134MB | Util  87% | Total 15109MB

At batch size 6 we are hitting the limits:
GPU RAM Free: 157MB | Used: 14952MB | Util  99% | Total 15109MB

So, when using auto for batch size, deepspeed determines optimal batch size to be 8 and runs out of CUDA memory. 
This is definitely not a transformer issue. 

I guess lesson learned is to ramp up the batch size as opposed to assuming that deepspeed will calculate the optimal size.

I'm glad to see you sorted it out, @xpact 

> So, when using auto for batch size, deepspeed determines optimal batch size to be 8 and runs out of CUDA memory.

The `auto` values in the ds config file are used differently depending on their key. their main purpose is to avoid situations where the command line args and ds_config mismatch, so basically they are just substituted with the command line args.

Only `auto`'s from `zero_optimization` section is actually "optimized" to the model size - this is the only exception.

Each `auto` key is documented here: https://huggingface.co/docs/transformers/main/main_classes/deepspeed#deepspeed-trainer-integration and this is exclusively HF Trainer feature (and of Accelerate) - i.e. deepspeed has no idea what to do with `auto` values.

-------------------

In general the batch size is typically one of the most importants hparam - and you always want to define it explicitly and not rely on any defaults.

You can also add `skip_memory_metrics=False` to your training args and it'll print you the full memory usage stats at the end of each run.