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 = 109 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
- 10 GB of memory (RTX 3080 10 GB)?
- 12 GB of memory (RTX 3080 12 GB and RTX 3080 Ti)?
- 16 GB of memory (RTX 4080)?
- 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.