Memory and its Importance for Large Models

While processing power from CPUs and GPUs is a primary concern, the amount of available memory often becomes the first practical bottleneck when working with large models. A processor is useless if it cannot access the data and model parameters it needs to compute. Understanding memory requirements is therefore essential for both training and deploying modern AI systems.

When we talk about memory in the context of AI hardware, we are primarily referring to two types:

1. System RAM (Random Access Memory): This is the main memory of the server, directly accessible by the CPU. It's used for holding the operating system, data preprocessing scripts, and entire datasets that are waiting to be fed to the accelerator.
2. GPU VRAM (Video RAM): This is high-speed memory located directly on the GPU card. It is significantly faster than system RAM and is where all the action happens for accelerated deep learning. The GPU can only operate on data that is present in its own VRAM.

For GPU-accelerated workloads, VRAM is almost always the more significant constraint. A training job will fail if its components cannot fit into the VRAM of a single GPU, or across the pooled VRAM in a multi-GPU setup.

What Occupies VRAM During Training?

During the training of a neural network, several components must reside in VRAM simultaneously. The total memory footprint is much larger than just the model itself.

The primary consumers of GPU VRAM during a model training cycle. The total required memory significantly exceeds the size of the model parameters alone.

Let's break down these components:

• Model Parameters: These are the weights and biases of the neural network. The size is determined by the number of parameters and the precision used to store them (e.g., 32-bit floating point, or FP32).
• Gradients: During backpropagation, a gradient must be calculated for every parameter in the model. Therefore, the memory required for gradients is typically equal to the memory required for the model parameters.
• Optimizer States: Modern optimizers like Adam maintain state to adapt the learning rate for each parameter. Adam, for example, stores two moving averages for each parameter (momentum and variance), effectively doubling the memory footprint of the model parameters again.
• Activations: These are the intermediate outputs of each layer that are computed during the forward pass. They need to be stored in memory to be used for gradient calculation during the subsequent backward pass. The size of the activations is dynamic and depends on the batch size and the sequence length of the input data.
• Input Data: The batch of data currently being processed must also be loaded into VRAM.

A Practical Memory Calculation

To make this tangible, let's estimate the VRAM needed to train a 7-billion parameter language model using the Adam optimizer and standard 32-bit precision (FP32), where each parameter requires 4 bytes.

1. Model Parameters: $$7 \text{ billion parameters} \times 4 \text{ bytes/parameter} = 28 \text{ GB}$$
2. Gradients: The gradients have the same size as the parameters. $$7 \text{ billion parameters} \times 4 \text{ bytes/parameter} = 28 \text{ GB}$$
3. Optimizer States (Adam): Adam stores two states per parameter. $$2 \times 7 \text{ billion parameters} \times 4 \text{ bytes/parameter} = 56 \text{ GB}$$

Just for these three components, the total VRAM required is:

$$28 + 28 + 56 = 112 \text{ GB}$$

This calculation doesn't even include the memory for activations or the CUDA kernel overhead, yet it already exceeds the capacity of high-end GPUs like the NVIDIA A100 80GB. This is why you frequently hear about "Out of Memory" (OOM) errors. An OOM error simply means you tried to allocate more data to the GPU's VRAM than it has available. This could be because the model is too large, the batch_size is too high (which increases activation memory), or a combination of factors.

Memory for Inference

When a model is deployed for inference (i.e., making predictions), the memory requirements are much lower. You no longer need to store gradients or optimizer states. The primary memory consumer is the model's parameters themselves, plus the activations for the current input.

For our 7B model, the base memory for inference would be the 28 GB for the model weights, plus memory for the activations of a single inference request. This is far more manageable and explains why it's possible to run inference for a model on a GPU that would be too small to train it from scratch.

Understanding this breakdown is a prerequisite for optimization. Techniques like mixed-precision training (using 16-bit floats) or choosing a different optimizer can drastically reduce this memory footprint, topics we will cover in a later chapter. For now, the main takeaway is that when planning your infrastructure, the model size and training configuration directly determine your minimum memory requirements.