GPU Optimization Techniques for Deep Learning

Introduction

GPU optimization is critical for training large-scale deep learning models efficiently. Without careful tuning, you can easily leave 60–80% of your GPU’s theoretical throughput on the table — which translates directly to longer training runs, higher cloud bills, and existential dread.

This article covers three key areas: memory management, kernel fusion, and mixed-precision training. Each section includes working code you can drop into your own projects today.

Memory Management Strategies

GPU memory is a precious, finite resource. Understanding how to manage it is the first step toward efficient training.

Gradient Checkpointing

Gradient checkpointing trades compute for memory by recomputing intermediate activations during the backward pass instead of storing them all. This can reduce activation memory by a factor of 10× or more for deep networks, at the cost of roughly 30% extra compute.

import torch
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    """Wraps a module with gradient checkpointing."""

    def __init__(self, layers: torch.nn.Sequential):
        super().__init__()
        self.layers = layers

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # use_reentrant=False is recommended for PyTorch >= 2.0
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

When applying checkpointing to a large transformer, wrap only the expensive residual blocks — embedding layers and the final projection typically aren’t worth checkpointing.

Memory Pool Configuration

Pre-allocating memory pools reduces fragmentation and allocation overhead:

• PyTorch: torch.cuda.memory.set_per_process_memory_fraction(0.8)
• JAX: XLA_PYTHON_CLIENT_MEM_FRACTION=0.8
• TensorFlow: tf.config.experimental.set_memory_growth(gpu, True)

Profiling Memory Usage

Before optimizing, measure. PyTorch ships with a built-in memory profiler:

import torch

def profile_memory(model, sample_input):
    torch.cuda.reset_peak_memory_stats()
    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CUDA],
        profile_memory=True,
        record_shapes=True,
    ) as prof:
        output = model(sample_input)
        output.sum().backward()

    print(prof.key_averages().table(
        sort_by="self_cuda_memory_usage", row_limit=15
    ))
    peak_mb = torch.cuda.max_memory_allocated() / 1e6
    print(f"Peak allocated: {peak_mb:.1f} MB")

profile_memory(model, torch.randn(8, 3, 224, 224, device='cuda'))

Kernel Fusion

Kernel fusion combines multiple GPU operations into a single kernel launch, reducing memory bandwidth overhead. Each separate GPU kernel requires a round-trip through global memory; fusing them keeps data in on-chip cache.

Why Fusion Matters

Consider a simple operation chain: LayerNorm → Dropout → ReLU. Without fusion:

1. LayerNorm reads input → writes output to DRAM
2. Dropout reads LayerNorm output from DRAM → writes result to DRAM
3. ReLU reads Dropout output from DRAM → writes final result to DRAM

That’s 6 DRAM round-trips for three cheap operations. With fusion, it’s just 2 (one read, one write).

Benefits of Fusion

1. Reduced kernel launch overhead
2. Better memory locality — intermediate values stay in registers / L1 cache
3. Lower GPU memory bandwidth consumption
4. Improved arithmetic intensity (FLOPs per byte)

Using torch.compile for Auto-Fusion

PyTorch 2.x’s compiler can fuse many patterns automatically:

import torch

# Eager mode — separate kernel launches
def naive_fused_op(x, weight, bias):
    return torch.nn.functional.layer_norm(x, x.shape[-1:]) @ weight + bias

# torch.compile triggers TorchInductor, which fuses + generates optimised CUDA
optimized_op = torch.compile(naive_fused_op, mode="max-autotune")

x = torch.randn(512, 768, device="cuda")
weight = torch.randn(768, 768, device="cuda")
bias = torch.zeros(768, device="cuda")

# Warmup (first call triggers compilation)
for _ in range(3):
    _ = optimized_op(x, weight, bias)

# Benchmark
import time
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(1000):
    _ = optimized_op(x, weight, bias)
torch.cuda.synchronize()
elapsed_ms = (time.perf_counter() - t0) * 1000
print(f"Avg latency: {elapsed_ms / 1000:.3f} ms")

Custom CUDA Kernels with Triton

For maximum control, OpenAI Triton lets you write fused GPU kernels in Python:

import triton
import triton.language as tl

@triton.jit
def fused_bias_relu_kernel(
    x_ptr, bias_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    bias = tl.load(bias_ptr + offsets, mask=mask)
    output = tl.maximum(x + bias, 0.0)  # ReLU(x + bias), fused
    tl.store(output_ptr + offsets, output, mask=mask)

This single kernel adds a bias vector and applies ReLU in one pass over memory — no intermediate tensor allocation required.

Mixed-Precision Training

Mixed-precision uses FP16 (or BF16) for forward/backward passes while maintaining FP32 master weights. The result: roughly 2× the throughput, half the memory, with negligible accuracy loss when done correctly.

AMP with PyTorch

The torch.amp API handles most of the complexity:

import torch
from torch.amp import GradScaler, autocast

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # scales loss to prevent FP16 underflow

for batch in dataloader:
    inputs, targets = batch
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # Scale loss, backward, unscale, clip, step
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

BF16 vs FP16

Format	Exponent bits	Mantissa bits	Dynamic range	Notes
FP32	8	23	±3.4×10³⁸	Baseline
FP16	5	10	±65,504	Needs loss scaling
BF16	8	7	±3.4×10³⁸	No scaling needed

BF16 is preferred on Ampere (A100) and later GPUs — same dynamic range as FP32, so no loss scaling required:

with autocast(device_type='cuda', dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

# BF16 doesn't need GradScaler — just backward directly
loss.backward()
optimizer.step()

Putting It All Together

A production training loop that combines all three techniques:

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

def train_epoch(model, dataloader, optimizer, criterion):
    model.train()
    scaler = GradScaler()
    total_loss = 0.0

    for step, (inputs, targets) in enumerate(dataloader):
        inputs, targets = inputs.cuda(non_blocking=True), targets.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)  # slightly faster than zero_grad()

        with autocast(device_type='cuda', dtype=torch.bfloat16):
            # Gradient checkpointing on the expensive encoder layers
            features = checkpoint_sequential(model.encoder, segments=4, input=inputs)
            outputs = model.head(features)
            loss = criterion(outputs, targets)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

        total_loss += loss.item()

    return total_loss / len(dataloader)

Key details worth calling out:

• non_blocking=True on .cuda() overlaps host→device transfer with compute
• zero_grad(set_to_none=True) avoids zeroing gradients in-place (saves a CUDA kernel)
• checkpoint_sequential splits the encoder into 4 segments for checkpointing

Summary

Effective GPU optimization is not a single trick — it’s a discipline. The three pillars covered here are:

1. Memory management — profile first, then apply gradient checkpointing and memory pooling where it helps
2. Kernel fusion — let torch.compile handle common patterns; reach for Triton for custom hot paths
3. Mixed-precision — BF16 on modern hardware is a near-free 2× win; FP16 with GradScaler works everywhere else

Start with profiling (you cannot optimize what you cannot measure), apply mixed-precision first (highest ROI), then address memory, and finally tune kernels once the easy wins are exhausted.