What size language model can you train on a GPU with x GB of memory?

Question

I'm trying to figure out what size language model I will be able to train on a GPU with a certain amount of memory. Let's for simplicity say that 1 GB = 10⁹ bytes; that means that, for example, on a GPU with 12 GB memory, I can theoretically fit 6 billion parameters, given that I store all parameters as 16-bit floats. However, in order to use a language model, you typically also need space for storing the input text and the activations of the current layer (and maybe also of the previous layer), and if you are going to train the model, you will typically need space to store the activations of all layers in order to be able to do backpropagation, and if you use an optimizer like Adam, you need space to store the running mean of the partial derivatives (of the loss function with respect to the various parameters, or in other words, the gradient), as well as the running mean of the squares of the partial derivatives.

So, given this complication, could someone tell me what size language models (that is, how many parameters) I will be able to train on a GPU with

1. 10 GB of memory (RTX 3080 10 GB)?
2. 12 GB of memory (RTX 3080 12 GB and RTX 3080 Ti)?
3. 16 GB of memory (RTX 4080)?
4. 24 GB of memory (RTX 3090 and RTX 3090 Ti)?

For example, Tim Dettmer mentioned in his blog that you should have at least 24 GB of memory if you do research on transformers. I'm guessing this translates roughly to a language model of a certain size.

Answer 1

Tldr; I’ve seen a good rule-of-thumb is about 14-18x times the model size for memory limits, so for a 10GB card, training your model would max out memory at roughly 540M parameters.

There is some really good information here: https://huggingface.co/docs/transformers/perf_train_gpu_one#anatomy-of-models-memory

Note that there are a ton of caveats, depending on framework, mixed precision, model size, batch sizes, gradient checkpointing, and so on. Just to summarize the above, rough memory requirements are: Model weights

• 4 bytes * number of parameters for fp32 training
• 6 bytes * number of params for mixed precision training.

Optimizer States

• 8 bytes * number of parameters for normal AdamW (maintains 2 states)
• 2 bytes * number of parameters for 8-bit AdamW optimizers like

bitsandbytes

• 4 bytes * number of parameters for optimizers like SGD with momentum (maintains only 1 state) Gradients:
• 4 bytes * number of parameters for either fp32 or mixed precision training (gradients are always kept in fp32) Other: Temporary memory, functionality specific memory, forward activations, and so on.