GPU acceleration and device management

Overview

GPU acceleration moves tensor computations from the CPU to specialized parallel processors (GPUs) that can perform thousands of operations simultaneously. Device management is the practice of explicitly controlling where tensors and models live (CPU vs GPU) and how data moves between devices.

Figure — GPU acceleration and device management

Core Concepts

The Physics of Why GPUs Win

Derivation: Speedup Factor

The time to compute a matrix multiplication $C = AB$ where $A \in \mathbb{R}^{m \times k}$ , $B \in \mathbb{R}^{k \times n}$ :

$T_{\text{CPU}} = \frac{2mkn}{f_{\text{CPU}} \cdot c_{\text{CPU}}}$

where $f_{\text{CPU}}$ is clock frequency (Hz), $c_{\text{CPU}}$ is cores (~16).

For GPU with $c_{\text{GPU}}$ cores (~10,000):

$T_{\text{GPU}} = \frac{2mkn}{f_{\text{GPU}} \cdot c_{\text{GPU}}} + T_{\text{transfer}}$

Speedup: $S = \frac{T_{\text{CPU}}}{T_{\text{GPU}}} \approx \frac{f_{\text{GPU}} \cdot c_{\text{GPU}}}{f_{\text{CPU}} \cdot c_{\text{CPU}}}$

Pluging in typical values: $f_{\text{GPU}} \approx 1.5$ GHz, $c_{\text{GPU}} = 10000$ , $f_{\text{CPU}} = 3$ GHz, $c_{\text{CPU}} = 16$ :

$S \approx \frac{1.5 \times 10000}{3 \times 16} = \frac{15000}{48} \approx 312\times$

But: This assumes $T_{\text{transfer}} \ll T_{\text{GPU}}$ . If you transfer data every operation, speedup colapses.

PyTorch Device Management: From First Principles

How Tensors Know Their Location

Every PyTorch tensor has a .device property. Under the hood, it's a pointer to memory in either:

CPU DRAM (accessible by CPU cores)
GPU VRAM (accessible by GPU cores, isolated from CPU)

Worked example Example 1: Creating Tensors on Devices

import torch
 
# Default: CPU
x = torch.randn(1000, 1000)
print(x.device)  # cpu
 
# Explicit GPU (if available)
if torch.cuda.is_available():
   device = torch.device('cuda')
   y = torch.randn(1000, 1000, device=device)
   print(y.device)  # cuda:0
   
   # Move existing tensor
   x_gpu = x.to(device)  # Copies data to GPU
   print(x_gpu.device)  # cuda:0

Why this step?

torch.device('cuda') creates a device handle.
device=device allocates memory in GPU VRAM.
.to(device) triggers cudaMemcpy under the hood (PCIe transfer).

The Movement Tax

Worked example Example 2: Quantifying Transfer Cost

import time
x_cpu = torch.randn(1000, 1000)  # 400 MB
# CPU computation
start = time.time()
z_cpu = x_cpu @ x_
print(f"CPU time: {time.time() - start:.4f}s")
 
# GPU computation (with transfer)
start = time.time()
x_gpu = x_cpu.to('cuda')  # Transfer cost
z_gpu = x_gpu @ x_gpu
torch.cuda.synchronize()  # Wait for GPU to finish
print(f"GPU time (with transfer): {time.time() - start:.4f}s")
# GPU computation (no transfer)
start = time.time()
z_gpu2 = x_gpu @ x_gpu  # Data already on GPU
torch.cuda.synchronize()
print(f"GPU time (no transfer): {time.time() - start:.4f}s")

Typical output:

CPU: 2.3s
GPU with transfer: 0.8s (transfer took ~0.7s, compute ~0.1s)
GPU no transfer: 0.1s (23× faster!)

Why this step? torch.cuda.synchronize() is crucial—GPU ops are asynchronous. Without it, timing is wrong.

Training Loop: Device-Aware Pattern

The Standard Recipe

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
# 1. Move model to device (once)
model = MyModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
 
# 2. In training loop, move batches to device
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # Move batch to GPU
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)  # Computed on GPU
        loss = criterion(output, target)  # Loss on GPU
        loss.backward()  # Gradients computed on GPU
        optimizer.step()  # Parameter update on GPU
        
        # Only move to CPU for logging
        if batch_idx % 100 == 0:
            print(f'Loss: {loss.item():.4f}')  # .item() copies scalar to CPU

Why each step?

.to(device) on model moves all parameters/buffers to GPU memory (one-time cost)
Moving batches to GPU: unavoidable (data loader reads from disk → CPU RAM → GPU VRAM)
.item() for logging: scalars are tiny (~4 bytes), transfer is negligible
Everything else stays on GPU: no unnecessary round-trips

Multi-GPU: Data Parallelism Derivation

The Problem

One GPU has limited memory (~24 GB). Large models or batches don't fit.

The Solution: Split the Batch

Formula Data Parallel Forward Pass

Given $N$ GPUs, batch size $B$ , split into $B/N$ per GPU:

Replicate model on all GPUs: $\theta_0 = \theta_1 = \cdots = \theta_{N-1}$
Scatter batch: GPU $i$ gets $X_i, Y_i$ where $|X_i| = B/N$
Forward on each GPU (parallel): $\hat{Y}_i = f_{\theta_i}(X_i)$
Compute loss per GPU: $L_i = \mathcal{L}(\hat{Y}_i, Y_i)$
Backward per GPU (parallel): $\nabla_{\theta_i} L_i$
All-Reduce gradients: $\nabla_\theta L = \frac{1}{N}\sum_{i=0}^{N-1} \nabla_{\theta_i} L_i$
Update (synchronized): $\theta \leftarrow \theta - \alpha \nabla_\theta L$

Communication cost: All-reduce requires $O(N)$ exchanges. For $P$ parameters, bandwidth needed: $\approx 2P \cdot \text{sizeof(float)}$ per iteration.

Worked example Example 3: PyTorch DataParallel

if torch.cuda.device_count() > 1:
   print(f"Using {torch.cuda.device_count()} GPUs")
   model = nn.DataParallel(model)  # Wraps model
model.to(device)
# Training loop unchanged!
# DataParallel automatically:
# - Splits batch across GPUs
# - Gathers outputs
# - Averages gradients

Why this works? nn.DataParallel intercepts forward(), splits input, runs replicas, gathers output. Gradients are averaged in .backward().

Modern alternative: DistributedDataParallel (DP) uses one process per GPU, faster all-reduce via NCCL.

Common Mistakes

Common mistake Mistake 1: Mixing CPU and GPU Tensors

x = torch.randn(100, 100).to('cuda')
y = torch.randn(100, 100)  # CPU
z = x + y  # RuntimeError: Expected all tensors on same device

Why it feels right: "Shouldn't PyTorch auto-transfer?" The reality: Auto-transfer would hide expensive ops. Explicit is better. GPU kernels can't access CPU memory.

Fix: y = y.to('cuda') before operation. Or use .to(x.device) to match.

Common mistake Mistake 2: Synchronous CPU Access in Tight Loop

for i in range(1000):
   x_gpu = torch.randn(1000, 1000, device='cuda')
   result = some_gpu_op(x_gpu)
   print(result[0, 0])  # Transfers ONE element 1000 times!

Why it feels right: "Just printing one number, how bad can it be?"

The reality: Each access forces GPU→CPU sync (cudaDeviceSynchronize). GPU sits idle waiting for PCIe.

Fix: Accumulate on GPU, transfer once: results.append(result) then cpu_results = torch.stack(results).cpu().

Common mistake Mistake 3: Forgetting to Move Loss to CPU for Logging

loss_history = []
for batch in train_loader:
   # ... training code ...
   loss_history.append(loss)  # loss is CUDA tensor!

Why it feels right: "Lists work with any Python object."

The reality: Every appended GPU tensor keeps GPU memory allocated. 1000 iterations = 1000 tensors in VRAM, OM crash.

Fix: loss_history.append(loss.item()) or loss.detach().cpu().

Mixed Precision Training: The 2× Speedup Trick

Why FP16?

Modern GPUs (Tensor Cores) can do FP16 matrix multiplies at 2× the FLOPS of FP32.

The problem: FP16 range is $\pm 65,504$ . Gradients often need FP32 precision (tiny values).

The solution: Automatic Mixed Precision (AMP)

Worked example Example 4: AMP in PyTorch

from torch.cuda.amp import autocast, GradScaler
 
scaler = GradScaler()
 
for data, target in train_loader:
   data, target = data.to(device), target.to(device)
   optimizer.zero_grad()
   # Forward in FP16
   with autocast():
       output = model(data)
       loss = criterion(output, target)
   
   # Backward with scaled loss
   scaler.scale(loss).backward()
   # Unscale, check for infs/NaNs, step
   scaler.step(optimizer)
   scaler.update()

Why this step?

autocast(): Casts eligible ops to FP16 (matmul, conv) but keeps numerically sensitive ops (softmax, loss) in FP32.
scaler.scale(loss): Multiplies by scale factor before .backward().
scaler.step(): Unscales gradients, checks for overflow (if found, skips update and reduces scale factor).

Memory Management: The Hidden Enemy

GPU Memory Lifecycle

Allocate → Use → Free (maybe)

PyTorch uses a caching allocator:

First allocation: cudaMalloc (slow, 1ms)
Subsequent: reuse cached blocks (fast, 1µs)
.empty_cache() returns memory to CUDA but doesn't free to OS

Worked example Example 5: Debugging OM

import torch
 
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
# Force garbage collection
import gc
gc.collect()
torch.cuda.empty_cache()
# Find large tensors
for obj in gc.get_objects():
   if torch.is_tensor(obj) and obj.is_cuda:
       print(type(obj), obj.size(), obj.element_size() * obj.nelement() / 1e6, "MB")

Why this helps: memory_allocated shows live tensors. memory_reserved includes cache. Gap = wasted space.

Gradient Accumulation for Large Batches

When batch size 512 doesn't fit, use gradient accumulation:

accumulation_steps = 4  # Effective batch size = 128× 4 = 512
 
for i, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    
    output = model(data)
    loss = criterion(output, target) / accumulation_steps  # Scale loss
    loss.backward()  # Accumulate gradients
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Derivation:

Normal: $\nabla_\theta L = \frac{1}{B}\sum_{i=1}^B \nabla_\theta \ell_i$
Accumulated: $\nabla_\theta L = \frac{1}{4} \sum_{j=1}^4 \left[\frac{1}{B/4}\sum_{i \in \text{batch}_j} \nabla_\theta \ell_i \right] = \frac{1}{B}\sum_{i=1}^B \nabla_\theta \ell_i$ ✓

Why divide loss? .backward() sums gradients. Without /accumulation_steps, final gradients would be 4× too large.

Recall Explain to a 12-Year-Old

Imagine you have math homework:1000 multiplication problems. CPU (your brain): You solve them one-by-one. You're smart and fast, but it takes time because you do one thing at a time.

GPU (a classroom of1000 kids): Each kid solves ONE problem. They're not geniuses, but together they finish in the time it takes one kid to solve one problem!

Neural networks are like this homework—millions of simple multiplications. GPUs have thousands of "kids" (cores) working in parallel.

The catch: Moving the homework from your desk (CPU memory) to the classroom (GPU memory) takes time. So we move all problems at once (batches), not one-by-one.

Mixed precision: Some problems can be solved with simpler math (FP16 = using fewer decimal places). It's faster but sometimes makes mistakes, so we keep a "master copy" of the answer in high precision (FP32).

Connections

3.3.05-Computational-graphsand-autograd: GPU accelerates both forward and backward passes through the computational graph
3.2.03-Backpropagation: Gradient computation is embarrassingly parallel, perfect for GPUs
3.4.02-Batch-normalization: Requires all-reduce in multi-GPU (sync mean/variance)
4.1.01-CNN-architectures: Convolutions are GPU bottleneck—optimizing kernels is critical
3.3.07-Distributed-training-strategies: Multi-node extends device management across machines

#flashcards/ai-ml

Why do GPUs outperform CPUs for deep learning? :: GPUs have thousands of cores optimized for parallel arithmetic (low clock speed but massive parallelism), while CPUs have ~16 complex cores for sequential tasks. Neural network operations (matrix multiplies) are embarrassingly parallel—each element computed independently—matching GPU architecture perfectly.

What is the GPU memory transfer bottleneck?

PCIe bandwidth (~16 GB/s) is 50× slower than GPU memory bandwidth (~900 GB/s). Frequent CPU↔GPU transfers destroy speedup. Solution: move data once, keep computations on GPU, only transfer results.

What does torch.cuda.synchronize() do and why is it needed?

GPU operations are asynchronous—Python continues while GPU computes. synchronize() blocks until all GPU kernels finish. Needed for accurate timing and when CPU must read GPU results. Without it, timing measures queue time, not compute time.

Explain data parallelism's all-reduce step :: Each GPU computes gradients on its batch partition independently. All-reduce synchronizes gradients: each GPU gets the average gradient across all GPUs. Formula: $\nabla\theta = \frac{1}{N}\sum_{i=1}^N \nabla\theta_i$ . Implemented efficiently via ring all-reduce (NCCL library), bandwidth cost $O(2P/N)$ where $P$ is parameter count, $N$ is GPUs.

Why does mixed precision training use loss scaling?

FP16 range is

\pm 65,504

with min

6 \times 10^{-8}

. Gradients in deep networks can be

10^{-10}

(underflow to zero in FP16). Scaling loss by

2^{16}

before backward pass scales gradients into FP16 range. After backward, unscale gradients before optimizer step. Preserves small gradients without overflow.

What is gradient accumulation and when to use it?

Simulating large batch sizes on limited GPU memory. Compute forward/backward on mini-batches, accumulate gradients (don't step optimizer), then step after N accumulations. Effective batch size = mini-batch × N. Critical: divide loss by N or optimizer sees N× larger gradients. Use when batch size B doesn't fit in memory but B/N does.

Why does appending GPU tensors to a list cause OOM?

Each tensor reference keeps GPU memory allocated. List of1000 tensors = 1000 separate allocations in VRAM. PyTorch's garbage collector doesn't free until Python object deleted. Solution: .item() for scalars (copies to CPU as Python float), or .detach().cpu() for larger tensors, severing GPU reference.

What's the difference between memory_allocated and memory_reserved? :: memory_allocated(): currently live tensor memory. memory_reserved(): memory PyTorch holds from CUDA (includes cache). Gap = cached free blocks. After del tensor, memory_allocated drops but reserved stays high (cache). .empty_cache() returns reserved memory to CUDA but not to OS (CUDA driver still holds it).

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, yahan core baat samajhne ki ye hai ki GPU aur CPU dono alag-alag tarah ke workers hain. CPU ek master chef ki tarah hai — bahut fast, par ek time pe thode se hi kaam kar sakta hai kyunki uske paas sirf 8-16 powerful cores hote hain. GPU thousands of simple cooks ki tarah hai — har ek chhota kaam simultaneously karta hai. Ab neural networks mein jo matrix multiplications hote hain, wo millions of independent multiply-add operations hote hain — matlab "embarrassingly parallel" kaam. Isiliye GPU pe ye kaam 100x se 300x tak fast ho jaata hai, kyunki saare operations ek saath parallel mein chal sakte hain.

Lekin ek catch hai jo bahut zaroori hai — data ko CPU (RAM) se GPU (VRAM) tak move karna padta hai, aur ye transfer SLOW hota hai (PCIe bottleneck ~16 GB/s), jabki GPU ki apni memory bandwidth ~900 GB/s hoti hai. To agar tum har chhoti operation ke liye baar-baar data idhar-udhar bhejte raho, to tumhara saara speedup advantage khatam ho jaata hai. Isiliye "Golden Rule" ye hai: data ko GPU pe hi rakho, transfers minimize karo, aur operations ko batch karke ek saath bhejo. PyTorch mein har tensor ka ek .device property hota hai jo batata hai wo CPU pe hai ya GPU pe, aur .to(device) se tum use move kar sakte ho.

Ye baat kyun matter karti hai? Kyunki jab tum real projects mein deep learning models train karoge, to sahi device management ke bina tumhara training bahut slow rahega ya crash bhi ho sakta hai. Bahut se beginners galti karte hain ki wo model GPU pe daal dete hain par data CPU pe chhod dete hain, ya baar-baar transfer karte hain — aur phir socha hua speed nahi milta. Agar tum ye intuition samajh loge — GPU ki power parallelism mein hai, aur transfer ek "tax" hai jise avoid karna hai — to tum efficient aur fast ML systems bana paoge, jo industry mein ek bahut valuable skill hai.

Overview

Core Concepts

The Physics of Why GPUs Win

Derivation: Speedup Factor

PyTorch Device Management: From First Principles

How Tensors Know Their Location

The Movement Tax

Training Loop: Device-Aware Pattern

The Standard Recipe

Multi-GPU: Data Parallelism Derivation

The Problem

The Solution: Split the Batch

Common Mistakes

Mixed Precision Training: The 2× Speedup Trick

Why FP16?

Memory Management: The Hidden Enemy

GPU Memory Lifecycle

Gradient Accumulation for Large Batches

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks

Connections