How to Train PyTorch Models Using Far Less GPU Memory

When training large deep‑learning models such as LLMs or Vision Transformers, GPU memory quickly becomes the primary bottleneck. The following techniques can be combined to reduce memory consumption by as much as 20× without sacrificing model performance.

1. Automatic Mixed Precision (AMP)

AMP mixes FP16 and FP32 tensors, halving activation memory while keeping critical operations in FP32 for numerical stability. PyTorch provides native support via torch.cuda.amp.autocast() and torch.cuda.amp.GradScaler:

import torch
from torch.cuda.amp import autocast, GradScaler
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
for data, target in data_loader:
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

2. Low‑Precision BF16

BF16 offers a larger dynamic range than FP16, reducing overflow risk while still cutting memory roughly in half. Modern GPUs (Ampere and later) support BF16. Check support with:

import torch
print(torch.cuda.is_bf16_supported())  # should print True

3. Gradient Checkpointing

Checkpointing stores only a subset of forward activations and recomputes the rest during back‑propagation, saving 40‑50% of activation memory at the cost of extra compute.

import torch
from torch.utils.checkpoint import checkpoint

def checkpointed_segment(input_tensor):
    # custom forward segment
    return model_segment(input_tensor)

output = checkpoint(checkpointed_segment, input_tensor)

4. Gradient Accumulation

Instead of reducing batch size directly, accumulate gradients over several small batches to emulate a larger effective batch. This preserves model accuracy but increases total training time.

5. Tensor Sharding & Fully Sharded Data Parallel (FSDP)

FSDP distributes model parameters, gradients, and optimizer states across GPUs, loading only the needed shards into memory. Combined with other tricks, it can lower memory demand up to ten‑fold.

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = MyLargeModel().cuda()
fsdp_model = FSDP(model)

6. Efficient Data Loading

Enable pinned memory and multiple workers to speed host‑to‑device transfers and avoid data‑loader bottlenecks.

from torch.utils.data import DataLoader
train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

7. In‑Place Operations

Modify tensors in place to avoid creating temporary copies, reducing fragmentation and overall memory usage.

import torch
x = torch.randn(100, 100, device='cuda')
y = torch.randn(100, 100, device='cuda')
# In‑place addition
x.add_(y)  # x is modified directly

8. Lightweight Optimizer

Adam stores two extra state tensors per parameter, roughly doubling memory. Switching to stateless SGD with a cosine‑annealing scheduler cuts optimizer memory by about two‑thirds.

# instead of Adam
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
# use SGD with cosine decay
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)

9. Memory Profiling & Cache Management

Use PyTorch utilities to inspect and free GPU memory.

import torch
print(torch.cuda.memory_summary(device=None, abbreviated=False))
torch.cuda.empty_cache()

10. TorchScript JIT Compilation

Convert Python models to TorchScript to reduce kernel launch overhead and improve memory efficiency.

import torch
scripted_model = torch.jit.script(model)
output = scripted_model(input_tensor)

11. Custom Kernel Fusion & torch.compile()

Fusing multiple operations into a single kernel and using torch.compile() can further shrink memory footprints and boost throughput, especially for Transformer‑style architectures.

Summary

By layering these strategies—AMP, BF16, checkpointing, accumulation, sharding, optimized data loading, in‑place ops, a lean optimizer, profiling tools, JIT compilation, and kernel fusion—researchers and engineers can train large models on a single GPU or modest workstation while keeping memory usage manageable and performance high.