Multi GPU training for Transformers with different GPUs

user_107726 0
Viewed 1586

I want to fine tune a GPT-2 model using Huggingface’s Transformers. Preferably the medium model but large if possible. Currently, I have a RTX 2080 Ti with 11GB of memory and I can train the small model just fine.

My question is: will I run into any issues if I added an old Tesla K80 (24GB) to my machine and distributed the training? I cannot find information about using different capacity GPUs during training and issues I could run into.

Will my model size limit essentially be sum of all available GPU memory? (35GB?)

I’m not interested in doing this in AWS.

huggingface-transformers machine-learning
1 Answers

You already solved your problem. That's great. I would like to point out a different approach and address a few questions.

Will my model size limit essentially be sum of all available GPU memory? (35GB?)

This depends on the training technique you use. The standard data parallelism replicates the model, gradients and optimiser states to each of the GPUs. So each GPU must have enough memory to hold all these. The data is splitted across the GPUs. However, the bottleneck is usually the optimiser states and the model not the data.

The state-of-the-art approach in training is ZeRO. Not only the dataset, but also the model parameters, the gradients and the optimizer states are splitted across the GPUs. This allows you to train huge models without hitting OOM. See the nice illustration below from the paper. The baseline is the standard case that I mentioned. They gradually split optimizer states, gradients and model parameter accross the GPU's and compare the memory usage per GPU.

enter image description here

The authors of the paper created a library called DeepSpeed and it is very easy to integrate it with huggingface. With that I was able to increase my model size from 260 Million to 11 Billion :)

If you want to understand in detail how it works, here is the paper: https://arxiv.org/pdf/1910.02054.pdf

More information on integrating DeepSpeed with Huggingface can be found here: https://huggingface.co/docs/transformers/main_classes/deepspeed

PS: There is a the model parallelism technique in which each GPU trains different layers of the model but it lost its popularity and is not being actively used.

user_80105680
Related