AI accelerators

1 · What it is

An AI accelerator is a chip designed to speed up the maths of AI models. Four families appear here: graphics cards (GPUs), Google’s tensor processing units (TPUs), neural processing units (NPUs) and Amazon’s Trainium and Inferentia.

The one sum AI keeps doing

Picture a photo app deciding that a picture shows a dog. The photo becomes a grid of numbers, and the model works through layers. A neural network is a model built from these layers of simple number steps. In each layer it takes a list of numbers, multiplies each one by a number it learned in training, and adds the results into a total. NVIDIA calls this multiply-add the most common operation in modern neural networks.

Here is a made-up example with tiny numbers. Say three inputs are 2, 5 and 1, and the learned numbers are 3, 1 and 4. The layer works out 2 times 3, plus 5 times 1, plus 1 times 4. That is 6 plus 5 plus 4, so the total is 15.

Many rows of these form a grid called a matrix. At the end, “dog” gets the highest score.

An accelerator is hardware built to run many of these sums at the same time.

Learning versus answering

AWS says it built Trainium to make fast AI training and inference cheaper at large scale. AWS made Inferentia chips to run trained AI models. AWS says Inferentia aims to be fast while keeping cost as low as it can on its cloud computers.

GPUs and TPUs

A GPU is built to do lots of work in parallel, which means side by side. It stays busy by running many threads at once. A thread is one small stream of work. NVIDIA added Tensor Cores to its GPUs. These are special parts of the chip that speed up grid-shaped multiply-and-add work.

Google’s TPUs are custom chips built to speed up machine learning, the way computers learn from examples. Their matrix unit, the part that does grid sums, is a 128 by 128 grid of multiply-add hardware. The software cuts big matrix jobs into blocks that fit it.

Think of a giant times-table sheet that is too big for one desk. You cut it into square pieces, hand each piece to a helper, and then combine their answers. Google describes the TPU software doing something like this: it tiles a big matrix multiply into smaller blocks so the matrix unit can work through them efficiently.

Google’s first TPU, described in a 2017 paper, was made to speed up inference. Compared with a CPU and GPU of the time, on Google’s own neural network work, the TPU was on average about 15 to 30 times faster. It also did far more work for each unit of electricity used.

NPUs in PCs

Accelerators also sit inside personal computers. Microsoft’s Copilot+ PCs are built around an NPU. The NPU uses less energy for AI work than a CPU or GPU, so batteries last longer.

Limits

An NPU needs software written to use it. Some models fit a chip badly. A batch is a group of inputs handled together; its shape is the size of its number grids. A TPU model is prepared, or compiled, for batches of one shape. If later batches come in a different shape, it stops working.

Each step has two parts: fetching numbers from memory and doing the sums. If fetching takes longer than the sums, the job is memory limited. Then the sum hardware spends its time waiting for numbers to arrive.