Articles

Maximize GPU Utilization During ML Training (2026)

April 18, 2026 7 min read

The primary driver of GPU utilization during machine learning training isn’t the model’s complexity, but rather the efficiency of data loading and preprocessing.

Let’s watch a typical training loop, stripped down to the essentials, and see where the GPU spends its time.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import time

# Dummy Dataset and Model
class DummyDataset(Dataset):
    def __init__(self, num_samples=10000, feature_size=1024):
        self.num_samples = num_samples
        self.feature_size = feature_size
        # Simulate some complex data generation or loading
        self.data = torch.randn(num_samples, feature_size)
        self.labels = torch.randint(0, 2, (num_samples,))

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # Simulate complex preprocessing
        data = self.data[idx]
        label = self.labels[idx]
        # Imagine this is reading from disk, augmenting, etc.
        time.sleep(0.0001) # Simulate I/O or CPU-bound work
        return data, label

class SimpleModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(input_size, num_classes)

    def forward(self, x):
        return self.fc(x)

# Configuration
dataset_size = 50000
feature_dim = 2048
batch_size = 128
learning_rate = 0.001
num_epochs = 1

# Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

dataset = DummyDataset(num_samples=dataset_size, feature_size=feature_dim)
# Crucially, num_workers > 0 is key for CPU-GPU overlap
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)

model = SimpleModel(input_size=feature_dim, num_classes=2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Training Loop
print("Starting training...")
start_time = time.time()
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(dataloader):
        # Move data to GPU - this is where CPU work finishes and GPU work begins
        inputs, labels = inputs.to(device), labels.to(device)

        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # Backward pass and optimize
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        # Optional: Print progress or log metrics
        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(dataloader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0

    print(f'Epoch [{epoch+1}/{num_epochs}] finished.')

end_time = time.time()
print(f"Training finished in {end_time - start_time:.2f} seconds.")

The core problem is that the GPU is incredibly fast at matrix multiplications and other tensor operations, but it can’t do anything without data. If your data loading and preprocessing pipeline (running on the CPU) can’t feed batches to the GPU fast enough, the GPU will sit idle, waiting. This idle time is what kills utilization.

To maximize GPU utilization, you need to ensure the data pipeline is a well-oiled machine that can keep the GPU fed. This involves several key areas:

1. Parallelize Data Loading with num_workers: The DataLoader in PyTorch is your primary tool. The num_workers argument tells it how many separate CPU processes to use for loading and preprocessing data.

  • Diagnosis: Monitor your CPU and GPU utilization. If CPU usage is low and GPU usage is fluctuating or low, num_workers is likely too low or zero.
  • Fix: Set num_workers to a value that keeps your CPU cores busy but not saturated. A common starting point is num_workers=4 or num_workers=8. For very demanding preprocessing, you might go higher, even up to the number of physical CPU cores.
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8)
  • Why it works: Each worker process loads and preprocesses a batch of data independently. As soon as one worker finishes a batch, it can start on the next, allowing batches to be ready for the GPU by the time the current one finishes computation. This creates an assembly line where data preprocessing happens concurrently with GPU computation.

2. Optimize Data Preprocessing on the CPU: The work done within Dataset.__getitem__ is crucial. If this is slow, num_workers can only do so much.

  • Diagnosis: Profile your __getitem__ method. Use time.time() or a profiler to identify bottlenecks. Are you doing complex image augmentations, reading from slow storage, or performing computationally expensive transformations?
  • Fix:
    • Pre-process offline: If possible, perform expensive transformations (like resizing, complex augmentations) once and save the processed data. Load the pre-processed data during training.
    • Use efficient libraries: Replace slow Python loops with optimized libraries like NumPy, OpenCV, or Pillow-SIMD for image operations.
    • Batch preprocessing: If your __getitem__ is called for single items, consider if some operations can be applied to a small batch of items within __getitem__ or by modifying the DataLoader to yield larger chunks for processing.
    • Avoid Python time.sleep(): The time.sleep(0.0001) in the example is a placeholder for I/O bound or CPU bound work. In reality, this could be disk reads, network fetches, or complex calculations. Ensure these are optimized.
  • Why it works: Faster preprocessing means each worker can generate batches more quickly, reducing the latency between GPU computation and the availability of new data.

3. Utilize pin_memory=True: When pin_memory=True is set in the DataLoader, it tells PyTorch to allocate data tensors in pinned (page-locked) memory.

  • Diagnosis: This is a general optimization. While not directly diagnosable with simple nvidia-smi, it complements num_workers. If you have high num_workers but still see GPU stalls, ensuring pinned memory is used can help.
  • Fix: Set pin_memory=True in your DataLoader.
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
  • Why it works: Tensors in pinned memory can be transferred to the GPU’s memory much faster and more efficiently via DMA (Direct Memory Access) compared to pageable memory. This reduces the overhead of data transfer between the CPU and GPU.

4. Choose Appropriate Batch Size: While not directly about data loading speed, batch size impacts GPU utilization by affecting how much computation is done per data transfer.

  • Diagnosis: If your GPU memory is not fully utilized, you might be able to increase the batch size. If you’re hitting OOM errors, you need to decrease it.
  • Fix: Experiment with larger batch sizes. If you increase batch size, you might need to adjust learning rate (e.g., linear scaling rule: multiply LR by K if batch size is multiplied by K).
    batch_size = 256 # Example increase
dataloader = DataLoader(dataset, batch_size=batch_size, ...)
  • Why it works: Larger batches mean more work for the GPU per data transfer. This amortizes the overhead of kernel launches and data transfers over more operations, potentially leading to higher sustained utilization if the GPU memory can accommodate it.

5. Mix Precision (FP16/BF16): Using half-precision floating-point numbers can significantly speed up computation and reduce memory bandwidth requirements.

  • Diagnosis: Use nvidia-smi to observe GPU memory utilization and compute utilization. If memory bandwidth is a bottleneck, mixed precision can help.
  • Fix: Use torch.cuda.amp.GradScaler and torch.cuda.amp.autocast.
    from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

# Inside the training loop:
with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
  • Why it works: Operations on FP16 or BF16 are typically twice as fast as FP32 on modern GPUs and use half the memory bandwidth. This allows the GPU to process more data per unit of time, and the reduced memory footprint can sometimes allow for larger batch sizes.

6. Overlap Data Transfer and Computation: The goal is to have data transfer to the GPU happening while the GPU is busy with the previous batch’s computation.

  • Diagnosis: This is the ideal state. You’ll see consistently high GPU utilization (e.g., >90%) with nvidia-smi. If you see the GPU usage drop significantly between backward pass and forward pass, it might indicate a bottleneck in data loading or transfer.
  • Fix: This is achieved by combining num_workers > 0, pin_memory=True, and efficient preprocessing. The PyTorch DataLoader is designed to do this automatically when configured correctly. The key is that the inputs, labels = inputs.to(device), labels.to(device) call should happen as early as possible in your loop, ideally after the previous batch’s optimizer.step() and before the current batch’s model(inputs) call.
  • Why it works: By moving data to the GPU in the background (via num_workers and pinned memory) while the GPU is crunching numbers on the previous batch, you minimize the time the GPU spends waiting.

A common pitfall is focusing solely on the model’s forward/backward pass. The entire pipeline, from data on disk to gradients computed, must be optimized. If your GPU utilization hovers around 30-50%, it’s almost certainly a data loading bottleneck.

Once you’ve maximized GPU utilization by ensuring the data pipeline is efficient, you’ll likely encounter the next bottleneck: GPU memory capacity limitations.

Want structured learning?

Take the full Gpu course →