The Four Pillars of VRAM Consumption

To accurately estimate memory, you must break down the GPU's workload into four distinct categories. Each serves a different purpose during the forward and backward passes of training. Ignoring even one of these will lead to immediate OOM errors once the training loop begins.

Parameter Memory and Precision Formats

1. Model Weights (Parameters): This is the static portion of your memory footprint. If you are running a 7B parameter model in FP32 (4 bytes per parameter), the weights alone take up 28 GB. In BF16 or FP16 (2 bytes per parameter), this drops to 14 GB. This memory remains constant regardless of your batch size.

Gradient Memory Overhead

2. Gradients: During the backward pass, the GPU stores gradients for every trainable parameter. These typically match the precision of your weights. If you are training in mixed precision, gradients may be stored in FP32 to maintain numerical stability, even if the forward pass uses BF16. Hugging Face's memory accounting assumes 4-byte FP32 gradients while ZeRO's assumes 2-byte 16-bit gradients, so the gradient dtype follows the framework's convention.

Optimizer States and Activation Memory

3. Optimizer States: This is often the most overlooked memory hog. Standard optimizers like Adam or AdamW store moving averages of the gradients. For every parameter, Adam keeps two additional states (momentum and variance). In a typical mixed-precision setup, these states are stored in FP32, requiring 8 bytes per parameter. When combined with the FP32 master copy of the weights, the optimizer states consume 12 bytes per parameter, which is the K=12 term in the ZeRO paper's accounting.

4. Activations: Unlike weights and optimizer states, activations scale with your input data. Every layer in your network stores intermediate outputs during the forward pass to calculate gradients during the backward pass. Activation memory scales linearly with batch size and hidden dimension, and both linearly and quadratically with sequence length through the attention term. This is the primary variable you can manipulate to fit a model into a specific GPU.

  • Weights: Static, depends on parameter count and precision.
  • Gradients: Static, usually matches weight precision.
  • Optimizer States: Static, but significantly larger than weights (for Adam).
  • Activations: Dynamic, depends on batch size and sequence length.

The Math of Precision and Data Types

Precision is the primary lever for reducing the static memory footprint. Moving from FP32 to BF16 or FP8 changes the byte-per-parameter calculation fundamentally. NVIDIA's mixed-precision training documentation sets out the standard recipe: compute in 16-bit while the optimizer keeps a primary copy of the weights in FP32. BF16 is the usual choice for LLM training because, as Kalamkar and co-authors report, its range of representable values matches FP32 at 2 bytes per parameter.

When calculating your requirements, use the following byte values for each data type:

  • FP32 (Full Precision): 4 bytes per parameter.
  • FP16 / BF16 (Half Precision): 2 bytes per parameter.
  • INT8 (Quantized): 1 byte per parameter.
  • FP8 (New Standard): 1 byte per parameter (requires H100 or newer hardware).

Consider a 70B parameter model. In FP32, the weights alone require 280 GB, which nearly fills a quad-H100 node (320 GB) before any gradients, optimizer states or activations. By utilizing 4-bit quantization for fine-tuning (QLoRA), you can reduce the weight footprint to roughly 35 GB, allowing the model to fit on a single A100 (80GB) with ample room for activations and gradients.

However, remember that training requires more than just weights. A common mistake is assuming a 7B model fits on a 16GB GPU because the weights are only 14GB in BF16. Once you add gradients (14GB) and optimizer states (8 bytes per parameter for Adam's FP32 momentum and variance, before any FP32 master copy), you are well beyond the limit. This is why distributed training and sharding are essential for modern AI development.

Calculating Activation Memory for Transformers

Activations are the most volatile component of GPU memory. While weights are fixed, activations grow as you increase the batch size or the context window. For a standard Transformer architecture, the activation memory per layer can be estimated with the formula published by Korthikanti and co-authors in NVIDIA's work on activation recomputation:

Memory = B * S * H * (34 + (5 * S * A / H))

Where:

  • B: Batch Size
  • S: Sequence Length (Context Window)
  • H: Hidden Dimension Size
  • A: Number of Attention Heads

This formula accounts for the storage of query, key, and value matrices, the attention matrix, and the feed-forward network outputs. Notice that the sequence length (S) appears twice, once as a linear factor and once as a quadratic factor in the attention calculation. This is why doubling your context window has a much more dramatic impact on VRAM than doubling your batch size.

To mitigate activation memory, engineers often use Gradient Checkpointing. This technique discards activations during the forward pass and recomputes them during the backward pass. Chen and co-authors showed that this costs O(sqrt(number of layers)) memory to train a network of that depth, at the price of one extra forward pass per mini-batch. If you are hitting VRAM limits on a single node, enabling checkpointing is often the first and most effective step before reducing batch size.

Distributed Strategies: ZeRO and Sharding

When a single GPU cannot hold the model and its states, you must distribute the memory across a cluster. The Zero Redundancy Optimizer (ZeRO), popularized by Microsoft's DeepSpeed and now integrated into PyTorch FSDP, is the gold standard for this. ZeRO eliminates memory redundancy by partitioning states across all available GPUs.

There are three stages of ZeRO, each offering progressively more memory savings:

  1. ZeRO-1: Partitions only the optimizer states. This reduces the footprint significantly while maintaining the same communication overhead as standard Data Parallelism.
  2. ZeRO-2: Partitions optimizer states and gradients. This is the sweet spot for many training runs, as it clears a massive amount of VRAM with minimal performance impact.
  3. ZeRO-3: Partitions weights, gradients, and optimizer states. This allows you to train models that are far larger than the memory of any single GPU, as the weights are only gathered when needed for a specific layer's computation.

At Lyceum, these distributed strategies run on multi-GPU nodes and clusters in European data centers. Frameworks like DeepSpeed and FSDP handle the underlying communication primitives, and Lyceum's S3-compatible storage carries no ingress or egress charges.

A Practical Decision Framework

Before starting your next training run, use this framework to determine if your hardware is sufficient. This logic prevents the "trial-and-error" loop that plagues many ML teams.

Step 1: Calculate Static Memory. Multiply your parameter count by the bytes required for your chosen precision, then add gradients and optimizer states. ZeRO's accounting for mixed-precision Adam totals 16 bytes per parameter: 2 for the weights, 2 for the gradients, and 12 for the FP32 master copy, momentum and variance. If this number is greater than your total cluster VRAM, you must use ZeRO-3 or offloading.

Step 2: Estimate Activation Peak. Use the transformer activation formula with your target batch size and sequence length. If the sum of static memory and activations exceeds 90% of your GPU capacity, you are in the danger zone for OOM errors due to fragmentation.

Step 3: Apply Mitigation. If you are over the limit, apply these techniques in order:

  • Enable Gradient Checkpointing (full recomputation keeps only each layer's input, 2sbh bytes per layer instead of the full sbh(34 + 5as/h)).
  • Use ZeRO-2 or ZeRO-3 to shard states across GPUs.
  • Reduce Batch Size (and use Gradient Accumulation to maintain effective batch size).
  • Use 8-bit or 4-bit optimizers (like bitsandbytes) to reduce optimizer state footprint.

By following this structured approach, you ensure that your training is stable from the first epoch. In the context of European data sovereignty, where compute resources must be used responsibly and efficiently, this level of engineering rigor is not just a best practice, it is a requirement.

Sources

[1] Hugging Face Accelerate, Model memory estimator (the accelerate estimate-memory CLI); [2] Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, arXiv:1910.02054, source of the 16 bytes per parameter total and the K=12 optimizer term; [3] Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, arXiv:2205.05198, source of the sbh(34 + 5as/h) formula and the 2sbh recomputation figure; [4] Chen et al., Training Deep Nets with Sublinear Memory Cost, arXiv:1604.06174; [5] Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training, arXiv:1905.12322; [6] NVIDIA, Train With Mixed Precision; [7] NVIDIA Technical Blog, GPU Memory Essentials for AI Performance. All retrieved 3 August 2026.