Eliminating CUDA OOM: Expert Memory Management for LLMs

To solve memory issues, you must first quantify where the bytes are going. In a standard LLM workload, memory is partitioned into four distinct areas. Model Weights are the most obvious, typically requiring 2 bytes per parameter in FP16/BF16 precision. A 70B model in half-precision requires 140GB of VRAM just to load the weights into memory before a single token is processed.

Optimizer States and Gradients dominate during the training and fine-tuning phases. While weights take 2 bytes, AdamW optimizers require an additional 12 bytes per parameter for the momentum and variance buffers. This means a 7B model, which seems small, can easily exceed the 80GB limit of an NVIDIA H100 if the training loop is not optimized. Activations are stored during the forward pass to calculate gradients during the backward pass, and their size scales linearly with batch size and sequence length.

The KV Cache is the silent killer of inference performance. As the model generates tokens, it stores the Key and Value vectors for all previous tokens to avoid redundant computations. In standard attention mechanisms, this cache grows quadratically with sequence length. According to a 2025 analysis of transformer architectures, the KV cache can consume up to 30% of total VRAM in long-context windows, often triggering OOM errors mid-generation even if the initial load was successful.

Quantization is the most effective lever for reducing the memory footprint of model weights. By moving from 16-bit (BF16) to 8-bit (INT8) or 4-bit (NF4) precision, you can fit significantly larger models on existing hardware. The bitsandbytes library has popularized 4-bit NormalFloat (NF4) quantization, which allows a 70B model to run on two A100 80GB GPUs with room for a substantial KV cache.

• FP8 (8-bit Floating Point)

  Supported natively by NVIDIA H100 and H200 GPUs. It offers a middle ground with minimal accuracy loss and significant speedups in throughput.

• 4-bit (NF4/AWQ)

  Reduces memory usage by 75% compared to FP16. While there is a measurable impact on perplexity, it is often negligible for task-specific fine-tuning.

• GGUF/IQ Quantization

  Useful for edge deployments or CPU-offloading scenarios, though less common in high-performance data center environments.

When choosing a quantization level, consider the Quantization Tax. A 2025 report from the OpenLM research group indicated that while 4-bit quantization is efficient, it can lead to a 1-3% drop in benchmark accuracy for reasoning tasks. For finance or healthcare applications where precision is non-negotiable, FP8 or INT8 is often the safer production standard. Lyceum's orchestration engine automates this selection by matching the model's precision requirements with the available hardware interconnects to ensure no performance is left on the table.

Standard attention is memory-inefficient because it materializes the full attention matrix in VRAM. FlashAttention-3, released in late 2024 and widely adopted in 2025, solves this by using tiling and recomputation to keep the attention calculation within the GPU's SRAM. This reduces memory reads/writes and allows for much longer context windows without the linear memory scaling typically seen in older implementations.

For inference-heavy workloads, PagedAttention (the core of the vLLM library) is the industry standard. Traditional memory allocation for the KV cache is contiguous, leading to internal and external fragmentation. PagedAttention treats GPU memory like virtual memory in an operating system, dividing the KV cache into non-contiguous blocks. This eliminates nearly 100% of memory waste, allowing you to increase batch sizes by 2x to 4x on the same hardware.

1. Identify the bottleneck: Use nvidia-smi or PyTorch Profiler to determine if OOM occurs during model loading or during the forward pass.

2. Implement Gradient Checkpointing

   This technique trades compute for memory by discarding activations during the forward pass and re-calculating them during the backward pass. It can reduce activation memory by up to 70%.

3. Enable Grouped Query Attention (GQA)

   If you are training a custom model, GQA reduces the number of Key and Value heads, significantly shrinking the KV cache size without sacrificing much model quality.

When a single GPU is insufficient, distributed strategies are required. The DeepSpeed ZeRO (Zero Redundancy Optimizer) ecosystem provides a tiered approach to memory management. ZeRO-1 shards optimizer states, ZeRO-2 shards gradients, and ZeRO-3 shards the model weights themselves across the entire GPU cluster. This allows for the training of models with trillions of parameters that would be impossible to fit on any single node.

For European enterprises operating in regulated sectors, scaling across clusters must be handled with strict data sovereignty. Lyceum's sovereign cloud infrastructure ensures that these distributed workloads remain within specific geographic boundaries while providing the high-bandwidth interconnects (like InfiniBand or RoCE) necessary for ZeRO-3 to function without massive latency penalties. Fully Sharded Data Parallel (FSDP) in PyTorch offers a similar benefit, sharding parameters, gradients, and optimizer states to ensure that the memory load is perfectly balanced across your fleet.

Common mistakes in distributed setups include failing to account for CPU RAM offloading. While DeepSpeed allows offloading optimizer states to the CPU, this introduces a massive bottleneck in training speed. It should be viewed as a last resort for fitting a model onto limited GPU resources rather than a primary scaling strategy. Instead, prioritize efficient sharding and pipeline parallelism to keep the workload on the GPU's high-speed memory bus.