Understanding Transformers · Part 16 of 16
Hardware Accelerators: GPUs, TPUs, and JAX
The chips built to multiply matrices fast, the operation behind most of a transformer's arithmetic.
Matrix multiplication
Look back at the code in the earlier parts: the Q, K, and V projections, Q KT, weights @ V, the feedforward's W₁x and W₂h, the final projection to logits. Each one is a matrix multiplication, and they account for most of a transformer's arithmetic. Softmax, layer normalization, and the embedding lookup are the exceptions. So the hardware underneath matters as much as the architecture on top of it: a chip that multiplies matrices quickly speeds up most of the model without changing a line of it.
Why a naive loop is slow
A matrix multiply of two n×n matrices is n3 individual multiply-and-add operations. Written as Python loops, each multiply-add goes through the interpreter one at a time. Written as a single vectorized call, the whole product goes to an optimized linear-algebra library (BLAS), which applies SIMD instructions that handle several numbers per cycle, uses several cores, and orders memory access to stay in cache. Same chip, same math, orders of magnitude apart:
import numpy as np, time
n = 200
A, B = np.random.randn(n, n), np.random.randn(n, n)
# Naive: one multiply-add at a time
start = time.time()
C = np.zeros((n, n))
for i in range(n):
for j in range(n):
s = 0.0
for k in range(n):
s += A[i, k] * B[k, j]
C[i, j] = s
loop_time = time.time() - start
# Vectorized: the same n^3 multiply-adds, done in bulk
start = time.time()
C_fast = A @ B
vector_time = time.time() - start
# loop_time is usually thousands of times vector_time (about 11,000x in one run on a laptop)Every NumPy snippet in these lessons has quietly relied on that gap: A @ B instead of a triple loop is already an appeal to hardware built for bulk arithmetic. GPUs and TPUs take the same idea much further.
GPUs
A CPU has a handful of powerful cores optimized for varied, branching, sequential work, the kind that runs a web server or a Python interpreter. A GPU was originally built to compute color values for millions of independent pixels at once, which turned out to be the same shape of problem as multiplying large matrices: enormous numbers of independent, identical arithmetic operations with little to no branching. A modern GPU has thousands of small arithmetic cores that all execute the same instruction on different data simultaneously (an approach called SIMT: single instruction, multiple threads). Recent NVIDIA GPUs also include tensor cores, units that multiply small matrix tiles directly. Handing a matmul to a GPU means splitting its n3 multiply-adds across thousands of cores instead of a handful, which is why training and running large transformers moved to GPUs once they got big enough to need it.
TPUs
A TPU (Tensor Processing Unit, Google's accelerator) goes a step further: instead of a general-purpose parallel chip that happens to be good at matmul, it's an ASIC, a chip designed for machine-learning workloads. Its matrix unit is a systolic array: a grid of multiply-accumulate units wired directly to their neighbors, so partial results flow from one unit to the next without being written back to memory between steps. Memory movement is often what limits speed at scale, and the systolic array avoids it for the matmul-heavy pattern a transformer produces. The trade-off is generality: Google's documentation says Cloud TPUs are poorly suited to programs with frequent branching or many element-wise operations, and a GPU handles those better. For large-scale transformer training and inference, a TPU can be more efficient per watt.
JAX and compilation
PyTorch code, including every PyTorch line in the companion notebook, runs eagerly by default: each operation executes the instant Python reaches it, one kernel launch at a time. JAX, also from Google, is built around a different idea. jax.jit traces a whole function the first time it is called with a given input shape, hands it to the XLA compiler (the compiler that turns the computation into TPU machine code, and also targets GPUs and CPUs), and gets back one compiled, fused program. Fusing a dozen small operations into one larger kernel cuts both the number of kernel launches and the amount of intermediate memory that has to be written out and read back: the same memory-movement bottleneck the systolic array attacks in silicon, handled here in software.
JAX offers other function transformations too. jax.grad takes an ordinary function and returns a new function computing its gradient, the same calculus from Part 2 expressed as a transformation, where PyTorch's .backward() walks a graph built while the forward pass ran. jax.vmap takes a function written for a single example and turns it into one that runs across a batch, with no explicit loop. The companion notebook's JAX training step uses it in place of a hand-written batch loop. jit, grad, and vmap compose: that training step wraps a vmap'd, gradient-taking function in a single jax.jit, and the whole thing compiles into one program.
Compiling through XLA has a side effect: the same program can target CPU, GPU, or TPU. PyTorch also picks a device with one line (device = torch.device("cuda" if torch.cuda.is_available() else "cpu"), in the companion notebook), so portability is no reason to choose JAX. What JAX makes the default is writing compiled, differentiable, auto-batched code.