Google Cloud TPU System Architecture: From First Principles to Pod-Scale Training

Prerequisites assumed: You know what a CPU, GPU, and TPU are at a conceptual level. You understand basic machine learning (forward pass, backward pass, gradients, optimizers) and basic distributed computing (nodes, networks, parallelism). You have never worked with Google Cloud TPUs or TPU Pods.

What this document is not: A restatement of the official TPU system architecture documentation. It’s a teaching document that builds the mental model the docs assume you already have, using analogies, diagrams, and worked examples — then connects that model back to the real terminology so the official docs become easy to read afterward.

A note on numbers: TPU generations (v2, v3, v4, v5e, v5p, v6e/Trillium, v7/Ironwood) change specs every year or two. This document teaches you the architecture and the mental model, and uses representative numbers from recent generations as illustrations. Always check Google’s current specs page for the exact numbers of the generation you’re using; the shapes and relationships explained here stay valid across generations.

1. The Big Picture

1.1 The problem Google is actually solving

Here’s a question that sounds simple but isn’t: how do you multiply matrices that are too big to fit on one chip, fast enough that a language model with hundreds of billions of parameters can be trained in weeks instead of decades?

A single accelerator chip hits a wall. Modern large models don’t just need fast compute, they need:

More memory than any single chip has (model weights, activations, optimizer states for hundreds of billions of parameters).
More raw compute than any single chip has (quadrillions of floating-point operations per training step).
A way for thousands of chips to act like one giant chip, exchanging partial results fast enough that communication doesn’t become the bottleneck.

GPUs solved this by clustering discrete cards with fast interconnects (NVLink) and networking (InfiniBand/Ethernet) between servers, orchestrated by a software stack (CUDA, NCCL) that grew up organically over 15 years.

Google’s TPU system took a different starting point. Instead of “take a graphics chip, adapt it for AI, then bolt on clustering,” Google asked: if you were designing a computer from scratch whose only job was to run neural network math at datacenter scale, what would you build?

The answer is TPU system architecture: purpose-built chips with a matrix multiplier as the literal centerpiece, connected by a purpose-built interconnect (ICI) that forms a physical mesh/torus network without needing a separate networking switch fabric for chip-to-chip traffic, all controlled by a compiler (XLA) that treats thousands of chips as a single compilation target rather than thousands of independent machines.

1.2 Why TPUs need a fundamentally different system architecture

A CPU-based system is built around flexibility. Every core is general-purpose, memory is designed for latency, and the network is a general datacenter network, because CPU workloads are enormously varied.

A GPU-based system is built around parallelism within a card, with high-speed links (NVLink) between cards within a server, then standard datacenter networking (InfiniBand/RoCE) to connect servers. This is hierarchical and somewhat heterogeneous: fast local links, slower global links, and the software (NCCL) has to know the difference.

A TPU-based system takes a different bet: build the interconnect into the chip itself as a first-class feature, so chips wire directly to their neighbors in a grid/cube pattern with no switch required for most links, and scale this pattern uniformly from 4 chips to thousands.

	CPU system	GPU system	TPU system
Design center	One flexible machine	One fast card, clustered	One coherent fabric of specialized chips
Scaling philosophy	Add more independent machines	Add more cards, then more servers	Add more chips to a uniform mesh/torus
What’s built in	General instruction set	Matrix cores + NVLink within a server	Matrix cores + interconnect wired chip-to-chip
Compiler’s view	Compiles for one core at a time	Compiles per-GPU; NCCL handles cross-GPU	Compiles for the whole slice as one program (SPMD)

That last row matters enormously and we’ll return to it in Section 6: on TPUs, the compiler (XLA) can treat hundreds or thousands of chips as if writing one program with data spread across them — a model called SPMD (Single Program, Multiple Data). This is only practical because the hardware topology is so uniform and predictable.

1.3 TPU Pods vs. a single TPU

A single TPU chip is already powerful, but it:

Has a fixed, finite amount of ultra-fast on-chip memory (HBM)
Has a fixed compute ceiling
Cannot hold a 70B, 175B, or trillion-parameter model’s weights, activations, and optimizer state by itself

A TPU Pod is Google’s answer: a large collection of TPU chips wired together via a dedicated high-speed interconnect (ICI, Inter-Chip Interconnect) into a physical mesh or torus topology, functioning as a single, enormous, distributed accelerator. A full pod might contain thousands of chips. You don’t have to rent a whole pod — you can rent a slice (a subset of a pod’s chips) sized to your job.

The key mental shift: with TPUs, you don’t think “I have one accelerator” — you think “I have a slice of a fabric,” and your program is compiled to run across that entire slice as a coordinated whole.

1.4 Where this fits in modern AI training

Training a large model requires, in order:

Loading and preprocessing enormous datasets (CPU-heavy, I/O-heavy work).
Running forward and backward passes with matrix multiplications at massive scale (the TPU’s core job).
Synchronizing gradients across every chip participating in training (the interconnect’s job).
Updating billions of parameters consistently across all chips (memory + compute + interconnect together).
Repeating this millions of times.

TPU system architecture is Google’s end-to-end answer to steps 2–4 at extreme scale: chip design (Section 3), topology (Section 4), the software stack that compiles your model onto that topology (Section 5), and the collective communication primitives that keep thousands of chips’ parameters consistent (Section 6). Everything in this document explains a different layer of that one stack.

2. TPU VM Architecture

2.1 What a TPU VM is

A TPU VM (also called a worker) is a Linux virtual machine that has direct access to one or more TPU chips physically attached to it. When you SSH into a TPU VM, you’re logging into a regular Linux box — you can install Python packages and run your training script directly on that machine, with the TPU chips visible as local devices.

The important detail: the TPU VM’s CPU, RAM, and the TPU chips it controls are physically co-located — the chips are connected to that host machine over PCIe, not accessed remotely over a generic network API.

2.2 Why TPU VMs exist — and what came before

The old model: TPU Node architecture. In the original design, you provisioned two separate things:

A regular Google Compute Engine (GCE) VM, which ran your Python code.
A separate, Google-managed TPU Node — a host machine + TPU chips you could not SSH into directly.

Your user-facing VM talked to the TPU Node over gRPC, across the network. This worked, but had real costs: an extra network hop for every step of data feeding, opacity (couldn’t profile or debug the host), and awkward access for low-level operations.

The new model: TPU VM architecture. Google removed the middleman. Now, the machine physically attached to the TPU chips is the machine you SSH into and run your code on.

Why this matters: imagine the difference between mailing your ingredients to a chef in another city and asking them to cook and mail the meal back, versus being the chef yourself, standing at your own stove. TPU Node was the former; TPU VM is the latter. You get lower latency for host-device communication, full visibility into the host machine, and the ability to write custom data-loading code that runs right next to the chips.

2.3 How the host CPU and TPU accelerator work together

Even in the TPU VM model, the CPU and TPU chips are two separate physical processors doing two very different jobs — they’re just co-located and tightly coupled.

The host CPU’s job:

Runs your Python interpreter, JAX/PyTorch/TensorFlow runtime, and training script’s control logic
Loads data from storage, decodes it, batches it, and stages it in host RAM
Invokes the XLA compiler to turn your model’s computation graph into an executable for the TPU
Hands the compiled program and input tensors to the TPU runtime, which triggers execution

The TPU chip’s job:

Executes the compiled program: matrix multiplications, activations, reductions, and collective operations over ICI with neighboring chips
Everything performance-critical happens on-chip, in HBM and the MXU, without going back to the host CPU on every operation

The CPU is the conductor, not a performer. For large training jobs, the CPU issues a large compiled program once per step, and the TPU chips run the whole step’s computation autonomously, exchanging data with each other directly over ICI — not routing through the host CPU.

2.4 Memory and interconnect map for a TPU VM

Three link types matter:

Host RAM ↔ Host CPU: ordinary DRAM access, used for data staging and running the Python/XLA control plane.
Host CPU ↔ TPU HBM (PCIe): relatively lower bandwidth and higher latency — used sparingly. Well-designed TPU programs minimize crossings of this link.
Chip ↔ Chip (ICI): the special sauce. High-bandwidth, low-latency, and does not go through the host CPU at all. Chip 1 and Chip 2 talk directly to each other.

Practical implication: if your training loop copies data back to the host every single step to compute a Python-side metric, you’re paying the slowest link on the hottest path. Idiomatic TPU code keeps the whole training step “on device” and only crosses the PCIe boundary at coarse intervals (e.g., every N steps for logging).

3. TPU Chip Architecture

3.1 The core idea: a chip built around one operation

Here’s the single most important fact about TPU chip design: almost the entire chip exists to serve one operation — matrix multiplication — as efficiently as physically possible. Everything else exists to feed that one unit with data fast enough that it’s never starved.

Compare this to a CPU (built around flexibility for arbitrary instructions) or a GPU (built around parallel flexible cores that happen to be good at matrix math among other things). A TPU inverts the priority: the matrix multiplier is the star of the show; everything else is supporting cast.

3.2 The Matrix Multiply Unit (MXU) and Systolic Arrays

The MXU is a grid of multiply-accumulate (MAC) units wired together as a systolic array. The intuition:

Consider matrix multiplication done the “normal” way: for every output element, you read a row from matrix A and a column from matrix B out of memory, multiply element-by-element, sum them up. If you do this naively, you’re constantly reading the same rows and columns from memory over and over — memory access becomes the bottleneck.

A systolic array solves this: arrange hundreds of small multiply-accumulate units in a grid, and stream the input data through the grid so that each value, once loaded, gets reused by many neighboring units before it’s discarded — like a bucket brigade where water (data) is passed hand-to-hand, and useful work (multiplication) happens at every hand-off, instead of everyone walking back to the well (memory) for every bucket.

Real MXUs are large grids — production TPU MXUs are square arrays on the order of hundreds by hundreds of MAC units. Inputs are typically stored and multiplied in bfloat16 (16-bit floating point), while running sums accumulate in full FP32 so numerical error doesn’t build up.

Why this matters for you as an engineer: matrix shapes close to the MXU’s native size (and its multiples) use the hardware efficiently. Oddly-shaped matrices — very “thin” matrix multiplications, or dimensions that don’t divide nicely — can leave part of the systolic array idle. This is why ML practitioners prefer batch sizes, hidden dimensions, and vocabulary sizes that are multiples of 128 or 256 when tuning models for TPUs.

3.3 The Vector Unit

Not everything in a neural network is a matrix multiplication. Activation functions (ReLU, GELU, softmax), elementwise operations (adding a bias, applying a mask, layer normalization), and various “glue” computations need the Vector Unit, which operates on whole vectors in parallel.

Think of the MXU as the heavy-machinery press that stamps out car body panels, and the Vector Unit as the finishing crew that sands edges and does quality checks on each panel right after it comes off the press.

3.4 The Scalar Unit

The Scalar Unit handles control flow (loops, conditionals within the compiled program), calculating memory addresses for where the next chunk of data should be read, and other orchestration work that doesn’t itself involve heavy numeric throughput.

If the MXU is the stamping press and the Vector Unit is the finishing crew, the Scalar Unit is the floor supervisor with the clipboard — deciding what happens next and where things go, without personally doing the heavy lifting.

Together, MXU + Vector Unit + Scalar Unit form what Google calls a TensorCore (not to be confused with Nvidia’s similarly-named “Tensor Cores,” which are different hardware). A TPU chip contains one or more TensorCores.

3.5 High Bandwidth Memory (HBM)

HBM is the chip’s main working memory — where model weights, activations, and gradients live while the chip is computing. The defining feature is bandwidth: it’s physically stacked very close to the compute logic so that enormous amounts of data can move between memory and the compute units every second.

This matters because the MXU is so fast at arithmetic that the real risk isn’t “not enough compute” — it’s starving the MXU of data. A systolic array ready to do 16,000+ multiply-accumulates per cycle is useless if memory can’t feed it operands fast enough. HBM exists specifically to solve this feeding problem.

3.6 Walking through a matrix multiplication

What happens when the TPU executes C = A @ B:

Data lives in HBM. Matrices A and B already sit in HBM.
Tiles are staged into on-chip SRAM. The Scalar Unit computes memory addresses for the specific tiles needed, and the chip streams those tiles from HBM into fast on-chip SRAM.
Data streams into the MXU. Rows of A and columns of B are pushed into the systolic array from two edges. Each cycle, values move one step through the grid, and every MAC unit performs a multiply-accumulate.
Partial sums accumulate in FP32. As data flows through the grid, partial products accumulate. Because inputs are bfloat16 but sums accumulate in FP32, precision is preserved.
Results flow out into SRAM and from there back to HBM.
The Vector Unit picks up the output for activation functions and bias additions.
The Scalar Unit loops for every tile until the full output matrix C is produced.
If this is a distributed matmul, ICI gets involved. If A or B is sharded across chips, partial results must be combined via collective operations over ICI.

4. TPU Topology

4.1 The vocabulary, built bottom-up

Chip: A single physical TPU integrated circuit with one or more TensorCores. The smallest unit of compute.
Device: In software (e.g., JAX), “device” usually refers to something addressable for computation — often mapping 1:1 with a chip or TensorCore, depending on generation and framework.
Host: A TPU VM — the CPU + RAM machine with a fixed set of TPU chips attached via PCIe (commonly 4 chips per host).
Slice: A collection of chips located inside the same TPU Pod, allocated together for your job. A slice can be single-host (all chips on one TPU VM) or multi-host (chips spanning multiple TPU VMs, all connected via ICI).
Pod: The maximum-size, fully-connected set of TPU chips of a given generation that Google’s datacenter interconnect supports as one fabric — the “whole fabric” a slice is carved from.

Analogy: think of a Pod as an entire apartment building wired with a shared high-speed internal network. A slice is “the 3rd floor, reserved for your team” — some floors are small (one apartment/host), others span many apartments (multi-host), but everyone on the building’s internal network can talk to each other fast.

4.2 Mesh and torus topology

Chips aren’t connected randomly or through a central switch — they’re wired directly to their physical neighbors in a regular grid pattern, like graph paper or a Rubik’s Cube.

A 2D mesh/torus connects chips in a flat grid — each chip connects to its neighbors up/down/left/right.
A 3D torus connects chips in a cube — each chip connects to neighbors along three axes (x, y, z).

The word “torus” specifically means the edges wrap around — the chip at the right edge of a row is also directly connected to the chip at the left edge of that same row (like Pac-Man’s screen). This wraparound halves the worst-case number of hops between any two chips and removes the “edge chips are worse off” asymmetry a plain grid would have.

For a 3D torus, imagine stacking several 2D torus layers and wiring corresponding chips between layers, with the front layer wrapping around to connect to the back layer. The result is a cube-shaped fabric where every chip has 6 direct neighbor connections, and the whole structure wraps in all three dimensions.

4.3 Why topology shapes are expressed as tuples

You’ll see TPU slice shapes written as tuples: a 2D slice might be 2x4 (2 chips by 4 chips = 8 chips total), a 3D slice might be 4x4x4 (4×4×4 = 64 chips — one “cube”). This notation directly describes the physical/logical grid dimensions of the torus the slice forms.

4.4 2D vs. 3D torus

2D torus: simpler physical wiring, cheaper, good enough when communication patterns don’t need the extra dimension.
3D torus: more direct neighbor connections per chip (6 instead of 4), which reduces the average number of hops to any other chip in a large slice. This matters more as slice sizes grow into the thousands of chips.

Analogy: a 2D torus is like a city laid out in a flat grid of streets that wrap at the edges. A 3D torus adds “elevators” connecting stacked floors of that same grid — a whole extra dimension of shortcuts, so the average trip across town gets shorter even as the city grows larger.

5. Data Flow: From Python to Silicon

5.1 The compilation stack

The path from your Python code to TPU execution:

Your Python code (model(x))
    → ML Framework (JAX / PyTorch / TensorFlow)
    → Computation Graph
    → XLA Compiler
    → HLO (High-Level Optimizer IR)
    → Compiler Optimization Passes (fusion, layout, sharding, scheduling)
    → Compiled Executable (TPU machine code)
    → TPU Runtime (on the Host CPU)
    → TPU Chip(s) (MXU / Vector / Scalar / HBM)
    → Result tensors

Key terms:

XLA (Accelerated Linear Algebra): Google’s domain-specific compiler for linear algebra. Takes a description of your computation and produces highly optimized machine code for a specific target.
HLO (High-Level Optimizer IR): the intermediate representation XLA uses internally — a structured, optimizer-friendly “assembly language” for tensor math. When engineers say “let’s look at the HLO,” they mean inspecting this intermediate form to understand what the compiler decided to do.
Compiled executable: the final, chip-specific machine code XLA produces after optimization passes.
TPU runtime: software on the host CPU that manages the compiled executable’s lifecycle — transferring it and input data to the chip, triggering execution, retrieving outputs.

5.2 Why compilation is central to TPU performance

If TPUs ran Python operations one at a time, every single tensor operation would need a round-trip: Python → host CPU → PCIe → chip → PCIe → host CPU → Python. Given that PCIe is the slowest link in the system, this would be catastrophically slow.

Instead, XLA compiles entire functions — often an entire training step (forward pass, loss, backward pass, and optimizer update, all together) — into a single executable ahead of time. The host CPU dispatches that whole compiled program to the chip once, and the chip runs potentially thousands of individual operations autonomously, only reporting back when the whole step is done.

Analogy: this is the difference between giving a chef one instruction at a time over a walkie-talkie versus handing the chef a complete recipe card once and letting them cook the entire dish before reporting back. XLA writes the recipe card from your Python code.

5.3 Caching and shape sensitivity

On the second call to model(x) with the same input shapes, almost everything above the “Dispatch executable” step is skipped — no re-tracing, no re-compiling. A new input shape means a new compilation, which is comparatively slow.

This is why you’ll see advice like “pad your batches to a fixed size” or “avoid Python-level control flow that depends on tensor values” when writing TPU-targeted code — both keep the shape and structure of the computation graph stable across calls, so the expensive compilation step only happens once.

5.4 What XLA’s optimization passes are actually doing

Operation fusion: instead of writing the matmul result to HBM, reading it back for bias-add, writing again, reading again for ReLU, XLA fuses these into a single compiled kernel that keeps intermediate results in fast on-chip SRAM and only writes the final result to HBM.
Layout assignment: deciding exactly how a tensor’s values are arranged in HBM so that when the MXU streams the tensor in, it arrives in the most efficient order.
Sharding decisions: for multi-chip programs, deciding which tensors get split across which chips and where ICI communication needs to be inserted.
Scheduling: ordering independent operations to maximize overlap — e.g., starting a communication operation on ICI while the MXU is still busy with unrelated compute, so communication latency gets “hidden” behind computation.

6. Distributed Training

How do thousands of chips cooperate to train one model? There are a handful of orthogonal strategies, and real large-scale training almost always combines several at once.

6.1 The core problem

A model has three things that can be “big”: the data (millions of training examples), the parameters (the weights themselves — potentially hundreds of billions), and the activations (intermediate values during the forward pass, which consume memory). Different parallelism strategies address different combinations.

6.2 Data Parallelism

The idea: every chip has a full copy of the model, but each chip processes a different slice of the batch. After each chip computes gradients on its own mini-batch, all chips’ gradients are averaged together (via AllReduce) so every replica ends up with identical combined gradients — keeping every copy of the model in sync.

When it’s enough: when the model fits comfortably in a single chip’s HBM, and you just want to process more data per second.

Its limit: does nothing to help when the model itself is too big to fit on one chip — every chip still needs a full copy.

6.3 Model Parallelism

When the model doesn’t fit on one chip, split the model itself across chips.

Tensor Parallelism

Split individual large operations — most often, individual weight matrices — across chips. For example, a huge weight matrix can be split column-wise across 4 chips; each chip computes a partial result using its slice of the weights, and the partial results are combined via ICI to form the full output. This happens within a single layer’s computation — chips are cooperating on one operation together, requiring frequent, low-latency communication (exactly what ICI is built for).

Pipeline Parallelism

Split the model by layer — chips 1–2 hold the first half of the network’s layers, chips 3–4 hold the second half. Data flows through the pipeline like an assembly line. To keep every chip busy, multiple micro-batches are kept “in flight” through the pipeline simultaneously.

Tensor vs. pipeline, intuitively: tensor parallelism is like several people jointly lifting one very heavy box together, in constant close coordination. Pipeline parallelism is like an assembly line, where each station does its own job on a different box passing through — less constant coordination, but you need to keep boxes flowing continuously.

6.4 SPMD: Single Program, Multiple Data

Rather than writing separate code for “what chip 1 does” vs. “what chip 2 does,” SPMD means you write one program describing the computation on the full, logical (unsharded) tensors, and the compiler (XLA) automatically figures out how to partition that computation across all the chips, inserting the necessary communication.

You annotate how you want your tensors split (e.g., “split the batch dimension across these 8 chips” or “split this weight matrix’s output dimension across these 4 chips”), and XLA’s partitioner does the rest — generating per-chip code and the ICI collective operations needed to keep results correct, automatically.

GSPMD (Generalized SPMD) is Google’s more advanced system: it takes user-provided sharding hints on a handful of tensors and propagates sharding decisions through the entire computation graph automatically, handling combinations of data, tensor, and pipeline parallelism within one unified framework.

Why this is a big deal: the same Python model code can, in principle, scale from 1 chip to thousands of chips by changing sharding annotations and slice size — not by rewriting the model’s logic.

6.5 Collective operations: the vocabulary of chip-to-chip communication

Broadcast: one chip sends the same data to all other chips. Example: distributing initial model weights to every chip.
AllReduce: every chip contributes a value (e.g., its locally computed gradient), all values are combined (summed or averaged), and every chip ends up with the same combined result. The workhorse of data-parallel gradient synchronization.
AllGather: every chip contributes a piece of data, and every chip ends up with the full collection of all pieces (concatenated). Used when every chip needs to see the complete tensor that was previously sharded.
ReduceScatter: every chip contributes a full-size value; these are combined (reduced), but instead of every chip getting the whole combined result, each chip ends up with only its own slice of the combined result. Often implemented internally as AllReduce = ReduceScatter followed by AllGather.

6.6 How TPU Pods cooperate during LLM training

A real large-model training step on a TPU Pod typically layers several ideas at once:

The full chip slice is organized into a logical mesh — e.g., an 8×32 arrangement where one axis is used for data parallelism and the other for tensor parallelism.
Each replica (a group of chips along the tensor-parallel axis) processes a different shard of the global batch.
Within each replica, large weight matrices and attention computations are split across chips — requiring frequent AllReduce/AllGather operations over ICI within that replica group.
After the backward pass, gradients are synchronized across replicas via AllReduce — crossing more of the mesh, but happening less frequently per step.
GSPMD’s compiler has already worked out exactly which collective operation goes where, based on the sharding annotations you provided.

7. Memory Architecture

7.1 The memory hierarchy

Host RAM (DRAM)         → Largest capacity, lowest bandwidth. Accessed over PCIe.
HBM (High Bandwidth Memory) → Large capacity, very high bandwidth. The TPU's main working memory.
On-chip SRAM / Registers → Small capacity, extremely high bandwidth. Right next to the MXU.
                         ↓ feeds directly ↓
                         MXU

This is the same fundamental idea as a CPU’s memory hierarchy (registers → L1/L2/L3 cache → RAM → disk), just with different names. The further data has to travel from the compute unit, the larger the tier tends to be, but the slower it is to access. Good performance engineering on a TPU is largely about keeping data in the fastest tier possible for as long as possible.

7.2 DMA transfers

DMA (Direct Memory Access) is a mechanism that lets data move between memory tiers without requiring the compute unit to babysit every byte. A DMA engine is told “move this block from address A to address B,” and it does so independently, freeing the CPU or TensorCore to keep working on something else. This is why data loading and computation can overlap: while the MXU is busy on the current batch, a DMA transfer can already be moving the next batch from host RAM into HBM in the background.

7.3 Memory bandwidth vs. FLOPs — why bandwidth is often the real bottleneck

FLOPs measure how much arithmetic a chip can do. Memory bandwidth measures how fast data can move between memory and the compute units. Marketing materials love to quote FLOPs because it’s a big, exciting number — but for many real workloads, FLOPs is not the bottleneck at all. Bandwidth is.

Consider adding two large vectors elementwise (c = a + b). For every single addition (1 FLOP), you need to read two values and write one — that’s 3 memory accesses for 1 unit of compute. No matter how fast your arithmetic units are, if memory can’t deliver operands fast enough, the arithmetic units sit idle. This ratio — how much data movement an operation requires per unit of compute — is called arithmetic intensity.

Matrix multiplication has naturally high arithmetic intensity when matrices are reasonably large. This is exactly why the systolic array design works so well: each value loaded gets reused across many multiply-accumulate operations before being discarded.

Practical consequence: techniques like operation fusion, larger batch sizes, and avoiding unnecessary intermediate tensors all exist to increase effective arithmetic intensity — keeping the MXU fed from fast on-chip memory rather than constantly waiting on HBM, and keeping HBM fed rather than waiting on the (much slower) PCIe link. When people say a model is “memory-bound,” this is what they mean.

8. Networking

8.1 ICI in detail

ICI (Inter-Chip Interconnect) is the dedicated, purpose-built network fabric that directly wires TPU chips to their neighbors in the mesh/torus topology. Key properties:

Direct chip-to-chip links. Chips connect to their physical neighbors without going through a general-purpose network switch for most communication.
Very low latency. Because the topology and routing are fixed and predictable, communication latency is far more predictable and typically far lower than general datacenter networking.
High, dedicated bandwidth. ICI links are engineered specifically for the traffic pattern of collective operations (AllReduce, AllGather) that distributed training generates constantly.
Scales with slice shape. Because it’s a mesh/torus, aggregate bandwidth across the fabric scales as you add chips, rather than funneling everything through one shared central switch.

8.2 Latency and bandwidth for different communication patterns

Communication-heavy strategies with lots of small, frequent exchanges (like tightly-coupled tensor parallelism) are more sensitive to latency.
Communication involving large tensors (like AllReducing gradients for billions of parameters) is more sensitive to bandwidth.

ICI is engineered to be strong on both axes because real large-model training needs both kinds of communication simultaneously.

8.3 Why networking matters so much for large models

As models grow, the fraction of total training time spent on communication tends to grow too — because parameter counts grow (larger gradients to synchronize) and splitting a model across more chips means more frequent cross-chip communication within each layer’s computation. If the network can’t keep up, adding more chips stops helping — you hit diminishing returns. This is why interconnect design isn’t a secondary detail; for large-scale training, it’s as architecturally central as the compute units themselves.

8.4 TPU ICI vs. GPU clusters (NVLink + InfiniBand)

GPU clusters solve the same problem with a two-tier approach:

NVLink: high-bandwidth, low-latency interconnect connecting GPUs within the same server (typically up to 8 GPUs).
InfiniBand (or high-speed RoCE Ethernet): connects servers to each other, routed through switches, for communication that needs to cross beyond a single server’s GPUs.

This creates a non-uniform communication picture: talking to a GPU in the same server (via NVLink) is much faster than talking to a GPU in a different server (via InfiniBand, through switches). Software (NCCL) has to be aware of this hierarchy.

TPU’s ICI mesh/torus, by contrast, aims for a more uniform fabric: whether two chips are on the same host or different hosts, they’re still neighbors (or a few hops away) in the same logical mesh/torus, connected by the same kind of link.

Neither approach is strictly “better” — GPU clusters benefit from an enormous, mature, general-purpose software ecosystem and flexible commodity networking; TPU Pods benefit from a more uniform, purpose-built fabric and a compiler stack (XLA/GSPMD) designed hand-in-hand with that fabric from the start.

9. Real Training Example: A Transformer, Start to Finish

Let’s walk through one complete training step for a Transformer model on a TPU slice — 64 chips across several hosts, using data parallelism across replicas and tensor parallelism within each replica.

Batch loading (Host CPU, every host in parallel). Each host’s CPU reads its shard of the training data from storage, tokenizes text, pads/truncates to a fixed sequence length (fixed shapes matter — Section 5.3), and assembles a batch in host RAM. Zero TPU activity here.
Sharding (decided at compile time). Based on your mesh and sharding specification, the framework determines which slice of the global batch goes to which chip. This was worked out once at compilation time.
Host-to-device transfer (PCIe DMA). Each host DMA-transfers its batch shard from host RAM into the HBM of its locally attached chips. The “slow link” from Section 2.4 — crossed once per step for input data.
Forward pass — embedding lookup (TPU chip, HBM + Vector/Scalar Units). Token IDs are used to look up embedding vectors in HBM. Relatively memory-heavy operation (low arithmetic intensity).
Forward pass — attention and feed-forward layers (MXU-dominated, with ICI for tensor-parallel layers). For every Transformer block: query/key/value projections, attention computation, and feed-forward matrix multiplications run as large matmuls on the MXU. Because weight matrices are tensor-parallel-sharded, each chip computes a partial result using its own weight shard, and an AllGather or AllReduce over ICI combines partial results — this repeats for every Transformer block, so ICI traffic is interleaved with MXU computation throughout the entire forward pass.
Loss computation (TPU chip, Vector/Scalar Units). The final layer’s logits are compared against target labels via cross-entropy to produce a scalar loss — computed on-chip, without returning to the host.
Backward pass (MXU + Vector/Scalar Units + ICI). Gradients are computed via more matrix multiplications and elementwise operations, mirroring the forward pass in reverse. Wherever the forward pass needed a tensor-parallel collective, the backward pass needs a corresponding one.
Gradient synchronization across data-parallel replicas (ICI, AllReduce). Once each replica has computed its local gradients, an AllReduce combines them across all data-parallel replicas so every replica ends up with the same, globally-averaged gradient.
Optimizer step (entirely local after synchronization). Each chip updates its own shard of the parameters (e.g., Adam’s moment updates and parameter adjustment). No further communication needed — every chip already has the correctly-synchronized gradient for the parameters it owns.
Loop. Updated parameters remain resident in HBM, used directly as input to the next training step — no round trip to the host in steady state.
Occasional host round-trips (logging/checkpointing). Periodically — not every step — scalar metrics or full parameter checkpoints are copied back from HBM to host RAM, crossing the slow PCIe link deliberately but infrequently.

The big takeaway: in steady-state training, the host CPU’s role shrinks to “feed data in, occasionally pull metrics out.” The overwhelming majority of both computation and communication happens entirely device-side — chip HBM, chip compute units, and chip-to-chip ICI links.

10. Comparisons

CPU vs. GPU vs. TPU

	CPU	GPU	TPU
Design center	General-purpose, flexible	Massively parallel, general-purpose-ish	Purpose-built for tensor/matrix math
Core count/style	Few, complex cores	Thousands of simpler cores	Few TensorCores, each built around a large MXU systolic array
Best at	Sequential logic, varied workloads	Parallel workloads: graphics, simulation, ML	Large-scale matrix multiplication for ML training/inference
Interconnect	Standard datacenter networking	NVLink (intra-server) + InfiniBand (inter-server)	ICI mesh/torus, uniform across intra- and inter-host
Programming model	Any general-purpose language	CUDA / ROCm + frameworks	XLA-compiled JAX/PyTorch/TensorFlow
Typical ML role	Orchestration, data loading, serving glue	Training and inference (extremely widely used)	Large-scale training and inference at Google/GCP scale

GPU Cluster vs. TPU Pod

	GPU Cluster	TPU Pod
Topology	Hierarchical: NVLink islands joined by InfiniBand/Ethernet	Uniform mesh/torus fabric (ICI) across the whole pod
Scaling pattern	Add servers connected via switched fabric	Add chips into the same physical mesh/torus
Software for cross-chip comms	NCCL, topology-aware	XLA/GSPMD, compiler-driven partitioning
Ecosystem maturity	Very broad — most ML software targets CUDA first	Narrower but deep — strong JAX/TensorFlow support, growing PyTorch/XLA support
Flexibility	Very high — GPUs run almost any GPU-compatible workload	Narrower — optimized specifically for ML tensor workloads

CUDA vs. XLA

	CUDA	XLA
What it is	A parallel computing platform + programming model for Nvidia GPUs	A domain-specific compiler for linear algebra, targeting TPU (and also CPU/GPU)
Programming style	Often lower-level; write custom kernels directly, or rely on cuDNN/cuBLAS libraries	Higher-level; describe tensor computations in JAX/PyTorch/TensorFlow, XLA compiles automatically
Where optimization happens	Partly hand-written kernels/libraries, partly compiler	Almost entirely automatic, compiler-driven (fusion, layout, sharding, scheduling)
Analogy	A general workshop full of power tools you can wield directly	A master craftsman who takes your blueprint and decides how to build it optimally

11. Glossary

TPU (Tensor Processing Unit): Google’s custom-built chip, purpose-designed to accelerate the matrix/tensor math used in machine learning.
TPU VM: A Linux virtual machine with direct, local access (via PCIe) to physically-attached TPU chips; you SSH directly into it.
TPU Node: The legacy Cloud TPU architecture where the host driving the chips was a separate, Google-managed machine you could not SSH into; superseded by TPU VM.
Pod: The largest fully-interconnected fabric of TPU chips of a given generation, forming one coherent mesh/torus.
Slice: A subset of a Pod’s chips, allocated together for a given job; can be single-host or multi-host.
HBM (High Bandwidth Memory): The TPU chip’s large, very-high-bandwidth main working memory, physically close to the compute units.
XLA (Accelerated Linear Algebra): Google’s compiler that turns tensor computation graphs into optimized machine code for TPU (and CPU/GPU).
HLO (High-Level Optimizer IR): XLA’s internal intermediate representation, between your framework’s computation graph and the final compiled machine code.
SPMD (Single Program, Multiple Data): A programming model where one program is written against full logical tensors, and the compiler automatically partitions it across many chips.
GSPMD (Generalized SPMD): Google’s system for automatically propagating sharding decisions across an entire computation graph from a small number of user-provided annotations.
MXU (Matrix Multiply Unit): The systolic-array hardware block at the heart of a TensorCore, performing large numbers of multiply-accumulate operations per cycle.
TensorCore: The combination of MXU + Vector Unit + Scalar Unit that makes up the compute portion of a TPU chip.
ICI (Inter-Chip Interconnect): The dedicated high-bandwidth, low-latency network fabric wiring TPU chips directly to their neighbors in the mesh/torus topology.
Systolic array: A grid of tightly-connected compute units that stream data through themselves, reusing loaded values across many computations — the design underlying the MXU.
bfloat16: A 16-bit floating-point format commonly used for TPU matrix multiplication inputs, trading precision for speed and memory savings, typically paired with FP32 accumulation.
DMA (Direct Memory Access): A mechanism for moving data between memory tiers without requiring the CPU or TensorCore to manage each byte individually, enabling overlap between data transfer and computation.
Arithmetic intensity: The ratio of compute operations to memory accesses for a given operation; low arithmetic intensity operations tend to be memory-bandwidth-bound rather than compute-bound.
AllReduce: A collective operation where every chip contributes a value, all values are combined (summed or averaged), and every chip ends up with the same combined result.
AllGather: A collective operation where every chip contributes a piece of data, and every chip ends up with the full collection of all pieces concatenated.
Sharding: Splitting a tensor (data, weights, or activations) into pieces distributed across multiple chips.
Mesh: The logical grid arrangement of chips used for reasoning about sharding and communication.

12. The Complete Journey: From Python to Hardware

Three lines of code:

loss = model(images)
loss.backward()
optimizer.step()

Let’s trace this end to end.

The moment before: compilation

The very first time this training step function executes, the framework doesn’t run it eagerly. It traces the sequence of operations and hands the resulting computation graph to XLA. XLA lowers this into HLO, and — because you’ve specified a mesh and sharding annotations — the GSPMD partitioner rewrites that HLO into a version that’s aware of exactly how every tensor is split across your slice’s chips, inserting collective operations at precisely the points where cross-chip communication is mathematically required. XLA’s optimization passes then fuse operations, assign memory layouts, and schedule everything (including overlapping ICI communication with MXU computation), producing one compiled executable — the entire training step as a single program the TPU can run start to finish without checking back in with Python.

`loss = model(images)` — the forward pass

Host CPU (on every host in the slice, in parallel) has already loaded and batched images; each host holds its shard in host RAM.
Each host DMA-transfers its batch shard across PCIe into the HBM of its locally-attached chips.
On each chip, the TPU runtime triggers execution of the compiled executable. Computation proceeds on-device, autonomously.
For each Transformer layer: the MXU performs large matrix multiplications, pulling tiles from HBM through fast on-chip SRAM. The Vector Unit applies activation functions and layer normalization. The Scalar Unit manages control flow and addressing.
Because weight matrices are tensor-parallel sharded, partial results are combined with peer chips via AllGather/AllReduce over ICI — direct chip-to-chip links, never touching the host CPU.

`loss.backward()` — the backward pass

The compiled executable’s backward portion computes gradients in reverse, walking the network in reverse order.
This mirrors the forward pass: more large matmuls on the MXU, with corresponding ICI collective operations inserted automatically by GSPMD wherever the forward pass required communication.
At the end, each chip holds a local gradient for the parameter shards it owns, computed from its shard of the batch.

`optimizer.step()` — synchronization and update

Before parameters can be updated correctly, gradients computed independently on each data-parallel replica must be combined — an AllReduce over ICI, crossing the wider part of the mesh.
Once synchronized, each chip runs the optimizer’s math locally, updating only the parameter shard it owns. No further communication needed.
The updated parameters remain resident in HBM, ready for the next call to model(images) — no round-trip to the host CPU required.

Why this architecture wins at scale

Every design decision in this document — the systolic array, the torus topology, the ICI fabric, the XLA compiler, GSPMD, the shift from TPU Node to TPU VM — serves a single goal: minimize the time the high-performance compute units spend waiting on anything else. Waiting on memory (solved by HBM and operation fusion). Waiting on communication (solved by ICI and compiler-driven overlap). Waiting on the host CPU (solved by compilation and direct device-side execution).

The result is a system where, in steady-state training, the overhead connecting your Python code to the silicon running your model is, by design, almost invisible.