VRAM — Module 6.3 | Codextroop

How This Lesson Fits the Module

CUDA cores need bits to crunch; those bits live in VRAM (Video RAM)—the GPU’s dedicated high-bandwidth memory. Out-of-memory (OOM) errors are the most common hard stop in deep learning. Before tuning batch size or buying a bigger card, you must know what consumes VRAM and how to measure it.

VRAM limits also drive the precision choices in later lectures (FP16, BF16, INT8)—smaller dtypes mean more fits in the same pool.

Learning Objectives

By the end of this lesson, students should be able to:

Define VRAM and contrast it with system RAM.
Break down VRAM usage into weights, activations, gradients, and optimizer states.
Estimate how batch size and model depth affect memory during training.
Diagnose OOM errors and apply mitigations (gradient checkpointing, smaller batch, mixed precision).
Monitor VRAM with nvidia-smi and PyTorch memory APIs.

VRAM vs System RAM

VRAM sits on the GPU board (GDDR6/HBM) with bandwidth often exceeding 1 TB/s on datacenter cards. System RAM is larger but slower to reach from the GPU across PCIe. During training, tensors that participate in GPU compute must live in VRAM. Keeping weights on CPU while computing on GPU forces constant transfers and kills performance.

Memory	Typical Size	Bandwidth	Deep Learning Role
System RAM	32–512 GB	~50–100 GB/s	Datasets, DataLoader buffers, checkpoints on disk path
VRAM	8–80 GB (consumer to H100)	500 GB/s–3+ TB/s	Model, activations, gradients, optimizer tensors on GPU

What Consumes VRAM During Training?

For a single training step, peak VRAM is roughly the sum of:

Model weights — one copy per parameter tensor (FP32 = 4 bytes/param).
Gradients — same shape as weights (another 4 bytes/param in FP32).
Optimizer states — Adam stores momentum and variance (often 8+ extra bytes/param).
Activations — saved for backprop; scale with batch size and network depth.
Workspace — temporary buffers for cuDNN convolutions and attention.

Component	Scales With	Training vs Inference
Weights	Parameter count	Both (inference often needs weights only)
Activations	Batch size × layers × feature maps	Training (stored); inference often smaller with `torch.inference_mode()`
Gradients + optimizer	Parameter count	Training only

Rule of Thumb Training memory ≈ 3–4× model weight size (FP32 + Adam) plus activation memory. A 7B-parameter model in FP32 weights alone is ~28 GB before activations.

Measuring VRAM in PyTorch

import torch
import torch.nn as nn

torch.cuda.reset_peak_memory_stats()
device = "cuda"

model = nn.Linear(4096, 4096).to(device)
x = torch.randn(256, 4096, device=device)
y = model(x)
loss = y.sum()
loss.backward()

alloc = torch.cuda.memory_allocated() / 1e6
peak = torch.cuda.max_memory_allocated() / 1e6
print(f"Current: {alloc:.1f} MB | Peak: {peak:.1f} MB")

Pair this with nvidia-smi for process-level view. Note: PyTorch caches freed blocks in a memory pool—memory_allocated() is what matters for OOM, not always nvidia-smi’s “used” after empty_cache().

OOM Mitigations

Quick Fixes

Reduce batch_size
Use gradient accumulation to simulate larger batches
Clear unused variables; avoid storing full activation graphs
torch.cuda.empty_cache() between runs (debug only)

Structural Fixes

Mixed precision (FP16/BF16) — halves many tensors
Gradient checkpointing — recompute activations, save VRAM
Smaller model or LoRA fine-tuning
Multi-GPU sharding (Module 6.2 distributed training)

# Gradient accumulation: effective batch 256 with micro-batch 32
accum_steps = 8
optimizer.zero_grad()
for i, (x, y) in enumerate(loader):
    loss = model(x.to("cuda"), y.to("cuda")) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Critical Mistake — Ignoring Activation Memory

Students often count only parameter size. Wide layers with large batch size can make activations dominate VRAM—especially in transformers and high-resolution CNNs (preview in Volume 07 CNNs).

Misconception — Freeing Python Variables Instantly Frees VRAM

PyTorch’s caching allocator may retain blocks for reuse. Use del tensor and torch.cuda.empty_cache() when debugging, but fix the root cause (batch size, precision) for production.

Knowledge Check

Short Answer: What is VRAM? Answer: High-bandwidth memory on the GPU that holds tensors during CUDA computation.
Short Answer: Name four VRAM consumers during training. Answer: Weights, gradients, optimizer states, activations (also workspace).
True/False: Inference always uses more VRAM than training. Answer: False—training needs gradients and optimizer states too.
Multiple Choice: First knob to turn on OOM: (a) learning rate, (b) batch size, (c) random seed. Answer: (b).
Short Answer: How many bytes per FP32 parameter? Answer: 4 bytes.
Short Answer: What does gradient accumulation achieve? Answer: Larger effective batch size without holding all samples’ activations at once.
Short Answer: What PyTorch call reports peak VRAM in the current run? Answer: torch.cuda.max_memory_allocated().
True/False: Adam uses extra VRAM beyond weights and gradients. Answer: True—momentum and variance buffers.
Multiple Choice: Halving precision from FP32 to FP16 on weights roughly: (a) doubles VRAM, (b) halves weight VRAM, (c) no change. Answer: (b).
Short Answer: Why is PCIe bandwidth relevant to VRAM? Answer: Data must be copied from system RAM into VRAM before GPU compute can use it.

Key Takeaways

VRAM is the scarce, fast memory pool on the GPU; OOM means something in that budget exceeded capacity.
Training memory = weights + gradients + optimizer + activations; activations scale with batch size.
Measure with PyTorch memory APIs and nvidia-smi; don’t guess.
Reduce batch size, use mixed precision, checkpointing, or sharding before buying hardware.
Next: Tensor Cores — hardware that multiplies matrices faster while using less VRAM.

Trainer’s Guide

Hands-on idea: Train a small CNN while logging max_memory_allocated() at batch sizes 8, 16, 32, 64 until OOM. Plot batch size vs peak MB.

Discussion prompt: You have 24 GB VRAM and a 13B model. What combination of precision and sharding strategies would you evaluate?

What’s Next Smaller dtypes save VRAM and unlock tensor cores designed for fast low-precision matrix math.