Hardware and acceleration

Overview

Almost every performance question in machine learning reduces to one of two numbers: how many bytes the accelerator can move per second, and how many floating-point operations it can do per second. The ratio of the two, together with the ratio of operations to bytes in your workload, tells you whether you are compute-bound or memory-bound before you write a line of code. The roofline model (Williams, Waterman and Patterson, 2009) is that argument drawn as a graph, and it is the single most useful idea in this hub.

This page is for engineers choosing hardware, sizing a job or wondering why a GPU reports 10 percent utilisation while training is slow. It covers the main accelerator families, the memory hierarchy, mixed precision and custom kernels. Buying advice goes stale, so the numbers here are the vendors’ published specifications at the time of writing and are given to show the shape of the problem, not to recommend a part.

Key ideas

Bandwidth is the budget

Take an NVIDIA H100 SXM. NVIDIA’s datasheet lists 80 GB of HBM3 at 3.35 TB/s and 989 TFLOPS of dense BF16 tensor-core throughput. Dividing those gives about 295 operations per byte. A matrix multiply with large batch sizes does far more than that per byte loaded and is compute-bound, which is why training saturates the chip. Decoding one token from a 7 B parameter model in bf16 reads 14 GB of weights to do roughly 14 GFLOP of work, an intensity of about one operation per byte, so it is bound entirely by memory: at most 3.35 TB/s divided by 14 GB, about 240 tokens per second, no matter how fast the arithmetic is. Batching requests, quantising weights and keeping the KV cache small all attack this same limit, which is why /topics/mlops-and-deployment/ and /topics/efficient-ml/ spend so long on them.

The same arithmetic ranks consumer parts. NVIDIA lists the RTX 4090 at 24 GB of GDDR6X and 1,008 GB/s. Apple lists the M4 Max at up to 128 GB of unified memory at 546 GB/s and the M3 Ultra at up to 512 GB at 819 GB/s. A Mac with 128 GB can hold a model that a 24 GB GPU cannot, and will decode it at roughly half the speed per token.

The accelerator families

GPUs are the general choice: thousands of simple cores, tensor cores for matrix multiply, mature compilers and the largest library ecosystem. NVIDIA dominates through CUDA. AMD’s Instinct parts run PyTorch through ROCm with a thinner but improving stack. Google’s TPUs are systolic-array machines built for dense matrix multiply, programmed through JAX and XLA, and sold only as cloud instances. Google’s documentation lists the TPU v6e at 32 GB of HBM at 1,640 GB/s per chip, with chips connected in pods for large training runs. Apple Silicon puts CPU, GPU and a neural engine on one die with a shared memory pool, so a model can be larger than any discrete consumer GPU’s memory, at the cost of lower bandwidth and no CUDA. NPUs in phones and laptops are small fixed-function matrix engines aimed at a few watts, good for on-device inference of models already quantised to int8 or lower. FPGAs sit at the far end: they are reconfigurable logic that wins on latency-bound, odd-precision or sensor-attached inference and loses on anything that looks like a large dense matrix multiply. The sister site has a full hub on that case, linked below.

Mixed precision and kernels

Training in full fp32 wastes most of the tensor cores. Mixed precision (Micikevicius et al., 2017) keeps a master copy of weights in fp32 and runs the forward and backward passes in 16-bit. bf16 keeps fp32’s exponent range and needs no loss scaling, which is why it replaced fp16 on hardware that supports it. FP8, introduced with the Hopper generation, halves the bytes again for the matrix multiplies, with per-tensor scaling to stay in range.

Kernels are where the remaining factor of two or three hides. A fused kernel does several operations per trip to memory instead of one. FlashAttention (Dao et al., 2022) is the canonical example: it computes exact attention in tiles that stay in shared memory and never materialises the full attention matrix. torch.compile and Triton make fused kernels reachable without writing CUDA.

Where to start

Read /learn/beginner/04-training-loop-and-loss/ and /learn/beginner/07-your-first-neural-network/ first, then time the training loop from /reference/pytorch-training-loop/ on CPU, on MPS or CUDA, and in bf16 autocast, and compare. /reference/tensor-shapes-and-broadcasting/ covers the layout questions that decide whether a kernel is fast. The /projects/lora-fine-tune-apple-silicon/ project is a realistic unified-memory workload. The /topics/deep-learning/ hub covers the models these numbers apply to.

Tools and platforms

Common mistakes

• Comparing peak TFLOPS. Peak assumes dense matrix multiply with every unit busy. Compare measured throughput on your workload, and check whether a quoted figure includes structured sparsity.
• Buying for compute when the job is memory-bound. For inference, bandwidth and capacity decide, and for training at scale, interconnect does.
• Training in fp16 without loss scaling, or in fp32 on hardware with bf16 tensor cores.
• Reading GPU utilisation as a measure of efficiency. A kernel waiting on memory shows 100 percent utilisation.
• Running a data loader on two CPU workers and blaming the GPU.
• Assuming a model that fits in memory will run. The KV cache, optimiser state and activations need room too.

Going deeper

• Williams, Waterman and Patterson, Roofline: an insightful visual performance model for multicore architectures, Communications of the ACM, 2009.
• Micikevicius et al., Mixed Precision Training, 2017.
• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022.
• NVIDIA H100 Tensor Core GPU datasheet, source of the capacity and bandwidth figures above.
• Google Cloud TPU system architecture documentation, source of the TPU figures.
• Sister site: the FPGA Guide hub on AI and ML acceleration on FPGAs, covering dataflow and overlay accelerators, quantisation for FPGAs and when an FPGA beats a GPU or NPU.