The Memory Game

When training large language models, they’re memory hogs.

A 7B parameter model in full precision? That’s 28 GB just for the weights. Add in the optimizer state (another 56 GB!), gradients (28 GB), and activations (20+ GB), and you’re looking at 130+ GB of memory.

Your RTX 4090 has 24 GB.

See the problem?

The good news: with the right tricks, you can train that 7B model on consumer hardware. This notebook is your guide to the memory optimization techniques that make it possible.

Understanding the Memory Problem¶

First, let’s break down where all that memory goes during training.

Think of GPU memory like a packed suitcase. You’ve got limited space, and you need to fit everything in. Here’s what’s taking up room:

The Memory Breakdown¶

Total GPU Memory During Training:
├── Model Weights        (~25-30%)  ← The model parameters themselves
├── Optimizer State      (~50-60%)  ← Momentum & variance (biggest offender!)
├── Gradients            (~25-30%)  ← One gradient per parameter
├── Activations          (~10-20%)  ← Saved outputs for backprop
└── Framework Overhead   (~5%)      ← PyTorch bookkeeping

Notice something? The optimizer state is typically the largest consumer of memory. Not the model itself!

Why? Because modern optimizers like AdamW keep track of two extra values per parameter: a momentum term and a variance term. That’s 2x the model size right there.

This is why a 7B parameter model needs way more than 28 GB of memory for training. It’s not just storing the weights—it’s storing all the machinery needed to update those weights effectively.

Let’s Do the Math¶

Numbers make this concrete. Let’s calculate the exact memory requirements for two models.

GPT-2 (124M parameters) - Full Fine-Tuning¶

Model weights (fp32):     124M params × 4 bytes = 496 MB
Optimizer (AdamW):        124M params × 8 bytes = 992 MB  ← 2x model size!
Gradients (fp32):         124M params × 4 bytes = 496 MB
Activations (batch=8):                            ~500 MB
Framework overhead:                               ~100 MB
─────────────────────────────────────────────────────────
Total:                                           ~2.6 GB

Not too bad! This fits comfortably on most GPUs.

Llama 7B - Full Fine-Tuning¶

Model weights (fp32):     7B params × 4 bytes = 28 GB
Optimizer (AdamW):        7B params × 8 bytes = 56 GB  ← Oof.
Gradients (fp32):         7B params × 4 bytes = 28 GB
Activations (batch=8):                          ~20 GB
Framework overhead:                             ~2 GB
─────────────────────────────────────────────────────
Total:                                         ~134 GB

Whoops! That’s not fitting on consumer hardware.

Breaking Down the Calculations¶

Why 4 bytes for fp32? Each 32-bit floating point number takes 32 bits = 4 bytes of memory.

Why 8 bytes for AdamW? AdamW stores two additional values per parameter (first moment and second moment), each in fp32. So that’s 4 bytes + 4 bytes = 8 bytes per parameter just for optimizer state.

What are activations? During the forward pass, we save the output of each layer. We need these saved values during backpropagation to compute gradients. More layers, longer sequences, and bigger batch sizes = more activations to store.

The good news? We can dramatically reduce these numbers with the right techniques.

Technique 1: Mixed Precision Training¶

This is the most impactful technique you can apply. It cuts memory usage roughly in half with just a few lines of code.

What is Mixed Precision?¶

Instead of using 32-bit floating point numbers (fp32) for everything, we use 16-bit numbers (fp16 or bf16) for most operations.

Think of it like this: you’re doing carpentry. Sometimes you need a micrometer for precise measurements. But most of the time? A ruler is fine. Mixed precision training uses the “ruler” (16-bit) for most work and pulls out the “micrometer” (32-bit) only when needed.

The Three Formats¶

fp16 (Float16): Traditional half precision. Small range, can overflow/underflow easily.

bf16 (BrainFloat16): Google’s format. Same range as fp32 but less precision. This is the sweet spot for modern GPUs (Ampere, Ada, Hopper). No overflow issues, no loss scaling needed.

Which should you use? If your GPU supports it (RTX 30-series or newer, A100, H100), use bf16. It’s simpler and more robust. Otherwise, fp16 works but requires loss scaling to prevent underflow.

Memory Savings¶

7B model in fp32:     7B × 4 bytes = 28 GB
7B model in bf16:     7B × 2 bytes = 14 GB
                                     ↓
                           50% reduction!

And training quality is essentially identical. You’re getting half the memory usage for free.

Technique 2: LoRA (Low-Rank Adaptation)¶

If mixed precision is a memory reducer, LoRA is a memory destroyer. In the best way.

Remember how I said the optimizer state is the biggest memory hog? LoRA solves this by freezing the base model and training only tiny adapter layers.

The Core Idea¶

Instead of updating all 7 billion parameters, we freeze them and add small “adapter” matrices that we train instead.

Think of it like this: you’ve got a massive reference book (the base model). Instead of rewriting the whole book, you add sticky notes (LoRA adapters) with corrections and additions. The book stays the same; only the notes change.

The Math¶

For each weight matrix W, LoRA adds two small matrices A and B:

Original:    W (full rank, millions of parameters)
LoRA adds:   W + (A × B)
             ↑    ↑   ↑
         frozen  rank r matrices
                 (tiny!)

If W is 4096×4096, that’s 16.7M parameters. But A and B with rank r=16? That’s only 4096×16 + 16×4096 = 131K parameters!

That’s 128x fewer parameters to train.

Memory Impact¶

Full Fine-Tuning (Llama 7B):
  Trainable params:   7,000,000,000
  Optimizer state:    56 GB  (8 bytes per param)
  Gradients:          28 GB  (4 bytes per param)

LoRA with r=16 (Llama 7B):
  Trainable params:   16,777,216  (0.24% of model!)
  Optimizer state:    134 MB      (418x reduction!)
  Gradients:          67 MB       (418x reduction!)

The base model stays frozen, so we only need optimizer state and gradients for those tiny adapter matrices.

It’s kind of absurd how well this works.

Technique 3: Gradient Accumulation¶

Okay, this one’s clever.

You know how larger batch sizes generally lead to more stable training? But bigger batches need more memory (you’re processing more samples at once).

Gradient accumulation lets you have your cake and eat it too.

The Trick¶

Instead of computing gradients and updating weights after every batch, we accumulate gradients across multiple small batches, then update once.

Effective batch size = per_device_batch_size × gradient_accumulation_steps

Memory usage = per_device_batch_size (not effective batch size!)

So if you can only fit batch_size=2 in memory, but you want the training stability of batch_size=32, you can do:

per_device_batch_size = 2
gradient_accumulation_steps = 16
# Effective batch size = 2 × 16 = 32

You get the benefits of a large batch size with the memory footprint of a small one.

The Trade-off¶

Pro: Train with larger effective batch sizes without OOM errors.

Con: Slightly slower (you’re doing more forward passes before each update).

But slower beats “doesn’t fit in memory” every time.

Technique 4: Gradient Checkpointing¶

This technique trades computation for memory. It’s brilliant in a slightly painful way.

The Problem¶

During the forward pass, we need to save the output of every layer. Why? Because during backpropagation, we need those saved values to compute gradients.

For a 32-layer transformer processing a batch of 8 sequences, that’s a LOT of saved activations. They can easily use 20+ GB of memory.

The Solution¶

What if we... didn’t save them?

Gradient checkpointing only saves activations at certain “checkpoint” layers (say, every 4th layer). During backpropagation, when we need an activation we didn’t save, we recompute it on the fly.

Think of it like taking notes during a lecture. You could transcribe everything (high memory), or you could write down key points and fill in the details later (low memory, more work).

Without Gradient Checkpointing:
  Forward pass:  Compute activations → Save all of them
  Backward pass: Use saved activations → Fast
  Memory: HIGH

With Gradient Checkpointing:
  Forward pass:  Compute activations → Save only checkpoints
  Backward pass: Recompute missing activations → Slower
  Memory: LOW

The Trade-off¶

You’re recomputing activations during backprop, so training is slower (typically 20-30% slower).

But the memory savings are huge: 50-80% reduction in activation memory.

When you’re memory-constrained, this is a lifesaver.

Technique 5: Model Quantization¶

Now we’re getting into the heavy artillery.

Quantization is like... aggressive compression for neural networks. Instead of storing weights in 16-bit or 32-bit precision, we use 8-bit or even 4-bit integers.

Sounds crazy, right? How can you possibly represent a neural network weight that might be 0.0123456789 using just 4 bits (16 possible values)?

The Magic¶

Modern quantization techniques (like NormalFloat4, or “nf4”) are calibrated to the distribution of neural network weights. Most weights cluster around zero, with a long tail. NF4 packs more precision where it matters and less where it doesn’t.

It’s like how JPEG compression works: throw away information humans won’t notice. Here, we throw away precision the model doesn’t need.

The Numbers¶

Yes, you read that right. A 7B model in 4-bit takes 3.5 GB. That fits on a laptop GPU.

The Catch¶

Quantized models can only be used for inference or as frozen base models for fine-tuning (like with LoRA).

You can’t train a quantized model directly. The low precision causes training to diverge.

But combined with LoRA (which adds trainable adapters on top), you get QLoRA: 4-bit base model + 16-bit LoRA adapters = magic.

Profiling Memory Usage¶

Before you can optimize memory, you need to measure it.

Here’s a simple profiling function that shows exactly where your memory is going. This is invaluable when debugging OOM errors or trying to squeeze more performance out of your GPU.

When Things Go Wrong: Debugging OOM Errors¶

The dreaded “CUDA out of memory” error. We’ve all been there.

Here’s your systematic debugging checklist.

The Quick Fixes (Try These First)¶

1. Cut your batch size in half

• If batch_size=8 OOMs, try batch_size=4

• Use gradient accumulation to maintain effective batch size

• This fixes 80% of OOM issues

2. Enable gradient checkpointing

• One line: model.gradient_checkpointing_enable()

• Saves 50-80% on activation memory

• Especially important for long sequences

3. Use gradient accumulation

• Smaller batches, same training dynamics

• Costs nothing but a bit of training time

4. Clear the CUDA cache

• torch.cuda.empty_cache()

• Defragments memory, sometimes helps

• Won’t free memory that’s actually in use

Common OOM Causes¶

The Nuclear Option¶

If nothing else works:

# QLoRA: 4-bit model + LoRA
quantization_config = BitsAndBytesConfig(load_in_4bit=True, ...)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config
)
model.gradient_checkpointing_enable()

# LoRA config
lora_config = LoraConfig(r=8, ...)  # Smaller rank if needed
model = get_peft_model(model, lora_config)

# Training config
training_args = TrainingArguments(
    per_device_train_batch_size=1,  # Minimum
    gradient_accumulation_steps=32,  # Large accumulation
    ...
)

This’ll fit almost anything, anywhere. It’s slow, but it works.

Putting It All Together: A Memory Optimization Strategy¶

Okay, we’ve covered a lot of techniques. How do you actually use them?

Here’s my recommended approach, in order. Think of it as a ladder—start at step 1, and only climb higher if you need to.

Step 1: The Essentials (Always Do This)¶

These are free wins. Apply them to every training run.

1. Mixed precision (BF16 or FP16)

• 50% memory reduction

• Often faster training

• Literally no downside on modern GPUs

2. LoRA (if training large models >1B params)

• 80-95% reduction in optimizer memory

• No speed penalty

• Quality is excellent for most tasks

3. Find your maximum batch size

• Start with batch_size=8

• If OOM, halve it

• If comfortable, double it

• Use gradient accumulation to hit your target effective batch size

Step 2: If You’re Still Running Out of Memory¶

4. Gradient checkpointing

• 50-80% reduction in activation memory

• 20-30% slower training

• Essential for long sequences

5. Gradient accumulation

• Lets you use smaller batches

• Maintains training stability

• Slight speed cost

Step 3: The Extreme Measures¶

6. 4-bit quantization (QLoRA)

• 75% reduction in model memory

• Small quality hit (1-2%)

• Can only be used with LoRA

7. CPU offloading (DeepSpeed ZeRO-3)

• Offloads optimizer state to CPU RAM

• 50-70% GPU memory reduction

• 60-80% slower training

• Use as last resort

The Rule¶

Start simple. Add complexity only when you hit limits.

Don’t start with QLoRA + checkpointing + accumulation + offloading just because you can. Start with BF16 + LoRA, see how far that gets you, and add techniques as needed.

Complete Memory Optimization Example¶

Here’s what a fully optimized training setup looks like. This is the configuration that lets you train 7B models on consumer GPUs.

Goal: Fine-tune Llama 7B on a 12 GB GPU (RTX 3060/3080)

The recipe:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
import torch

# Step 1: Load model in 4-bit (QLoRA)
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",              # NormalFloat4
    bnb_4bit_use_double_quant=True,         # Double quantization
    bnb_4bit_compute_dtype=torch.bfloat16   # Compute in bf16
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=quantization_config,
    device_map="auto"
)

# Step 2: Add LoRA adapters
lora_config = LoraConfig(
    r=16,              # Rank (balance capacity vs memory)
    lora_alpha=32,     # Scaling factor
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Which layers to adapt
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# Step 3: Enable gradient checkpointing
model.gradient_checkpointing_enable()

# Step 4: Configure training with small batches + accumulation
training_args = TrainingArguments(
    output_dir="./llama-7b-finetuned",
    
    # Memory settings
    per_device_train_batch_size=4,      # Small batch per GPU
    gradient_accumulation_steps=8,       # Effective batch = 4 × 8 = 32
    
    # Precision
    bf16=True,                           # Use BF16 (if supported)
    
    # Training
    num_train_epochs=3,
    learning_rate=2e-4,
    
    # Checkpointing
    save_strategy="steps",
    save_steps=100,
)

# Step 5: Train!
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    max_seq_length=512,  # Reasonable sequence length
)

trainer.train()

Memory Breakdown:

Base model (4-bit):     ~3.5 GB
LoRA params (bf16):     ~67 MB
Optimizer state:        ~268 MB
Gradients:              ~67 MB
Activations (bs=4):     ~5 GB (with checkpointing)
Framework overhead:     ~500 MB
──────────────────────────────
Total:                  ~9.4 GB

Fits comfortably on a 12 GB GPU with room to spare!

This same approach scales:

• 13B model on 24 GB GPU (RTX 4090)

• 70B model on 48 GB GPU (A6000)

The QLoRA revolution made this all possible.

Quick Reference: Memory Optimization Techniques¶

Here’s your cheat sheet. Bookmark this cell.

The Impact Hierarchy¶

Most impact for least effort:

1. Mixed precision (BF16)

2. LoRA

3. Gradient checkpointing

When you’re desperate: 4. Quantization (4-bit) 5. CPU offloading

Model Size → GPU Requirements¶

With all optimizations (QLoRA + gradient checkpointing + BF16):

These are realistic fine-tuning numbers, not theoretical minimums.

Remember¶

The goal isn’t to use every technique. The goal is to use the minimum number of techniques needed to fit your model in memory while maintaining reasonable training speed.

Start simple. Add complexity only when needed.

What’s Next?¶

You now have the tools to fit large models on limited hardware. But memory is just one piece of the puzzle.

Next up: hyperparameter tuning. Learning rate schedules, warmup steps, weight decay—all the knobs you can turn to make your model actually learn well (not just fit in memory).

Because a model that fits in memory but doesn’t train well is... not very useful.

(But hey, at least it runs!)