Batch size < GPU number when training with Trainer and deepspeed.
Author: zhaowei-wang-nlp
Pull request: no

# 🚀 Feature request
Hi, I am fine-tuning T5-11b using Trainer with deepspeed feature. I use the deepspeed zero3 stage to split T5-11b and its gradients to different GPUs. However, when I try to use multi GPUs, I found that the argument `per_device_train_batch_size` must be an integer, which means it at least is 1. So when I use more GPUs, the batch size must increase at the same time, which will cost must more GPU memory. Thus, it turns out that I can't fine-tune T5-11b with 2, 4 or 8 A100 (40G) GPUs. So, in general, deepspeed feature doesn't solve the memory issue if the model's size is similar to or larger than the memory of one of the multi GPUs.
So, I request the feature that batch size is smaller than the number of GPUs like train batch size of 2 on 4 GPUs.
here is the link to the  argument`per_device_train_batch_size`
https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L122

## Motivation

Fine-tune an extremely large model with a few small memory GPUs.

## Your contribution
I think the deepspeed package supports this feature already. So, adding this feature to Trainer is not hard.


Comments:
cc @sgugger 

cc @stas00 for DeepSpeed ;-)

@zhaowei-wang98, could you please try again explaining what is the problem that you're running into? Please show the actual command line / config you're running as I have a hard time understanding your Issue.

> So when I use more GPUs, the batch size must increase at the same time, which will cost must more GPU memory

There is no such requirement. 

In general there is no problem running T5-11B on A100 (40GB) w/ Deepspeed ZeRO-3 - or at least it worked last time I run it - It was done already more than a year ago, perhaps have a look at this old thread https://github.com/huggingface/transformers/issues/9996 and then if you're still stuck tell us more details about your particular setup?

> 

Hi @stas00,
I am trying to fine-tune t11-3b without CPU offload. So, all the parameters in the model and momentum in the optimizer are loaded on the GPUs. I do this because I found it is very slow to use CPU offload (I have 500k data with an average length of 32 for both input and output). 
In other words, I deleted:
"offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
in the deepspeed configuration file: https://github.com/huggingface/transformers/blob/main/tests/deepspeed/ds_config_zero3.json

In contrast, your old thread #9996 used CPU RAM to store the model.

OK, but I still can't help you since I don't know how to reproduce your issue as you gave no specific instructions to do so, nor you shared anything about your setup other than the type of GPUs.

But I can probably do some guessing:

### understanding the memory requirements

To train t5-11b you need at least `11*18=200`GB of memory just for the optim/grads/weights (I assume mixed precision) plus memory for activations and temps, so let's say roughly 240GB. 

With 40GB GPUs, that means at least 6 gpus. So it should be possible to load it on 8x 40GB gpus using deepspeed w/o any offload.

### use the sharded checkpoint

Also I recommend for you to switch to the sharded version of t5-11b which I have just [made](https://github.com/huggingface/transformers/issues/16884), by passing to the trainer:  `--model_name_or_path t5-11b --model_revision sharded` and use `huggingface@main` as this feature hasn't yet been released.

Because if you don't shard you would need 44GB of CPU memory per process, just to load the checkpoint (deepspeed shards it directly to gpus). And with 8 gpus you'd need 352GB of CPU memory just to load 8 checkpoints concurrently.

I think with 10GB shards some 100GB of CPU memory should be enough to load the checkpoints concurrently in 8 processes, but then there are extras to copy things and temps.

----------------

I'd be very happy to help you sort it out, but you need to help me first. To continue please be very specific:

1. here is my hardware setup
2. here is my software setup
3. here is my command line (using public data) and ds config file to reproduce the problem with 
4. here is the traceback

Thank you!


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.