Solving CUDA Out of Memory Errors in Llama Fine-Tuning

The Anatomy of a CUDA OOM Error

When you trigger a fine-tuning script for Llama 3.1 70B, your GPU memory is not just holding the model weights. The total VRAM consumption is a sum of four distinct categories: model weights, optimizer states, gradients, and activations. Understanding this breakdown is the first step toward a successful run.

Model weights are the most obvious. A 70B model in 16-bit precision (BF16) requires roughly 140GB of VRAM just to sit in memory. However, the real killer is the optimizer states. If you are using the standard AdamW optimizer, it stores a momentum and a variance term for every single trainable weight, and mixed-precision training keeps an FP32 master copy of the weights on top. Hugging Face puts that at 12 bytes per parameter [1], six times what the BF16 weights themselves cost. Set against 2 bytes of weights and 2 bytes of gradients, optimizer states therefore account for up to 75% of the total memory footprint in a non-quantized full fine-tune.

Activations are the next hurdle. These are the intermediate values stored during the forward pass to calculate gradients during the backward pass. Their size scales linearly with batch size, and with context length linearly under Flash Attention but quadratically in a standard attention implementation that materializes the full score matrix. If you are pushing Llama 3.1's 128k context window, your activations will likely exceed the capacity of an 80GB H100 before you even finish the first forward pass. To manage this, we look at three primary levers: precision, sharding, and recomputation.

• Model Weights: 2 bytes per parameter (BF16).
• Optimizer States: 8 bytes per parameter for AdamW's two FP32 moments, 12 once the FP32 master weights are counted.
• Gradients: 2-4 bytes per parameter.
• Activations: Dependent on sequence length and batch size.

Quantization and PEFT: The 4-Bit Revolution

If you are working with limited hardware or trying to maximize efficiency on sovereign cloud nodes, Parameter-Efficient Fine-Tuning (PEFT) is your best friend. Specifically, LoRA (Low-Rank Adaptation) and its quantized cousin, QLoRA, have changed the economics of fine-tuning.

Low-Rank Adaptation (LoRA) Fundamentals

LoRA works by freezing the original model weights and only training a small set of adapter layers. This drastically reduces the number of gradients and optimizer states you need to store. QLoRA takes this further by quantizing the frozen base model to 4-bit precision. Hugging Face estimates a QLoRA fine-tune of Llama 3.1 70B at around 48GB of VRAM, against roughly 500GB for a full 16-bit fine-tune [4], which puts the job on a single 48GB card rather than a multi-GPU node. That matches the QLoRA authors' own headline result: they finetune a 65B parameter model on a single 48GB GPU while preserving full 16-bit finetuning task performance [3]. Predicting memory overhead before you launch is still worth the effort, because adapter rank and the target module list move the number.

Common Quantization Pitfalls

However, do not fall into the trap of thinking LoRA is a silver bullet. If your task requires deep structural changes to the model's knowledge base, LoRA might underperform. For most instruction-tuning and domain-adaptation tasks, it is the professional standard. When implementing this, ensure you are using Flash Attention 3. As noted in the 2024 technical release from the FlashAttention team, version 3 runs 1.5 to 2.0 times faster than FlashAttention-2 in FP16 on Hopper GPUs and reaches up to 740 TFLOPS, while its tiling approach avoids materializing the full attention matrix and so keeps attention memory linear in sequence length [2].

Common Mistake: Forgetting to target all linear layers in your LoRA config. Many developers only target the 'q_proj' and 'v_proj' layers. While this saves memory, it limits the model's ability to learn complex patterns. Targeting all linear layers (gate_proj, up_proj, down_proj) provides better results with only a marginal increase in VRAM usage.

Distributed Training: FSDP vs. DeepSpeed ZeRO

When a single GPU is not enough, you must shard your model across multiple nodes. This is where Fully Sharded Data Parallel (FSDP) and DeepSpeed ZeRO come into play. These protocols are essential for scaling Llama 3.3 70B across a sovereign GPU cluster.

DeepSpeed ZeRO (Zero Redundancy Optimizer) offers three stages of optimization. ZeRO-1 shards optimizer states, ZeRO-2 shards gradients, and ZeRO-3 shards the model weights themselves. For Llama 70B, ZeRO-3 is often mandatory. It ensures that no single GPU holds the entire model, instead fetching the necessary shards over the interconnect (like NVLink) just in time for computation. This allows you to scale to models that are technically larger than the memory of any individual GPU in your cluster.

PyTorch FSDP is the modern alternative and is often preferred for its native integration with the PyTorch ecosystem. FSDP implements a similar sharding logic but offers more granular control over wrapping policies. By wrapping specific layers of Llama, you can control exactly how the model is partitioned. According to the PyTorch 2025 documentation, FSDP2 has introduced significant improvements in communication-computation overlap, which is critical when training on high-performance European infrastructure where low-latency networking is a bottleneck.

The Hidden VRAM Eaters: Context and Checkpointing

Even with QLoRA and FSDP, you can still hit an OOM if you don't manage your activations. The most effective tool here is Gradient Checkpointing. Instead of storing all intermediate activations during the forward pass, gradient checkpointing discards them and re-calculates them during the backward pass. This is a classic engineering trade-off: you save a massive amount of VRAM at the cost of roughly 30% more compute time. If you are hitting OOM at the start of the backward pass, this is usually the fix.

Another critical factor is the KV Cache. During fine-tuning, especially with long sequences, the memory required to store Key-Value pairs can explode. Using Grouped-Query Attention (GQA), which is native to Llama 3, helps, but you should also consider PagedAttention if you are doing any form of interactive evaluation during your training loops. This prevents memory fragmentation, which is a silent killer of long-running training jobs.

Finally, look at your Micro-Batch Size. Many developers try to set a large batch size to speed up training. This is a mistake. Instead, set your per_device_train_batch_size to 1 or 2 and use gradient_accumulation_steps to reach your desired effective batch size. This decouples your memory usage from your optimization logic, allowing you to train with large effective batches even on 24GB or 40GB cards. At Lyceum, we recommend a 'start small, scale slow' approach to batching to ensure stability before pushing the limits of the hardware.

Sovereign Infrastructure for Llama Scaling

The technical fixes mentioned above are only half the battle. The infrastructure you run on determines your ultimate performance and shapes your compliance posture, but a hosting provider cannot confer compliance on its customer: controller obligations such as lawful basis, transparency and data-subject rights remain yours. For European enterprises, fine-tuning Llama on US-based hyperscalers introduces data residency risks and potential latency issues. Lyceum provides a sovereign alternative, with high-performance European data centers in Spain, Paris and the Nordics that are built for these workloads.

Instances come with dashboard, CLI, API, Jupyter and VS Code access, and the support level is agreed per contract and typically set during the PoC, so you can focus on your model rather than the underlying driver versions or interconnect bottlenecks. When you run a Llama 3.3 fine-tuning job on Lyceum, your workload runs in European data centers and bills per second with no base fee. Engineering excellence is not just about the code you write; it is about the sovereignty of the environment where that code executes.

Sources

[1] Hugging Face: Optimizing Memory Usage for Training LLMs; [2] FlashAttention-3: Fast and Accurate Attention with Asynchrony; [3] arXiv: QLoRA, Efficient Finetuning of Quantized LLMs; [4] Hugging Face: Llama 3.1, 405B, 70B and 8B with Multilinguality and Long Context