Optimize GPU Memory

You know the feeling: your training run hits the halfway mark, the loss curve is dropping beautifully, and then — CUDA out of memory. The process dies, your GPU is idle, and you've lost hours of compute. For anyone training deep learning models, GPU memory is the scarcest resource you'll fight over. It's not just about buying a bigger card; it's about understanding how memory is allocated and reclaimed during training, and then systematically optimizing it. This lesson gives you the mental model and a practical toolkit to squeeze more training into the same GPU — essential for moving from toy models to real-world applied AI.

The problem this lesson solves

Larger models, bigger batches, and longer sequences are the lifeblood of modern deep learning. But every increase in model size, batch size, or input length translates directly into increased GPU memory demand. When your process exceeds the GPU's available VRAM, you get the dreaded CUDA out of memory error, and your training run crashes. This is doubly frustrating because

• You may not be able to afford a larger GPU.
• Shared GPU environments (like cloud instances) impose strict memory limits.
• Simply reducing batch size can hurt convergence and slow down training.

The problem is that developers often treat GPU memory as a black box — they don't know what's consuming it, why, or how to reduce it without sacrificing performance. Without a systematic approach, you'll end up guessing, over-allocating, and leaving performance on the table.

By the end of this lesson, you'll be able to explain the core allocators of GPU memory during training and apply a suite of optimization techniques to fit larger models or batches into the same memory footprint.

Core concept / mental model

Think of GPU memory as a construction site. During training, you have several distinct jobs running simultaneously:

• Model weights: The scaffold — the parameters your model needs to make predictions.
• Gradients: The blueprint changes — the updates computed by backpropagation.
• Optimizer states: The tool shed — additional values like momentum and variance that the optimizer tracks.
• Activations: The temporary workbenches — the intermediate outputs of each layer from the forward pass.
• Temporary buffers: Spare tools for operations like batchnorm statistics and CUDA workspace.

Your goal is to manage this site efficiently: reuse the workbenches, clean up unused space, and move heavy equipment to a cheaper location when needed.

Key definitions:

• Batch size: Number of samples processed together. Larger batches need more memory for activations and gradients.
• Gradient checkpointing: Recomputing activations during backprop instead of storing them all — a classic time-memory tradeoff.
• Mixed precision: Using float16 instead of float32 for certain tensors, cutting memory nearly in half.

How it works step by step

Optimizing GPU memory follows a clear progression. Start with the highest-impact, lowest-effort changes, then move to more complex strategies.

1. Profile your memory usage

You can't optimize what you can't measure. Before making changes, identify exactly what's consuming memory:

import torch

def print_memory_usage():
    print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
    print(f"Reserved:   {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
    print(f"Max (peak): {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

# Call this at different points (before, during, after training)
print_memory_usage()

You can also use PyTorch's memory profiler for a detailed breakdown per tensor and operation.

2. Reduce batch size first

A smaller batch size directly reduces activations and gradient memory. This is the simplest lever — but it may slow convergence. Later you can use gradient accumulation to keep effective batch size.

3. Enable gradient checkpointing

Instead of storing all activations, PyTorch can recompute them during backward pass. This trades compute for memory, often reducing activation memory by 60–80%.

model = YourModel()
model.gradient_checkpointing_enable()  # Hugging Face style
# Or for raw PyTorch:
from torch.utils.checkpoint import checkpoint_sequential

4. Use mixed precision

With automatic mixed precision (AMP), you store activations in float16 while keeping a float32 master copy of weights for stability. This halves memory usage for those tensors and can speed up training on Tensor Cores.

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
for data, target in dataloader:
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

5. Offload and watch your data pipeline

• Move data loading to CPU with num_workers and pin_memory=True — keep the GPU clear of I/O spikes.
• Use with torch.no_grad() for inference-only calculations.
• Delete tensors you no longer need and call torch.cuda.empty_cache() only when necessary (it's slow).

Hands-on walkthrough

Let's apply these techniques to a real training loop. We'll train a simple CNN on CIFAR-10, first with default settings and then optimized, and compare memory usage.

Step 1: Setup

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Simple CNN
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc = nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = Net().cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()

Step 2: Baseline profiling

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True)

torch.cuda.reset_max_memory_allocated()
for i, (data, target) in enumerate(trainloader):
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
    if i == 5:
        break

# After the loop:
print(f"Peak memory used: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Expected output (approximate): Peak memory used: 0.24 GB

Step 3: Optimized training with gradient checkpointing + AMP

from torch.cuda.amp import autocast, GradScaler
from torch.utils.checkpoint import checkpoint_sequential

# Enable gradient checkpointing for convolutional sequence
seq = [model.conv1, model.conv2, model.fc]

def optimized_forward(x):
    x = checkpoint_sequential(seq, 3, x)
    return x

# Override forward method (or create a wrapper)
model.forward = optimized_forward

scaler = GradScaler()

torch.cuda.reset_max_memory_allocated()
for i, (data, target) in enumerate(trainloader):
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    if i == 5:
        break

print(f"Peak memory used (optimized): {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Expected output (approximate): Peak memory used (optimized): 0.14 GB

Notice the improvement — roughly 40% reduction in peak memory in this simple example. In larger language models, the savings can be dramatic.

Step 4: Gradient accumulation to recover effective batch size

When you reduce batch size, you might lose convergence quality. Gradient accumulation fixes that:

accumulation_steps = 4
effective_batch_size = batch_size * accumulation_steps

for i, (data, target) in enumerate(trainloader):
    data, target = data.cuda(), target.cuda()
    output = model(data)
    loss = loss_fn(output, target) / accumulation_steps  # scale down
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

This way, your optimizer updates as if you had used a batch size of batch_size * accumulation_steps, but the memory usage stays low.

Compare options / when to choose what

Different techniques have different tradeoffs. Here's a comparison table:

When to choose what: - Start with mixed precision — it's almost always a win. - Add gradient checkpointing when activations dominate memory (common in transformers and CNNs with large feature maps). - Use gradient accumulation to recover batch size without memory cost. - Consider offloading only when you need to train a model that vastly exceeds GPU memory and you can tolerate speed loss.

Troubleshooting & edge cases

• CUDA out of memory despite optimizations: Check for memory leaks — for instance, accumulating loss values or storing gradients. Use torch.cuda.memory_summary() to inspect.
• Mixed precision causes NaN/Inf: This can happen with very small gradients. Make sure you scale the loss via GradScaler and keep a float32 master copy of the weights.
• Gradient checkpointing slows training too much: If your model is small, the recomputation overhead may dominate. Use checkpointing only on large activation blocks.
• torch.cuda.empty_cache() doesn't reduce memory: It releases cached memory but doesn't free tensors still referenced. Delete variable references first.
• Multi-GPU training still OOM: Each GPU holds its own copy of the model; consider model parallelism or ZeRO (from DeepSpeed) to shard the model.

What you learned & what's next

You've learned how to explain the core ideas behind GPU memory optimization and apply a practical toolkit to reduce memory usage during training. You can now profile memory, use mixed precision, gradient checkpointing, and gradient accumulation to train larger models within your hardware limits.

In the next lesson of the Applied AI engineering track, you'll apply these techniques to fine-tune Large Language Models efficiently — fitting billion-parameter models on a single consumer GPU using parameter-efficient fine-tuning methods like LoRA.