Learn AI Series (#125) - GPU Programming Basics

avatar

Learn AI Series (#125) - GPU Programming Basics

variant-b-09-blue.png

What will I learn

  • You will learn why GPUs are the beating heart of modern AI -- and why "a lot of small cores" beats "a few fast cores" for the kind of math neural networks actually do;
  • the CUDA programming model -- kernels, grids, blocks, threads and warps -- and how those words map onto real silicon (without you ever having to write a line of C);
  • the GPU memory hierarchy (registers, shared, L2, global) and the single most useful realisation in this whole episode: most of your operations are memory-bound, not compute-bound;
  • what cuBLAS and cuDNN are, and why every fast thing PyTorch does secretly runs through them;
  • mixed precision and Tensor Cores -- the free 2-3x speedup you're probably leaving on the table;
  • when the GPU is the wrong tool, and how a single stray .cpu() call can quietly wreck your training loop;
  • and the three flavours of multi-GPU -- data parallel, model parallel, pipeline parallel -- and which one you actually want (spoiler: usually the boring one).

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.10+) distribution with PyTorch (pip install torch) -- a GPU is a nice-to-have, but not required: we're after the concepts here, and those hold whether or not you own an RTX card;
  • You've been through episode #42 (PyTorch fundamentals, where we first met tensors and .cuda()) and episode #120 (model optimization, where we made models fast) -- today we go one layer deeper and look at the machine underneath the speed.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#125) - GPU Programming Basics

Last time, at the very end of episode #124, I admitted I'd been cheating. All those train steps in our CI/CD pipeline just... ran. I said nothing about where the heavy lifting actually happens, or why a training run that grinds for eight hours on your laptop finishes in eight minutes on the right hardware. Time to pay that debt. Because the whole reason modern AI is even possible lives one layer beneath everything we've built in this series so far -- down in the silicon that does thousands of multiplications at once.

Every time you type model.cuda() and training suddenly goes 50x faster, there is an absurd amount of parallel machinery whirring away under the hood. Most of it PyTorch hides from you, and honestly, thank goodness it does -- you do NOT want to be writing raw CUDA kernels to train a classifier. But understanding how the GPU works, even at a hand-wavy level, is one of those things that quietly turns you from someone who copies training scripts into someone who understands why a given script is slow. You'll know why a certain operation crawls, why you got that dreaded CUDA out of memory at 3am, and how to arrange your code so the hardware actually earns its keep. Let's crack it open.

CPU vs GPU: different tools, different jobs

Here's the mental model I want you to carry out of this episode. A CPU has a handful of big, clever, powerful cores -- somewhere between 8 and 64 on a decent machine -- each optimised to fly through complicated, branchy, sequential work as fast as humanly possible. A GPU takes the opposite bet entirely: thousands of small, comparatively dumb cores, each one no genius on its own, but all capable of doing the same operation on different data at the same time. A CPU is a few Formula 1 drivers. A GPU is ten thousand people on bicycles. For getting one important person across town, the F1 car wins every time. For delivering ten thousand pizzas across the same city -- well, now you want the bicycles ;-)

Neural networks are the pizza-delivery problem. Almost everything they do reduces to matrix multiplication -- the same multiply-and-add, repeated across millions of numbers with no dependency between them. That's the exact shape the GPU was built for. Let's measure it:

import torch
import time

def benchmark_matmul(size, device, n_runs=100):
    """Matrix multiplication: the bread and butter of neural networks."""
    a = torch.randn(size, size, device=device)
    b = torch.randn(size, size, device=device)

    if device == 'cuda':
        torch.cuda.synchronize()   # wait for setup before we start the clock

    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        c = torch.mm(a, b)
        if device == 'cuda':
            torch.cuda.synchronize()   # GPU work is async -- WAIT for it
        times.append(time.perf_counter() - start)

    median_ms = sorted(times)[len(times) // 2] * 1000
    return median_ms

# Typical results (varies wildly by hardware):
# Size 256:   CPU ~0.5ms,    GPU ~0.1ms    (GPU 5x faster)
# Size 1024:  CPU ~30ms,     GPU ~0.3ms    (GPU 100x faster)
# Size 4096:  CPU ~4000ms,   GPU ~5ms      (GPU 800x faster)

Look at that scaling and let it sink in. The GPU advantage is not fixed -- it grows with the size of the problem. At a tiny 256x256 the GPU is barely ahead, but crank it up to 4096x4096 and it's leaving the CPU somewhere near the horizon. There's a lesson buried in there about batch size: processing 256 samples in one fat batch is vastly more efficient than dribbling them through one at a time, because the big batch gives the GPU enough parallel work to actually fill all those cores.

And notice that torch.cuda.synchronize() call -- it matters more than it looks. GPU operations are asynchronous: when you call torch.mm, PyTorch queues the work and returns to Python immediately, before the GPU has finished (or even started). If you time GPU code without synchronising, you are timing how fast Python can ask, not how fast the GPU can answer. This is the number one way people accidentally "prove" their GPU code is instant. It isn't. It's just polite about hiding the wait ;-)

The CUDA programming model

CUDA (Compute Unified Device Architecture) is NVIDIA's platform for programming their GPUs. You will almost never touch it directly as an ML practitioner, but knowing its vocabulary makes error messages and profiler output stop looking like alphabet soup. The model organises work into a strict hierarchy:

# The CUDA execution hierarchy (PyTorch hides all of this from you)
cuda_model = """
GRID  (the entire job -- one kernel launch)
|
+-- BLOCK 0   (a group of threads that can share fast memory)
|     +-- Thread 0
|     +-- Thread 1
|     +-- ...
|     +-- Thread 255
+-- BLOCK 1
|     +-- Thread 0 ... Thread 255
+-- ...

Vocabulary you'll actually see in the wild:
- KERNEL : a function that runs ON the GPU, launched FROM the CPU
- GRID   : all the blocks for one kernel launch
- BLOCK  : 32-1024 threads that run together and can share memory
- THREAD : one single unit of execution (does one element of work)
- WARP   : 32 threads that execute in LOCKSTEP -- the real hardware unit
"""
print(cuda_model)

The one term worth tattooing on the inside of your eyelids is warp. Threads don't really run independently -- the hardware herds them into groups of exactly 32, called a warp, and every thread in that warp executes the same instruction at the same time. This is called SIMT (Single Instruction, Multiple Threads), and it explains a lot of weird GPU performance later on. For now, just remember: 32 is a magic number, and the GPU loves it when your work divides neatly into warps.

What does a kernel actually look like? Since CUDA kernels are written in a C dialect, here's one for adding two vectors, just so the abstraction has a face:

# What a real CUDA kernel looks like (this is C/CUDA, not Python):
cuda_kernel = """
__global__ void vector_add(float *a, float *b, float *c, int n) {
    // Each thread figures out WHICH element it is responsible for
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {          // guard: don't run off the end of the array
        c[idx] = a[idx] + b[idx];
    }
}

// Launch it: 1 million elements, 256 threads per block
// vector_add<<<(1000000 + 255) / 256, 256>>>(a, b, c, 1000000);
"""

# ...and here is that EXACT same computation, in PyTorch:
a = torch.randn(1_000_000, device='cuda')
b = torch.randn(1_000_000, device='cuda')
c = a + b          # PyTorch launches a CUDA kernel for you, internally

That's the whole trick, really. That one innocent line c = a + b compiles, under the hood, into a kernel launch just like the C above -- a million threads, each adding one pair of numbers, all at once. When you write output = model(input), PyTorch fires off hundreds of these kernels back to back: matrix multiplies, activation functions, normalisations, reductions. You get to think in tensors; the library thinks in warps. That division of labour is exactly why PyTorch won.

The memory hierarchy (this is where performance lives)

Right, if you remember one technical thing from this episode, make it this section. GPU memory is not one thing -- it's a layered pyramid, and each layer trades capacity for speed:

gpu_memory = {
    'registers': {
        'speed':    'fastest (basically zero latency)',
        'capacity': '~256 KB per SM (Streaming Multiprocessor)',
        'scope':    'per thread -- nobody else can see them',
        'use':      'local variables, loop counters',
    },
    'shared_memory': {
        'speed':    'very fast (~5 ns)',
        'capacity': '48-228 KB per SM',
        'scope':    'per block -- threads in the same block can share',
        'use':      'reused data, reductions, thread-to-thread comms',
    },
    'l2_cache': {
        'speed':    'fast (~50-100 ns)',
        'capacity': '6-50 MB total',
        'scope':    'all SMs',
        'use':      'automatic caching of global memory',
    },
    'global_memory': {
        'speed':    'slow (~400-800 ns)',
        'capacity': '8-80 GB -- THIS is what people call "GPU memory"',
        'scope':    'all threads, survives between kernels',
        'use':      'model weights, activations, your inputs and outputs',
    },
}

When someone whines "I need more GPU memory", they mean the bottom row -- global memory, the VRAM, the 24GB on the box. That's where your weights and activations live. But here's the counter-intuitive bit that trips up quit some people: raw compute is almost never the wall. A modern GPU can crunch 300+ TFLOPS of arithmetic but only move around 2 TB/s of data. Do the ratio. For a lot of operations, the cores finish their sums and then sit there twiddling their thumbs, waiting for the next chunk of data to arrive from global memory. That operation is memory-bound, and throwing more compute at it does exactly nothing.

This is the insight, so let me say it plainly: a huge fraction of real ML workloads are limited by memory bandwidth, not by how fast the GPU can multiply. It's why fusing operations helps (fewer trips to global memory), why float16 can double your speed (half the bytes to move), and why the shape of your tensors matters. Keep it in your back pocket. Meanwhile, you can always ask the GPU what it's holding:

if torch.cuda.is_available():
    dev = torch.cuda.current_device()
    props = torch.cuda.get_device_properties(dev)
    total     = props.total_memory        / 1e9
    allocated = torch.cuda.memory_allocated() / 1e9
    reserved  = torch.cuda.memory_reserved()  / 1e9

    print(f"Device:    {props.name}")
    print(f"Total:     {total:.1f} GB")
    print(f"Allocated: {allocated:.1f} GB   (tensors you actually hold)")
    print(f"Reserved:  {reserved:.1f} GB   (PyTorch's cache pool)")
    print(f"Free-ish:  {total - reserved:.1f} GB")

Nota bene: reserved is usually bigger than allocated, and that's on purpose -- PyTorch grabs memory from the driver in chunks and keeps a pool around so it doesn't have to keep asking (asking is slow). So when nvidia-smi says your process is using 20GB but your tensors only add up to 12GB, you're not leaking -- you're seeing the cache. torch.cuda.empty_cache() hands the spare back, though it rarely helps as much as people hope.

cuBLAS and cuDNN: the libraries doing the real work

You will basically never write a CUDA kernel by hand, and here's the honest reason: NVIDIA already wrote the fast ones, and you will not beat them over a weekend. PyTorch leans on two heavyweight libraries for the operations that matter:

# cuBLAS  = CUDA Basic Linear Algebra Subroutines
#   Every torch.mm(), torch.bmm(), and nn.Linear runs through cuBLAS.
#   It picks the fastest multiplication algorithm for YOUR exact
#   matrix shapes on YOUR exact GPU. You could not hand-tune this.

# cuDNN   = CUDA Deep Neural Network library
#   Every nn.Conv2d, nn.BatchNorm, nn.LSTM runs through cuDNN.
#   It can AUTO-TUNE: try several algorithms, keep the fastest.

# Switch cuDNN auto-tuning on, once, at startup:
torch.backends.cudnn.benchmark = True
# First forward pass gets SLOWER (it's benchmarking algorithms),
# every pass after that gets faster (it found the winner).

# Need bit-for-bit reproducibility instead (remember episode #119)? Then:
# torch.backends.cudnn.benchmark = False
# torch.backends.cudnn.deterministic = True

That cudnn.benchmark = True flag is close to free money for convolutional networks -- it can speed up convs by 2-3x. But read the fine print, because there's a catch that bites people. It tunes for a specific input shape. If your batch size or image resolution keeps changing, cuDNN has to re-benchmark every time it sees a new shape, and you'll spend all your time tuning and never reap the reward. This, by the way, is one more reason fixed input sizes are faster -- a theme that also showed up when we talked about serving in episode #121. Consistent shapes make the whole stack happy.

Mixed precision: the speedup you're probably not using

Back in episode #120 we made models smaller and faster with tricks like quantization. There's a sibling technique that lives squarely in GPU-land and costs you almost nothing to adopt: mixed precision training. The idea is that you don't need full 32-bit floats for most of the work. Modern GPUs have dedicated Tensor Cores that chew through 16-bit matrix multiplies several times faster than 32-bit ones -- and half the bit-width also means half the bytes to shovle around (remember, we're often memory-bound). PyTorch makes this genuinely a few lines:

from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()          # keeps tiny gradients from vanishing in fp16

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()

    with autocast():                       # run the forward pass in fp16
        preds = model(x)
        loss = loss_fn(preds, y)

    scaler.scale(loss).backward()          # scale UP so gradients survive
    scaler.step(optimizer)                 # unscale and step
    scaler.update()                        # adjust the scale for next time

Two things carry the weight here. autocast() automatically runs the safe operations (the big matrix multiplies) in 16-bit while quietly keeping the numerically touchy ones (like softmax and the loss) in 32-bit -- you get the speed without the instability. And GradScaler solves a subtle problem: 16-bit floats can't represent very small numbers, so tiny gradients would just round to zero and your model would stop learning. The scaler multiplies the loss up by a big factor before the backward pass (so the gradients land in a representable range), then divides back out before the optimizer step. On the right hardware this buys you a 2-3x speedup and roughly halves your memory use, for maybe six lines of code. If you're training on a modern NVIDIA GPU and not using mixed precision, you are leaving a genuine pile of performance on the floor.

When the GPU is the wrong tool

I've been selling the GPU hard, so let me be honest about where it falls flat on its face. It is NOT a magic "go faster" button -- there are whole categories of work where it loses to a humble CPU:

slow_on_gpu = {
    'tiny_tensors':     'Kernel-launch overhead (~10us) dwarfs the actual work',
    'sequential_ops':   'The GPU is a parallel beast -- serial code wastes it',
    'cpu_gpu_transfer': 'PCIe is the bottleneck: only ~12-32 GB/s across the bus',
    'random_access':    'GPU memory loves sequential reads, hates random ones',
    'branchy_code':     'if/else DIVERGENCE serialises a warp (see below)',
}

That last one deserves a word, because it's pure warp physics. Remember, 32 threads in a warp execute in lockstep. So what happens when your code says if (x > 0) do_A() else do_B() and half the warp goes one way and half the other? The hardware can't run both branches at once, so it runs branch A with the B-threads switched off, then runs branch B with the A-threads switched off. Both branches, back to back, at half efficiency. That's warp divergence, and it's why heavily branchy code (which the CPU handles beautifully) can crawl on a GPU.

But the single most common way people accidentally cripple their GPU code is the humble, innocent-looking data transfer. Watch:

if torch.cuda.is_available():
    x_cpu = torch.randn(100, 100)

    # (A) Pure CPU
    start = time.perf_counter()
    for _ in range(1000):
        y = x_cpu @ x_cpu
    cpu_time = time.perf_counter() - start

    # (B) GPU, but transferring data every single iteration (the SIN)
    start = time.perf_counter()
    for _ in range(1000):
        x_gpu = x_cpu.cuda()   # copy TO the gpu
        y = x_gpu @ x_gpu
        y_cpu = y.cpu()        # copy BACK -- ouch
    gpu_transfer = time.perf_counter() - start

    # (C) GPU, data already resident on the device
    x_gpu = x_cpu.cuda()
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(1000):
        y = x_gpu @ x_gpu
    torch.cuda.synchronize()
    gpu_resident = time.perf_counter() - start

    print(f"CPU:                  {cpu_time*1000:.1f} ms")
    print(f"GPU (with transfers): {gpu_transfer*1000:.1f} ms   <- often SLOWER than CPU!")
    print(f"GPU (data resident):  {gpu_resident*1000:.1f} ms")

Run that and the middle case will frequently be slower than plain CPU, which surprises everyone the first time. The compute on the GPU is lightning; the problem is you spent all your time hauling data back and forth across the PCIe bus, which is molasses compared to on-device memory. The moral -- and it's the practical takeaway of this whole episode -- is keep your data on the GPU. The biggest performance killers in real ML code are almost always transfer-related: a sneaky .cpu() inside a training loop, a Python for loop iterating over a GPU tensor element by element (each access is a mini-transfer), or preprocessing on the CPU that could have happened on the device. Move the data up once, do all your work, move the results down once. Having said that, don't over-rotate on tiny models either -- if your whole network is smaller than the transfer overhead, the CPU genuinely is the right answer, and there's no shame in it.

Going multi-GPU: three ways to split the work

Eventually one card isn't enough -- the model won't fit, or a training run that takes a week needs to take a day. You have three strategies, and they split the work along different axes:

import torch.nn as nn

# 1. DATA PARALLELISM -- same model on every GPU, split the batch.
#    This is the one you almost always want.
if torch.cuda.device_count() > 1:
    # DistributedDataParallel (DDP) is the good one -- use it.
    model = nn.parallel.DistributedDataParallel(model)
    # nn.DataParallel also exists but is slower -- avoid it for real work.

# How data parallelism works, conceptually:
#   GPU 0 processes batch[  0 : B/N ]
#   GPU 1 processes batch[ B/N : 2B/N ]
#   ... each computes gradients on its slice ...
#   gradients are AVERAGED across all GPUs (an "all-reduce")
#   every GPU updates to identical weights. Rinse, repeat.


# 2. MODEL PARALLELISM -- split the MODEL across GPUs (for models too big to fit).
class ModelParallel(nn.Module):
    def __init__(self):
        super().__init__()
        self.first_half  = nn.Linear(1024, 1024).to('cuda:0')  # lives on GPU 0
        self.second_half = nn.Linear(1024, 1024).to('cuda:1')  # lives on GPU 1

    def forward(self, x):
        x = torch.relu(self.first_half(x.to('cuda:0')))
        x = self.second_half(x.to('cuda:1'))   # hop the data across to GPU 1
        return x


# 3. PIPELINE PARALLELISM -- like model parallel, but keep every GPU BUSY.
#    While GPU 1 works on micro-batch #1, GPU 0 already starts micro-batch #2.
#    This is what GPipe and PipeDream implement.

For 95% of people the answer is data parallelism, and specifically DistributedDataParallel (DDP) rather than the older DataParallel. Why DDP? Because it runs one process per GPU, which sidesteps Python's GIL entirely, and it overlaps the gradient communication with the backward computation so the GPUs aren't sitting idle waiting to sync. The naive DataParallel runs everything from one process and bottlenecks on exactly the things DDP was built to avoid. Model and pipeline parallelism only come out to play when the model itself is too fat to fit on a single card -- and that's precisely how the billion-parameter language models from episode #60 get trained: you slice the model across many GPUs because no single card on Earth could hold the whole thing.

Profiling: measure, never guess

I'll leave you with the most important habit in performance work: you do not guess where the time goes, you measure it. Your intuition about what's slow is almost always wrong -- mine certainly is. PyTorch ships a profiler that tells you the truth:

if torch.cuda.is_available():
    model = nn.Sequential(
        nn.Linear(1024, 4096), nn.ReLU(),
        nn.Linear(4096, 4096), nn.ReLU(),
        nn.Linear(4096, 10),
    ).cuda()

    x = torch.randn(64, 1024, device='cuda')

    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        record_shapes=True,
    ) as prof:
        for _ in range(10):
            model(x)

    # Sort by total GPU time and show the worst offenders
    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

The output is a table of every operation and how much GPU time it ate. The surprises are where the value is. Common ones I've watched people discover, jaws slightly open: data loading eating 60% of total time (fix: more DataLoader workers, or preprocess on-GPU) while everyone was busy optimising the model that was never the problem; hundreds of tiny kernel launches for small operations (fix: fuse them, or bump the batch size); and memory allocation churn (fix: reuse buffers, lean on the caching allocator). This dovetails straight into episode #123's monitoring mindset -- in production you watch the model's behaviour, and in development you watch the hardware's behaviour, but it's the same discipline: instrument, look at real numbers, then act. Guessing is how you spend a day speeding up the 3% that didn't matter ;-)

TL;DR - what we cracked open today

  • GPUs win at parallel math -- thousands of small cores doing the same operation on different data -- and their advantage over CPUs grows with problem size, which is exactly why big batches are efficient;
  • CUDA organises work as a grid of blocks of threads, and 32 threads run in lockstep as a warp -- a fact that quietly explains warp divergence and a lot of GPU performance weirdness;
  • the memory hierarchy (registers -> shared -> L2 -> global) is where performance actually lives: many workloads are limited by memory bandwidth, not compute, so faster arithmetic won't save you;
  • cuBLAS and cuDNN are the hand-tuned libraries every fast PyTorch op runs through -- flip on cudnn.benchmark = True for a free conv speedup when your shapes are fixed;
  • mixed precision (autocast + GradScaler) uses Tensor Cores and 16-bit math for a 2-3x speedup and half the memory, for about six lines of code -- use it;
  • keep data on the GPU: CPU-GPU transfers over PCIe are the most common self-inflicted bottleneck in ML code, so a stray .cpu() in a loop can undo everything;
  • data parallelism via DistributedDataParallel is the multi-GPU strategy you want; model and pipeline parallelism are for when the model won't fit on one card at all.

Exercises

Three to get your hands (and your GPU) dirty before the next one. As always, I'll walk through full solutions next time.

  1. Find your own crossover point. Take the benchmark_matmul function and run it for matrix sizes 64, 128, 256, 512, 1024, 2048 on both CPU and GPU (if you have one; if not, reason it through). Plot or table the two curves and find the size at which the GPU overtakes the CPU. Write one sentence explaining why the GPU loses at small sizes -- and connect your answer back to the kernel-launch overhead we discussed.

  2. Catch a transfer bug in the act. Write a small training loop that does something silly on purpose: move a tensor to the GPU, then call .item() or .cpu() on a value inside the loop every iteration (a very common real bug). Time it. Now fix it so the transfer happens only once, outside the loop, and time again. Report the speedup and write a paragraph on why calling .item() inside a loop is secretly a synchronisation point that stalls the whole GPU pipeline.

  3. Turn on mixed precision. Take any small model you've built earlier in this series and wrap its training loop in autocast() + GradScaler as shown above. Measure training time and peak memory (torch.cuda.max_memory_allocated()) with and without it. Then deliberately remove the GradScaler while keeping autocast, train for a bit, and describe what goes wrong (hint: watch the gradients and the loss). Explain in a sentence what the scaler was protecting you from.

Right -- we now understand the machine that does the multiplying, how to feed it, and how to spread work across several cards in one box. But there's a ceiling hiding in plain sight in that last sentence: one box. What happens when the model, or the dataset, is so enormous that even a machine stuffed with eight GPUs isn't enough, and you have to make many separate computers cooperate on a single training run across a network? That's a whole different animal -- new bottlenecks, new failure modes, new tricks -- and it's exactly what we tackle next time.

De groeten!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
1 comments
avatar

Congratulations @scipio! You have completed the following achievement on the Hive blockchain And have been rewarded with New badge(s)

You received more than 45000 upvotes.
Your next target is to reach 50000 upvotes.

You can view your badges on your board and compare yourself to others in the Ranking
If you no longer want to receive notifications, reply to this comment with the word STOP

0
0
0.000
Reply