Memory Management › Memory Estimation

Estimating total memory before running a training job lets you choose the right GPU configuration, set batch sizes appropriately, and avoid wasting hours waiting for an OOM crash that was predictable from the start. A good estimate also helps you make principled decisions about which memory reduction techniques to apply and in what order.

For a standard mixed-precision training setup with Adam, the total memory $M$ breaks down as:

$$M = \underbrace{2P}_{\text{FP16 params}} + \underbrace{4P}_{\text{FP32 grads}} + \underbrace{8P}_{\text{Adam states}} + \underbrace{M_{\text{act}}}_{\text{activations}} + \underbrace{M_{\text{other}}}_{\text{overhead}}$$

where $M_{\text{other}}$ covers CUDA context, framework overhead, and temporary buffers. The CUDA runtime itself requires roughly 300-500 MB when initialized. PyTorch maintains a memory caching allocator that retains freed memory for reuse; its overhead depends on the number and sizes of allocations. Temporary workspace for cuBLAS and cuDNN operations can require several hundred megabytes more. In practice, $M_{\text{other}}$ typically adds 500 MB to 2 GB depending on the framework version and GPU model.

The activation memory $M_{\text{act}}$ depends heavily on whether gradient checkpointing is used:

$$M_{\text{act}} = \begin{cases} L \cdot B \cdot T \cdot \left(34d + 5hT\right) & \text{without checkpointing} \\ B \cdot T \cdot \left(34d + 5hT\right) & \text{with full checkpointing} \end{cases}$$

where the factor of 2 bytes per element (FP16) is incorporated into the constants 34 and 5.

Consider a GPT-2 style model with the following architecture:

• Vocabulary: $V = 50{,}257$
• Hidden dim: $d = 768$
• Layers: $L = 12$
• Heads: $h = 12$
• FF dim: $d_{ff} = 3{,}072$

The parameter count is approximately:

$$P \approx 50{,}257 \times 768 + 12 \times \left(4 \times 768^2 + 2 \times 768 \times 3{,}072\right) \approx 117M$$

In mixed precision with Adam:

$$M_{\text{non-act}} = 14 \times 117 \times 10^6 \text{ bytes} \approx 1.6 \text{ GB}$$

For a batch size of 8 and sequence length 512:

$$M_{\text{act}} \approx 12 \times 8 \times 512 \times \left(34 \times 768 + 5 \times 12 \times 512\right) \text{ bytes} \approx 2.1 \text{ GB}$$

Total estimated memory: roughly 3.7 GB, well within a single 16 GB GPU's budget. Scaling to a 7B model with batch size 16 and sequence length 2048 pushes this into the multi-hundred-gigabyte range, requiring distributed training.

One additional nuance in the estimation is memory fragmentation. GPU memory is allocated and freed in blocks of varying sizes. Over time, freed blocks of one size may not be reusable for allocations of a different size, causing apparent "wasted" memory that shows as free but is not contiguous enough for large allocations. PyTorch's CUDA memory allocator tries to mitigate this through block pooling and splitting, but fragmentation can cause an OOM even when the total free memory nominally exceeds the requested allocation. In practice, you should maintain a 10-20% safety margin in your memory estimates to account for fragmentation.