Kahibaro
Discord Login Register

14.3 GPU Utilization and Memory Management

What “Good GPU Utilization” Means

A GPU speeds up training when it stays busy doing large, parallelizable math, while the CPU feeds it data fast enough and the program avoids unnecessary synchronization. In practice, low utilization usually means the GPU is waiting, often on the input pipeline, Python overhead, small batch sizes, or frequent transfers between CPU and GPU. High utilization does not automatically mean fast training, because you can also be bottlenecked by memory bandwidth or by inefficient kernels, but it is a strong signal that your workload is large enough and well fed.

The key mental model is that the GPU is a separate device with its own memory. Work submitted to it is typically asynchronous, meaning PyTorch queues operations and the CPU continues. Any time you force the CPU to wait for a GPU result, such as printing a CUDA tensor value or calling certain timing functions incorrectly, you can accidentally stall the pipeline and reduce utilization.

Moving Work to the GPU Correctly

To run on the GPU, both the model parameters and the input tensors must be on the same CUDA device. A common pattern is to create a single device object and move everything once at setup time, then move each batch right before the forward pass.

You should treat device transfers as relatively expensive operations. Move tensors in large chunks rather than many tiny transfers, and avoid repeatedly calling .to(device) on the same tensors inside tight loops unless necessary. When you do move a batch from CPU to GPU, you can overlap transfer with computation by using pinned memory in the DataLoader and non blocking copies, which is covered later in this chapter.

Rule: A CUDA model and CPU inputs will error or silently trigger expensive transfers in some workflows. Ensure model and batch tensors are on the same device before the forward pass.

Diagnosing Underutilization Quickly

Before changing code, distinguish between “GPU waiting for data” and “GPU busy but slow.” If the GPU utilization is near zero or spiky, suspect the input pipeline, too small batches, or synchronization points. If utilization is high but throughput is low, suspect memory pressure, too many small ops, or kernel inefficiencies.

In PyTorch, frequent causes of accidental synchronization include calling .item() on a CUDA tensor inside the training loop, converting CUDA tensors to NumPy, and printing CUDA tensors often. These operations require the GPU to finish pending work so the CPU can read the result.

Timing on GPU must also be done carefully because of asynchrony. If you measure time with Python’s wall clock around CUDA code without synchronizing, you may underestimate time and make misleading conclusions. If you need accurate step timing, synchronize around the timed region.

Rule: .item(), .cpu(), .numpy(), and many printing or logging patterns can synchronize CUDA and stall the pipeline if used inside the hot path.

Memory Basics: Parameters, Activations, Gradients, Optimizer State

GPU memory is consumed by several categories. Model parameters live on the GPU. During the forward pass, intermediate activations are stored for backpropagation unless you are in inference mode or using checkpointing. During the backward pass, gradients for parameters are created. Many optimizers also keep extra state, for example Adam keeps running estimates, effectively adding roughly two extra buffers per parameter, which increases memory significantly.

The practical implication is that “model size” is not the same as “training memory.” Training memory is usually dominated by activations, especially with large batch sizes or large feature maps in CNNs and long sequences in NLP. That is why you can often fit a model for inference but run out of memory during training.

Rule: Training uses more memory than inference because activations and optimizer state must be stored. Do not estimate memory needs from parameters alone.

Choosing Batch Size With Memory in Mind

Batch size is one of the fastest levers for memory and throughput. Larger batches increase activation memory roughly linearly and can improve GPU utilization by giving kernels more work. If you hit out of memory errors, reducing batch size is the simplest fix. If a small batch causes low utilization, you can simulate a larger effective batch using gradient accumulation, which has its own dedicated chapter, but the memory takeaway here is that accumulation increases compute time while keeping per step activation memory closer to the smaller micro batch.

If your batch size is very small due to memory limits, consider whether your model has unusually large activations. In vision, large input resolution and early feature maps are common culprits. In transformers, sequence length drives activation memory strongly.

Avoiding Out of Memory Errors and Recovering Cleanly

A CUDA out of memory error means the allocator could not find a contiguous block large enough for the requested tensor. PyTorch uses a caching allocator, which holds freed blocks for reuse instead of returning them to the driver immediately. This is good for speed but can confuse newcomers because memory may appear “used” even after tensors are freed.

If you hit out of memory, first remove large references so tensors can be freed, then consider clearing the cache. Clearing the cache does not free memory held by live tensors, it only releases cached blocks back to the driver, which may help other processes or subsequent allocations.

If an out of memory happens mid step, the safest recovery in a notebook is often to restart the kernel, because partial graphs and references can keep memory alive. In scripts, you generally fix the underlying cause rather than trying to recover at runtime.

Rule: torch.cuda.empty_cache() does not reduce memory used by tensors you still reference. It only clears cached free blocks.

Understanding the CUDA Caching Allocator and Fragmentation

PyTorch’s caching allocator improves performance by reusing memory blocks, but it can lead to fragmentation, where there is enough total free memory but not in the right sized contiguous blocks. Fragmentation becomes more likely when tensor sizes vary a lot across iterations, such as variable image sizes, variable sequence lengths without bucketing, or changing batch sizes.

To reduce fragmentation, keep shapes consistent across steps when possible. For variable length sequences, padding and batching strategies can make memory usage more regular. For images, standardizing resize and crop sizes stabilizes allocations. If your workload must vary, you may still run fine, but if you see intermittent out of memory errors at the same nominal batch size, fragmentation is a suspect.

Monitoring Memory Usage in PyTorch

You can query memory stats from PyTorch to see current allocated memory and cached memory. Allocated is memory currently used by tensors, reserved is memory held by the allocator including cached blocks. When reserved is much larger than allocated, it means the cache is large, which is normal after peaks.

When debugging, watch for patterns where allocated memory grows every iteration. That usually indicates a memory leak, typically caused by storing tensors that require gradients in a Python list, keeping references to computation graphs, or logging tensors without detaching.

A common safe logging pattern is to detach tensors before storing them and to store small CPU scalars rather than full GPU tensors.

Rule: If GPU allocated memory increases every iteration, you are likely retaining computation graphs by storing tensors without .detach() or by accumulating outputs for logging.

Data Transfer and Pinned Memory

Even with a fast GPU, training can be slow if each batch transfer from CPU to GPU is a bottleneck. Two common steps help. First, use a DataLoader with pin_memory=True, which allocates page locked host memory that enables faster DMA transfers to the GPU. Second, when moving batches to the GPU, use non blocking transfers by calling .to(device, non_blocking=True) on tensors that originate from pinned memory.

This works best when you also use multiple DataLoader workers, so the next batch is prepared while the GPU is computing on the current batch. The goal is overlap, not just faster copies.

You should also avoid bouncing tensors back and forth between CPU and GPU. Keep tensors on GPU for as long as they are needed there, and only move results to CPU when necessary, ideally after detaching and only for small summaries.

Reducing Activation Memory: Inference Mode and No Grad

During evaluation or inference, you usually do not need gradients. Disabling gradient tracking reduces memory use and speeds up execution. In PyTorch, use torch.no_grad() for evaluation loops, and consider model.eval() to switch layer behavior like dropout and batch norm.

There is also torch.inference_mode(), which is even more restrictive and can be faster and lighter for pure inference workloads. The main idea for beginners is simple: do not build computation graphs when you do not need them.

Rule: Wrap validation and inference in torch.no_grad() or torch.inference_mode() to reduce memory and improve speed.

Reducing Memory With Mixed Precision and Choosing Numeric Types

Using lower precision can reduce memory footprint and increase throughput on many GPUs. Mixed precision typically stores many activations in float16 or bfloat16 while keeping certain computations in float32 for stability. This can cut activation memory substantially, which often enables larger batches.

Mixed precision has its own chapter, so the key point here is to recognize it as a memory management tool, not just a speed trick. If you are memory bound, mixed precision is often one of the highest impact changes.

Practical Patterns That Accidentally Waste Memory

A frequent beginner pitfall is storing per batch losses or predictions as GPU tensors in a Python list for later plotting. Each stored tensor may keep a computation graph alive, preventing intermediate buffers from being freed. Another pitfall is computing metrics on GPU tensors and repeatedly converting them to Python numbers with .item() inside the loop, which can synchronize and slow training.

A better habit is to detach, move small summaries to CPU, and keep only what you need. If you need to accumulate predictions for a full epoch, detach them first so they do not retain the graph, and consider moving them to CPU if they are large.

Putting It Together: A Mental Checklist

If training is slow, first check that the model and data are truly on the GPU and that you are not synchronizing every step through logging. Next, check whether the DataLoader can keep up, then consider increasing batch size until you approach memory limits. If you run out of memory, reduce batch size, switch to a lighter optimizer, enable mixed precision, or reduce input size or sequence length. If memory usage keeps climbing, look for retained graphs and tensors stored without detaching. Finally, stabilize shapes across iterations to reduce fragmentation and unpredictable out of memory errors.

Key statement: The fastest GPU code is often the code that minimizes host device transfers, avoids synchronization in the hot loop, and keeps shapes and memory usage predictable across iterations.

Views: 60

Comments

Please login to add a comment.

Don't have an account? Register now!