What the hardware is doing

Everything above this layer runs on a GPU, and a rough model of what one is doing answers the two questions you will be asked constantly: why is this slow, and will it fit. This section covers why matrix maths suits thousands of simple cores, why memory and the bandwidth to it usually bind before arithmetic does, what mixed precision buys and what it costs, and how to read nvidia-smi without being misled — high utilisation can mean the device is busy doing very little.

Many small operations at once#

A CPU has a few very capable cores. A GPU has thousands of simple ones. Matrix multiplication is thousands of independent multiply-and-add operations, which is exactly the shape a GPU is built for.

It is worth being specific about what 'capable' and 'simple' mean, because the difference is not that GPU cores are slower. A CPU core spends most of its transistors on making a single sequence of instructions finish sooner: predicting which branch will be taken, reordering instructions, and holding several megabytes of cache per core so it rarely waits for memory. All of that is machinery for code where the next step depends on the last one. A GPU spends those transistors on more arithmetic units instead, and handles waiting by keeping thousands of threads in flight so there is always one ready to run while others wait.

That bargain only pays when the work is thousands of identical independent operations with no branching. Multiplying two matrices is exactly that — every output element is an independent dot product of a row and a column — which is why the same hardware serves graphics and neural networks. Give a GPU a loop with a data-dependent branch in it and it will lose to a CPU comfortably.

Two consequences follow immediately and explain a lot of otherwise mysterious behaviour. Anything you write as a Python loop over items is the wrong shape and will be slow no matter what device it runs on, because it is a sequence of tiny operations rather than one large one. And a small operation on a GPU costs roughly the same as a medium one, because the fixed cost of launching the work dominates — which is why a batch of 1 and a batch of 8 often take almost the same time.

Memory is usually the constraint#

• VRAM decides whether a model fits at all. Exceeding it fails immediately rather than degrading.
• Bandwidth decides speed for large models, because generation is dominated by reading weights.
• Moving data between host and device is slow. Do it once per batch, never per item.

The second point is the one people find surprising, and it has a simple form. Every operation has an arithmetic intensity: the number of floating-point operations it performs divided by the number of bytes it has to move to and from memory. A chip has a ratio too — its peak arithmetic rate divided by its memory bandwidth. If your operation's intensity is below the chip's ratio, you are memory-bound and buying more arithmetic buys you nothing.

The gap between the two regimes is much larger than intuition suggests. There is no GPU on the machine these words were written on, so here it is measured on its CPU, where the same physics applies with smaller numbers.

matmul  1024x1024     1,194 GFLOP/s      intensity  171 FLOP/byte
matmul  4096x4096       769 GFLOP/s      intensity  683 FLOP/byte
add     4096x4096         4.0 GFLOP/s    intensity 0.08 FLOP/byte    48 GB/s
add     8192x8192         4.0 GFLOP/s    intensity 0.08 FLOP/byte    48 GB/s

This single ratio explains most of what you will be asked about performance. Training with a large batch does big matrix multiplications and is compute-bound. Generating one token at a time reads every weight in the model to do a handful of arithmetic per weight, so it is heavily memory-bound, which is why token generation speed tracks memory bandwidth rather than FLOPs, and why batching many requests together is the largest single lever on inference cost. And it is why fusing operations helps: the arithmetic is unchanged, the round trips to memory are what disappeared.

What is actually in the memory during training#

Ask someone how much memory it takes to train a model and they will usually answer with the parameter count times a couple of bytes. That answer is wrong by roughly a factor of eight, and knowing why is the difference between planning a run and being surprised by it. There are four things in memory, and the weights are the smallest of them.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(1024, 4096), nn.GELU(), nn.Linear(4096, 1024))
opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
model(torch.randn(8, 1024)).pow(2).mean().backward()
opt.step()                                   # the state exists only after a step

n = sum(p.numel() for p in model.parameters())
params = sum(p.numel() * p.element_size() for p in model.parameters())
grads = sum(p.grad.numel() * p.grad.element_size() for p in model.parameters())
state = sum(v.numel() * v.element_size()
            for s in opt.state.values() for v in s.values() if torch.is_tensor(v))
for label, b in [("parameters", params), ("gradients", grads), ("Adam state", state)]:
    print(f"{label:12} {b / 2**20:8.1f} MiB   {b / n:5.1f} bytes per parameter")

# parameters       32.0 MiB     4.0 bytes per parameter
# gradients        32.0 MiB     4.0 bytes per parameter
# Adam state       64.0 MiB     8.0 bytes per parameter
#                                16.0 in total, and the model is 4 of them

In mixed precision the bookkeeping changes but the total does not, which catches people out. Weights and gradients drop to 2 bytes each, but a 32-bit master copy of the weights has to be kept for the optimiser to update accurately, and Adam's two moments stay in 32-bit. Two plus two plus four plus four plus four is still sixteen bytes per parameter. Mixed precision buys speed and cheaper activations; it does not shrink the optimiser.

A 6.74B-parameter model (llama-7B shape), mixed precision, AdamW:

  bf16 weights          2 B/param     12.6 GiB
  bf16 gradients        2 B/param     12.6 GiB
  fp32 master weights   4 B/param     25.1 GiB
  Adam m                4 B/param     25.1 GiB
  Adam v                4 B/param     25.1 GiB
                       --------------------------
  total                16 B/param    100.4 GiB     before a single activation

  optimiser's share: 75 per cent.
  inference on the same model: 12.6 GiB.

The fourth item is activations: everything the forward pass computed that the backward pass needs in order to differentiate. Unlike the other three it does not scale with the parameter count at all — it scales with batch size and sequence length, which means it is the one you control.

Measured by summing every layer output in a 162M transformer, fp32:

  batch 1, 256 tokens      150.3 MiB
  batch 2, 256 tokens      299.9 MiB
  batch 4, 256 tokens      599.1 MiB
  batch 8, 256 tokens    1,197.4 MiB      -> 149.7 MiB per sample, flat

So the batch size you can use is decided by what is left after the fixed 16 bytes per parameter, and the three ways out are all forms of the same trade. Activation checkpointing stores only some activations and recomputes the rest in the backward pass, buying a large memory saving for roughly a third more compute. Gradient accumulation runs several small batches and steps once, so the effective batch is large while the resident activations are small. Sharding the optimiser state across devices attacks the other 12 bytes. Reach for them in that order.

Mixed precision#

Using 16-bit numbers for most operations and 32-bit where precision matters roughly halves memory and can double throughput, with essentially no quality loss. It is close to free and frequently left switched off.

There are two 16-bit formats and they are not interchangeable. Both use 16 bits; they spend them differently, one on range and the other on precision, and which one you have decides whether you need loss scaling.

import torch

for dt in (torch.float32, torch.float16, torch.bfloat16):
    i = torch.finfo(dt)
    print(f"{str(dt):16} max {i.max:.2e}  smallest normal {i.tiny:.2e}  eps {i.eps:.2e}")

# torch.float32    max 3.40e+38  smallest normal 1.18e-38  eps 1.19e-07
# torch.float16    max 6.55e+04  smallest normal 6.10e-05  eps 9.77e-04
# torch.bfloat16   max 3.39e+38  smallest normal 1.18e-38  eps 7.81e-03

g = torch.tensor([2.0e-8, 1.0e-5, 3.0e-3])          # plausible gradient magnitudes
show = lambda t: [f"{v:.2e}" for v in t.tolist()]
print(show(g.half().float()))
print(show(g.bfloat16().float()))
print(show((g * 1024).half().float() / 1024))       # the same in fp16, loss-scaled

# ['0.00e+00', '1.00e-05', '3.00e-03']   fp16: the first gradient became zero
# ['2.00e-08', '1.00e-05', '3.01e-03']   bf16: kept, with range to spare
# ['2.00e-08', '1.00e-05', '3.00e-03']   fp16 scaled by 1024 first: kept

print(float(torch.full((65536,), 1.0, dtype=torch.float16).sum(dtype=torch.float16)))
# inf  -- fp16 cannot represent 65536; bfloat16 and float32 both return 65536.0

The rule follows directly. On hardware with bfloat16 support, use it and do not use a gradient scaler — there is nothing to scale for. On older hardware you have float16, and you need the scaler: multiply the loss by a large constant so small gradients land inside the representable range, then divide it back out before the optimiser step. PyTorch's GradScaler does that and also skips any step whose gradients came back as infinity, adjusting the constant as it goes.

# Needs a CUDA device. model, opt, loss_fn, batch and target are yours.

# On modern hardware, with bfloat16: no scaler needed.
with torch.autocast("cuda", dtype=torch.bfloat16):
    loss = loss_fn(model(batch), target)
loss.backward()
opt.step()
opt.zero_grad()

# On older hardware, with float16: the scaler is required.
scaler = torch.amp.GradScaler()

with torch.autocast("cuda", dtype=torch.float16):
    loss = loss_fn(model(batch), target)

scaler.scale(loss).backward()      # multiply the loss so gradients stay in range
scaler.step(opt)                   # unscale, then skip the step if any gradient is inf
scaler.update()                    # adjust the scale for next time
opt.zero_grad()

Utilisation is not the whole story#

nvidia-smi                       # memory used, utilisation, temperature
watch -n 1 nvidia-smi            # live
nvidia-smi --query-gpu=memory.used,memory.total,utilization.gpu,power.draw \
           --format=csv -l 1     # just the four numbers, once a second

Four numbers are worth reading, and the useful skill is knowing what each one rules out rather than what it proves. Memory used against memory total tells you how much headroom the batch size has. Utilisation tells you what fraction of the last sampling interval had any kernel running. Power draw is the honest proxy for whether the arithmetic units are busy, because a GPU doing real matrix multiplication draws close to its rated power and one shuffling memory around does not. Temperature matters only when the card is throttling.

utilisation   power      what it means                      where to look

  low         low      the GPU is waiting                 data loading, or Python
                       for work                           between kernel launches
  high        low      kernels run constantly but         many tiny operations,
                       do little arithmetic               or a memory-bound workload
  high        high     the hardware is genuinely busy     the model, the batch size,
                                                          or accept it and scale out
  any        at cap    thermal or power throttling        cooling, or the power limit

So the diagnostic order is: check memory first, since it decides what is possible at all; then check utilisation, since a low number is unambiguous and almost always means data loading; and only then reach for a profiler, because a high utilisation number tells you nothing on its own.

How deep to go#

For engineering, this level is sufficient: what fits, what makes it slow, and how to see it. Writing CUDA kernels is a specialisation that very few roles require, and it is a poor place to start.

What is worth memorising is short: sixteen bytes per parameter to train with Adam, two bytes to serve; activations scale with batch and sequence and nothing else; a chip balances at roughly a hundred and fifty FLOPs per byte, so anything below that is memory-bound; and bfloat16 has range while float16 has precision. What can be looked up on the day: the datasheet numbers for any particular card, the CUDA and driver version matrix, the flag names for any given framework's mixed precision, and everything about writing kernels.