Skip to content
Ashish.
All posts
Diagram comparing memory usage between Full Fine-Tuning and PEFT showing optimizer states and activation checkpoints.

Full Fine-Tuning vs PEFT: The Memory Math

An examination of the memory math behind full fine-tuning versus Parameter-Efficient Fine-Tuning (PEFT), covering optimizer states and activation checkpointing.

By Ashish KumarPart 2 of Fine-Tuning LLMs Hands-On

When engineers attempt to fine-tune large language models (LLMs), they frequently hit a hard wall: out-of-memory (OOM) errors. The common intuition is that "bigger models need more GPU," but this is incomplete. The reality is that training memory consumption is driven by three distinct data structures: model weights, gradients, and optimizer states. Understanding the mathematical relationship between these components explains why Full Fine-Tuning (FFT) is often impossible on consumer hardware, while Parameter-Efficient Fine-Tuning (PEFT) makes it feasible.

This analysis dissects the memory mechanics of both approaches, focusing on the role of optimizer states and activation checkpointing. This is Part 2 of the Fine-Tuning LLMs Hands-On series.

The Memory Budget of Full Fine-Tuning

In full fine-tuning, every parameter in the model is updated. To do this, the GPU must store three things simultaneously in High Bandwidth Memory (VRAM):

  1. Model Weights: The current values of the parameters.
  2. Gradients: The partial derivatives calculated during backpropagation.
  3. Optimizer States: Auxiliary variables used by the optimization algorithm (typically AdamW) to adjust learning rates.

For a model using mixed precision (FP16 or BF16), weights and gradients each take 2 bytes per parameter. AdamW, however, maintains two additional moments for each parameter: the first moment (mean) and the second moment (uncentered variance). These are stored in FP32 (4 bytes) for numerical stability.

Therefore, the memory footprint for a single parameter in FFT with AdamW is: 2 bytes (weights)+2 bytes (gradients)+4 bytes (moment 1)+4 bytes (moment 2)=12 bytes2 \text{ bytes (weights)} + 2 \text{ bytes (gradients)} + 4 \text{ bytes (moment 1)} + 4 \text{ bytes (moment 2)} = 12 \text{ bytes}

If you use 8-bit optimizers (like AdamW 8-bit), you can reduce the optimizer state to 1 byte per moment, dropping the total to 6 bytes per parameter. However, the dominant factor remains: training requires roughly 4x to 6x the VRAM of the model's static weight size.

For a 7-billion parameter model (7B) stored in BF16:

  • Static Weights: 7×109×2 bytes≈14 GB7 \times 10^9 \times 2 \text{ bytes} \approx 14 \text{ GB}
  • Training Memory (AdamW): 7×109×12 bytes≈84 GB7 \times 10^9 \times 12 \text{ bytes} \approx 84 \text{ GB}

Add activation memory (discussed below), and a 7B model typically requires an A100 80GB or multiple A100s just to start training, excluding the overhead for the context window.

Technical diagram showing a bar chart of memory usage for Full Fine-Tuning. Categories : Weights (2 bytes), Gradients (2 bytes), Optimizer State M1 (4 bytes), Optimizer State M2 (4 bytes). Style : Clean vector illustration, dark mode background, bright blue and orange bars, mi…

The Mechanism of PEFT (LoRA)

Parameter-Efficient Fine-Tuning methods like Low-Rank Adaptation (LoRA) change the memory equation by freezing the pre-trained weights. Instead of updating the full weight matrix W∈Rd×kW \in \mathbb{R}^{d \times k}, LoRA decomposes the update ΔW\Delta W into two smaller matrices AA and BB:

ΔW=B×A\Delta W = B \times A

Where B∈Rd×rB \in \mathbb{R}^{d \times r} and A∈Rr×kA \in \mathbb{R}^{r \times k}, and rr is the "rank" (typically 4, 8, or 16), which is much smaller than dd or kk (often 4096 or higher).

The number of trainable parameters drops from d×kd \times k to r(d+k)r(d + k). For a typical attention layer with d=4096,k=4096d=4096, k=4096 and rank r=8r=8:

  • Original params: $16,777,216$
  • LoRA params: 8×(4096+4096)=65,5368 \times (4096 + 4096) = 65,536

This represents a reduction of roughly 256x in trainable parameters for that layer. Since optimizer states are only stored for trainable parameters, the memory required for AdamW drops proportionally. The frozen weights still occupy VRAM, but they do not require gradient or optimizer state storage.

The Activation Trap

A common misconception is that PEFT eliminates the need for large GPUs entirely. This is false because PEFT does not reduce activation memory.

Activations are the intermediate outputs computed during the forward pass. They must be stored in VRAM so that the backward pass can calculate gradients. Activation memory scales linearly with:

  1. Batch Size: More samples per step = more activations.
  2. Sequence Length: Longer context windows = more tokens = more activations.
  3. Model Depth: Deeper models have more layers, each storing activations.

For a 7B model, activation memory can easily exceed 20–30 GB even if the model weights are frozen. If you attempt to use a large batch size or long context window with LoRA on a single consumer GPU (e.g., 24GB VRAM), you will still OOM, not because of optimizer states, but because there is no room to store the activations for the forward pass. This is precisely why activation checkpointing is critical even for PEFT; without it, the dynamic nature of activation memory can quickly overwhelm consumer-grade hardware, regardless of how few parameters you are actually training.

Activation Checkpointing (Gradient Checkpointing)

To solve the activation bottleneck, engineers use activation checkpointing (also known as gradient checkpointing or recomputation). This technique trades compute time for memory space.

Instead of storing all activations for the entire forward pass, the framework stores only a subset of "checkpoints." During the backward pass, when a gradient is needed for a layer that wasn't checkpointed, the framework re-runs the forward pass for that specific segment to recompute the activations.

The memory savings are significant: by checkpointing every L\sqrt{L} layers (where LL is the total number of layers), activation memory usage becomes proportional to the square root of the number of layers (O(L)O(\sqrt{L})) rather than linearly proportional. However, this is a trade-off choice; checkpointing more frequently reduces memory further but increases compute overhead, while checkpointing less frequently saves compute but uses more memory. This comes at a cost: the forward pass is executed twice for non-checkpointed segments, increasing training time by approximately 20–30%.

Technical flowchart illustrating activation checkpointing. Show a neural network layer sequence. Forward pass highlights stored checkpoints. Backward pass shows recomputation arrows for intermediate layers. Style : Clean vector illustration, dark mode background, green for sto…

Practical Synthesis

Let's compare the VRAM requirements for fine-tuning a 7B model (BF16) on a single GPU with a batch size of 1 and sequence length of 2048.

Full Fine-Tuning (AdamW):

  • Weights: 14 GB
  • Gradients: 14 GB
  • Optimizer States: 56 GB (4 bytes each for the two Adam moments)
  • Activations: ~15 GB (estimated)
  • Total: ~99 GB (Requires multiple A100 80GB GPUs)

LoRA (Rank 8) + Activation Checkpointing:

  • Frozen Weights: 14 GB (read-only)
  • LoRA Weights: Negligible (~0.5 GB)
  • LoRA Gradients & Optimizer States: ~1 GB (tiny parameter count)
  • Activations: ~8 GB (reduced by checkpointing)
  • Total: ~23.5 GB (Fits on RTX 3090/4090 24GB)

The math reveals that PEFT is not just about saving disk space; it is about shifting the memory bottleneck from the optimizer state (which is fixed per parameter) to the activation memory (which is dynamic per batch). By combining LoRA's parameter reduction with activation checkpointing, engineers can fit large models into consumer-grade hardware, albeit with a slower training loop due to recomputation. For implementation details on these optimizations, refer to the Hugging Face PEFT documentation and PyTorch Gradient Checkpointing docs.

Conclusion

The choice between full fine-tuning and PEFT is ultimately a memory constraint problem. FFT requires massive VRAM because it stores optimizer states for every parameter. PEFT reduces the trainable parameter count, drastically lowering optimizer memory, but leaves activation memory as the primary constraint. Activation checkpointing then manages that remaining constraint by trading compute for space. Understanding this triad—weights, optimizer states, and activations—is essential for selecting the correct hardware and hyperparameters for LLM training.

Common Pitfalls

  1. Ignoring Activation Memory for LoRA: Engineers often assume that because LoRA reduces parameter count, it eliminates VRAM issues. They fail to account for activation memory, which scales with batch size and sequence length, leading to unexpected OOM errors on consumer GPUs.
  2. Assuming 8-bit Optimizers Solve All Issues: While 8-bit optimizers reduce memory footprint, they introduce computational overhead and may not always be compatible with all hardware or frameworks. Relying on them exclusively can create bottlenecks elsewhere in the training loop.
  3. Over-Checkpointing: Applying checkpointing too aggressively (e.g., checkpointing every single layer) can significantly slow down training without providing necessary memory savings, especially if the GPU has ample VRAM.

Practical Takeaways

  • FFT is 4-6x Model Size: Full fine-tuning requires approximately 4 to 6 times the VRAM of the model's static weight size due to gradients and optimizer states.
  • LoRA Reduces Optimizer State, Not Activations: LoRA drastically lowers the memory needed for optimizer states by reducing trainable parameters, but activation memory remains a critical constraint.
  • Checkpointing Trades Compute for Memory: Activation checkpointing reduces memory usage by recomputing activations, but it increases training time. Tune the checkpoint frequency based on your available VRAM and time constraints.
  • Monitor Context Length: Activation memory scales linearly with sequence length. Keep an eye on context window sizes when sizing GPUs for PEFT workflows.

FAQ

Does LoRA reduce activation memory? No. LoRA reduces the number of trainable parameters, which lowers the memory required for optimizer states and gradients. However, activation memory, which is determined by the model architecture, batch size, and sequence length, remains unchanged unless activation checkpointing is used.

How much VRAM do I need for a 7B model with LoRA? With LoRA and activation checkpointing, a 7B model typically requires around 20–24 GB of VRAM for training with a batch size of 1 and a moderate sequence length. This allows it to fit on consumer GPUs like the RTX 3090 or 4090.

What is activation checkpointing? Activation checkpointing (or gradient checkpointing) is a technique that reduces memory usage by storing only a subset of intermediate activations during the forward pass. During the backward pass, the missing activations are recomputed on the fly, trading increased compute time for reduced memory consumption.

Related posts