TPU Explained from the Ground Up, Part 1: Why Is TPU Computing So Cost-Efficient?

From AllReduce to 3D Torus

When training large AI models, a single chip is usually not enough. The dataset and model parameters are too large, and the memory of a single TPU cannot hold everything. Therefore, multiple TPUs need to be organized into a cluster and compute together.

However, a cluster is not just many machines placed together.

When multiple machines are running, the first problem to solve is: after each machine finishes computing its own part, how should the results be synchronized?

If the synchronization cost is too high, it does not matter how fast the chips are. A large amount of time will be stuck in network communication.

One key reason TPU clusters can be cost-efficient is that they do not rely only on expensive networking equipment to brute-force bandwidth. Instead, they reduce synchronization cost through collective algorithms and network topology design.

So what exactly is being synchronized inside a TPU cluster?

The basic idea of distributed computing is simple: split a large piece of data into many smaller partitions and assign them to different machines for computation.

For example, in AI training, the input can be split into multiple batch shards, and different TPUs process different batch shards. After each machine finishes computing its own part, the results need to be merged in some way. Otherwise, each machine only knows its own local result, and the cluster does not share a common state.

When talking about data synchronization in distributed computing, the most intuitive concept is MapReduce.

What is MapReduce? Does AI need MapReduce?

Let’s first build the intuition using SQL.

Assume we want to compute:

select uid, sum(amount)
from orders
group by uid;

When the orders table is very large, MapReduce first partitions the data and distributes it to different machines for local computation. At this stage, each machine may hold partial sums for different uids. Then the system enters the shuffle phase.

The purpose of shuffle is to send data with the same uid to the same reducer. Finally, each reducer computes the result for the keys it is responsible for.

For example:

sum | uid | machine
10  | 1   | 1
20  | 2   | 2
30  | 3   | 3

In other words, in MapReduce, different machines, or reducers, only hold the complete results for the keys they are responsible for.

AI training is different. It needs AllReduce.

That means after the reduce step, the result still needs to be broadcast to everyone, so every machine has the aggregated result for all keys.

That is:

sum       | uid     | machine
10,20,30  | 1,2,3   | 1
10,20,30  | 1,2,3   | 2
10,20,30  | 1,2,3   | 3

Therefore, you can first understand AllReduce as:

MapReduce + broadcast result to all machines

But why does AI training need AllReduce?

An AI model can first be simplified as:

y = w1*x1 + w2*x2 + ... + wN*xN = wT x

x is the input, w represents the model weights, and y is the model output. The training objective is to adjust w so that the model output y becomes closer and closer to the correct answer.

The adjustment process is roughly:

1. Input x and compute the prediction y.

2. Use a loss function to measure the gap between the prediction and the correct answer.

3. Use differentiation to calculate how much each weight affects the loss. This is the gradient.

4. Use the gradient to update the weights, so the loss becomes smaller in the next iteration.

Note: More precisely, the gradient represents how much the loss changes when a certain weight changes, either increasing or decreasing.

When there is a large amount of training input data, the input is split into multiple batch shards and assigned to different TPUs. Each TPU holds the same AI model, but processes a different batch shard. Therefore, each TPU computes a different local gradient vector.

For example, three TPUs may compute:

machine | w1_gradient | w2_gradient | w3_gradient
a       | 1           | 2           | 3
b       | 3           | 2           | 1
c       | 5           | 3           | 1

If each TPU directly uses its own local gradient to update the model, the models on the three TPUs will quickly diverge.

Therefore, before updating the model, all local gradients need to be aggregated into one global gradient vector, and every TPU needs to receive the same result.

This is similar to:

select weight_id, avg(gradient_value)
from gradient_vectors
group by weight_id;

Then broadcast the result to all machines.

That is:

machine | w1_gradient_avg | w2_gradient_avg | w3_gradient_avg
a       | 3               | 2.33            | 1.67
b       | 3               | 2.33            | 1.67
c       | 3               | 2.33            | 1.67

After that, each machine uses the same gradient vector to update its local model, ensuring that all models remain consistent.

But if every training step needs gradient synchronization, is simple broadcast practical?

The most intuitive way to implement AllReduce is all-to-all broadcast:

Each machine sends its full gradient vector to every other machine, then sums or averages the vectors it receives.

Assume there are n machines and each gradient vector has size D.

Each machine needs to receive roughly (n - 1) * D amount of data. If the number of TPUs increases, n becomes larger. If the model becomes larger, the gradient vector size D also becomes larger.

Therefore, all-to-all communication quickly explodes network bandwidth requirements and increases cost. This is especially problematic because gradient synchronization happens in every training step, so network communication can become the main bottleneck.

To solve this problem, we need a different AllReduce algorithm: Ring AllReduce.

What is Ring AllReduce?

Ring AllReduce is a collective communication algorithm.

Collective communication means that multiple devices perform the same type of communication operation at the same time.

Ring AllReduce is split into two steps: ReduceScatter and AllGather.

Assume there are three machines physically connected as a ring:

a -> b -> c -> a

The three machines each compute their own gradient vector:

      w1 w2 w3
a:   [1, 2, 3]
b:   [3, 2, 1]
c:   [5, 3, 1]

ReduceScatter first splits the gradient vector into shards. For example:

[w1, w2, w3] => [w1], [w2], [w3]

It also defines which machine will finally hold the reduced result for each shard.

For example:

a is responsible for w1
b is responsible for w2
c is responsible for w3

Then all machines decide which shard should be sent to their neighbor based on the current network topology and the final reducer of each shard. When a neighbor receives a shard, it adds its own local value to the shard, producing a partial sum, and forwards it to the next neighbor.

The logic is roughly:

a first identifies that the node at distance n - 1 is c, and c is responsible for the w3 shard. Therefore, a sends its w3 value to its neighbor b.

b receives it, adds its own w3 value to a’s w3, and forwards the partial sum to c.

Finally, c receives a + b for w3, adds its own w3, and obtains the final reduced result for w3, which is 3 + 1 + 1 = 5.

Collective communication means that a, b, and c perform this type of shard transfer at the same time.

First communication round:

a - w3(a) -> b   w3(a+b)
b - w1(b) -> c   w1(b+c)
c - w2(c) -> a   w2(c+a)

Second communication round:

b - w3(a+b) -> c   w3(a+b+c)
c - w1(b+c) -> a   w1(b+c+a)
a - w2(c+a) -> b   w2(c+a+b)

After ReduceScatter finishes, each machine only holds the reduced shard it is responsible for:

a: w1(b+c+a) = 1 + 3 + 5 = 9
b: w2(c+a+b) = 2 + 2 + 3 = 7
c: w3(b+c+a) = 3 + 1 + 1 = 5

Then AllGather begins. Each machine sends the reduced shard it holds along the ring, so the other machines can receive it.

First communication round:

a - w1 -> b
b - w2 -> c
c - w3 -> a

Second communication round:

a - w3 -> b
b - w1 -> c
c - w2 -> a

Finally, every machine receives the complete result:

a: [9, 7, 5]
b: [9, 7, 5]
c: [9, 7, 5]

The benefit of Ring AllReduce is that each machine only sends one shard of the gradient vector in each communication round, instead of sending the full vector.

For n machines and a vector of size D, the per-machine traffic of Ring AllReduce is approximately:

2 * (n - 1) * (D / n)

When n is large, this is approximately close to:

2*D

Compared with all-to-all communication, where each machine receives (n - 1) * D, Ring AllReduce significantly reduces the amount of data each machine needs to handle.

However, Ring AllReduce still has limitations.

If all machines use a single global ring, each machine only has two neighbors: the previous node and the next node. Data can only be passed along this fixed ring to the next neighbor. As the number of machines increases, the number of communication rounds also increases.

In other words, Ring AllReduce reduces the amount of data transferred, but if the topology is only one long ring, the number of communication rounds still limits overall efficiency.

So how does TPU reduce the number of communication rounds?

Instead of using a single global ring, TPU uses 2D or 3D torus topology.

For example, a 3D torus allows each machine to have neighbors in three different dimensions: x, y, and z. Each machine can have six neighbors in total. Therefore, in each communication round, messages can be transmitted to multiple neighbors along different dimensions.

Let’s use a 2D torus with dimensions x and y as an example to see how having more neighbors can reduce communication rounds.

Assume there are four machines, and each machine has neighbors in both the x and y directions:

a —— b  X-axis
|    |
c —— d  X-axis
Y-axis

Note: The diagram looks like a 2D mesh, but it is actually a 2D torus. The difference is that a torus wraps around like a ring. For example, the top row is not simply a - b; it is effectively a - b - a.

In a single global ring, each machine usually pushes data to only one fixed neighbor in each round.

But in a 2D torus, a machine can send data along both the x and y dimensions. Assume the gradient vector length is four, and during ReduceScatter:

a is responsible for w1
b is responsible for w2
c is responsible for w3
d is responsible for w4

For a to collect the complete w1 result, the process can be:

First communication round, parallel along the X-axis:

b -> w1(b) -> a   (a+b)
d -> w1(d) -> c   (c+d)

Second communication round, parallel along the Y-axis:

c -> w1(c+d) -> a   (a+b+d+c)

With only two communication rounds, a can collect the complete w1 result. A single global ring would need three rounds.

Note: The example only demonstrates the path for w1. In the actual ReduceScatter process, w2, w3, and w4 are also being transmitted at the same time along other links.

A 3D torus provides even more dimensions and neighbors. The collective algorithm can use more parallel paths to transmit data. In a large-scale cluster, this can effectively reduce the number of communication rounds.

Therefore, the key reason TPU clusters can be cost-efficient is not that they rely on expensive network devices to brute-force bandwidth, but that they reduce synchronization cost through:

Ring AllReduce, which reduces the amount of data each machine needs to transmit.

2D / 3D torus topology, which provides more multidimensional neighbors and increases parallel data movement.

However, topology only explains how TPU chips are connected. To actually run computation on the cluster, we also need the compiler and runtime.

How is computation executed over a TPU cluster network topology?

We can use Hive as an analogy.

In the early days, if you wanted to write distributed computation with MapReduce, you had to write mappers and reducers yourself. The MapReduce runtime would distribute different data partitions to different machines and execute your mapper and reducer logic.

But if you wanted to write a more complex query, such as:

select uid, sum(amount)
from orders
join products on orders.pid = products.id
where products.name = 'XXX'
group by uid;

Writing mappers and reducers by hand would be very cumbersome. Therefore, Hive provided a higher-level interface: users write SQL, and Hive compiles the SQL into MapReduce stages.

TPU training has a similar abstraction layer.

Users write high-level tensor computations in TensorFlow or JAX, such as forward pass, loss computation, backward pass, and optimizer update. However, TPU hardware itself only supports lower-level tensor operations, such as matrix multiplication, reduce, AllReduce, and tensor load/store.

Therefore, XLA is needed as the compiler. XLA converts the high-level tensor computation graph into HLO, which stands for High Level Optimizer IR. It then optimizes the graph and generates an executable program for TPU.

Similar to Hive, XLA is not just doing syntax translation. It also performs optimizations such as sharding, operator fusion, device placement, and collective communication planning.

For example, in distributed TPU training, XLA can arrange collective operations such as ReduceScatter and AllGather based on the device mesh and TPU network topology information.

Therefore, the essence of TPU design is a co-design of hardware, network, and compiler. It is not merely a standalone accelerator chip.

Preview for the next article:

Why is TPU computation so fast? Understanding ICI and internal data movement inside a TPU Pod.

This article focuses on why TPU computation can be cost-efficient.

The key point is not that a single TPU chip is magically cheap. Instead, TPU clusters reduce the network synchronization cost of each training step through Ring AllReduce, 2D / 3D torus topology, and the collective schedule planned by XLA and the runtime.

However, this article only covered topology and computation runtime. The next question is: when data is exchanged inside a TPU Pod, how does the system reduce the latency of data movement?