Anatomy of Model's Memory
The components on GPU memory are the following:
1. model weights
2. optimizer states
3. gradients
4. forward activations saved for gradient computation
5. temporary buffers
6. functionality-specific memory
A typical model trained in mixed precision with AdamW requires 18 bytes per model parameter plus activation memory.
For inference there are no optimizer states and gradients, so we can subtract those. And thus we end up with 6 bytes per model parameter for mixed precision inference, plus activation memory.
Let's look at the details.
Model Weights
4 bytes * number of parameters for fp32 training
6 bytes * number of parameters 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
