TPU Worker

A TPU worker is a virtual machine (VM) running Linux that has direct access to one or more Tensor Processing Unit (TPU) chips and executes the actual TPU computation on that attached hardware.^[1][15] In Google Cloud's TPU architecture, each worker serves as the computational host responsible for executing machine learning workloads on the TPU chips it is wired to.^[1] Google Cloud documentation defines it plainly: "A TPU VM, also known as a worker, is a virtual machine running Linux that has access to the underlying TPUs."^[15] Workers load and preprocess data, submit compiled programs to TPU chips via the XLA compiler, and coordinate with other workers during distributed training.^[1] The term "TPU worker" is used interchangeably with "TPU VM" in Google Cloud documentation, and it represents the fundamental unit of execution in both single-host and multi-host TPU configurations.^[2][15]

What is a TPU worker?

A TPU worker is the host process and its CPU-based virtual machine that drives the math on the attached TPU chips. It is the counterpart to the "master" or coordinator that, in the original Cloud TPU programming model, hands a computation graph to the worker. Where the coordinator decides what to run, the worker is where the TPU computation actually happens: it feeds inputs into the chips, triggers execution of the compiled program, and collects the results.^[1][16] In multi-host TPU pods, every host runs its own TPU worker, and the workers are coordinated so that one large model or batch is processed across many chips at once.^[2][3] Because the worker is a full Linux VM, in the current TPU VM architecture users SSH directly into it and run their training code on the same machine that holds the TPUs.^[15]

Explain like I'm 5 (ELI5)

Imagine you have a special calculator that is really fast at doing math homework (that is the TPU chip). But the calculator cannot read the homework by itself. It needs a helper to read the problems, write them down for the calculator, and then collect the answers. That helper is the TPU worker. Sometimes the homework is so big that you need many helpers, each with their own calculator, all working together on different parts of the homework at the same time. The helpers pass notes to each other so they all stay on the same page.

What does a TPU worker do?

A TPU worker is a CPU-based virtual machine physically connected to TPU hardware via PCIe.^[1] The worker handles tasks that the TPU chips themselves cannot perform: reading data from storage, running preprocessing pipelines, managing control flow, and communicating with other workers over the data center network. The actual tensor computations (matrix multiplications, convolutions, and other deep learning operations) are offloaded to the TPU chips.^[2]

The relationship between a TPU worker and its TPU chips follows a host-device model. The worker (host) prepares computation graphs using a framework such as JAX, PyTorch, or TensorFlow. These graphs are compiled by the XLA (Accelerated Linear Algebra) compiler into optimized TPU machine code. The compiled program is then dispatched to the TPU chips (devices) for execution. Results flow back to the worker through an outfeed queue, while input data enters through an infeed queue.^[1]

Each TPU worker can access up to 8 TPU chips, depending on the machine type and TPU generation.^[1] For example, a TPU v5e worker using the ct5lp-hightpu-8t machine type has access to 8 chips, while a ct5lp-hightpu-1t configuration provides only 1 chip per worker.^[6]

What is the difference between a TPU worker and a TPU master?

In the master/worker model that comes from the TensorFlow 1 era of Cloud TPU, the TPU master is the coordinator process that receives a computation graph from the client, partitions and optimizes it, and dispatches the work; the TPU worker is the process that runs the math on the TPU chips.^[16] The user's host connects to a remote TPU worker over a gRPC address, conventionally on port 8470 (for example, grpc://10.1.2.3:8470), which is how the coordinator feeds and steers the accelerator.^[16] In short: the master decides and directs, the worker computes.

Scroll sideways for more →

The explicit "master" terminology has evolved. It originated in the TensorFlow 1 distributed runtime, and in modern JAX and TensorFlow 2 workflows the explicit master string has been largely superseded by cluster resolvers and runtime services such as the TPUClusterResolver, though the underlying split between a coordinating host and the TPU that executes the computation remains.^[16] The shift to the TPU VM architecture also collapsed the user-facing distinction: instead of a separate user VM talking to a remote worker over gRPC, the user now logs directly into the worker VM itself.^[2][15] For the control role and its on-chip and host-level senses, see TPU master.

TPU node vs. TPU VM architecture

Google Cloud has offered two distinct architectures for accessing TPU hardware. Understanding the difference is important because the term "TPU worker" carries slightly different meaning in each.

Scroll sideways for more →

In the older TPU node architecture, the "worker" referred to the remote TPU host that the user could not directly access. In the current TPU VM architecture, the worker is the VM that the user directly logs into and runs code on.^[2] The TPU VM approach eliminates the overhead of a separate user VM and enables training setups (such as distributed reinforcement learning) that were not feasible with the TPU node model.^[11]

Hardware components of a TPU worker

Each TPU worker is connected to TPU chips that contain specialized processing elements.

TensorCore

A TensorCore is the primary compute unit within a TPU chip. Each TensorCore contains:

• Matrix multiply units (MXUs): Systolic arrays that perform the bulk of tensor computation.^[10] In TPU generations prior to v6e, each MXU uses a 128x128 array of multiply-accumulators. Starting with TPU v6e (Trillium), the MXU was expanded to 256x256, quadrupling FLOPS per cycle. Each MXU performs 16,000 multiply-accumulate operations per cycle, taking bfloat16 inputs and accumulating results in FP32.^[10]
• Vector processing unit (VPU): Handles element-wise operations such as activations, softmax, and other non-matrix computations.
• Scalar unit: Manages control flow and memory addressing.

The number of TensorCores per chip varies by generation. TPU v4 has two TensorCores per chip (each with four MXUs), while TPU v5e has one TensorCore per chip (with four MXUs).^[5]

Memory system

TPU chips use high-bandwidth memory (HBM) for storing model parameters and activations. Capacity and bandwidth vary by generation:

Scroll sideways for more →

In addition to HBM, each TensorCore has on-chip vector memory (VMEM) that serves as a software-controlled scratchpad. VMEM bandwidth is roughly 22 times higher than HBM bandwidth, making it valuable for operations on smaller tensors that fit in local storage.^[10]

SparseCore

Starting with TPU v4, Google introduced SparseCores: specialized dataflow processors designed to accelerate embedding operations common in recommendation and ranking models.^[8] TPU v4 includes four SparseCores per chip, each with 2.5 MB of scratchpad memory. SparseCores accelerate embedding-heavy models by 5x to 7x while using only about 5% of die area and power.^[8] TPU v5p and Ironwood (v7) also include SparseCores.^[14]

Communication hierarchy

TPU workers participate in a layered communication system. The bandwidth at each level determines how training workloads should be partitioned across devices.

Scroll sideways for more →

The steep drop in bandwidth from ICI to DCN (roughly 10x or more) has major implications for how parallelism strategies are chosen.^[10] Operations that require frequent, low-latency communication (such as tensor parallelism) work well over ICI within a single slice, while strategies that tolerate higher latency and lower bandwidth (such as data parallelism) are better suited for DCN communication across slices.^[10]

Single-host and multi-host configurations

TPU workers are deployed in one of two modes, depending on the scale of the workload.

Single-host (single worker)

A single-host configuration uses one TPU VM with its attached chips. This is appropriate for smaller models that fit within the memory and compute capacity of a single worker. For TPU v5e, a single host can access up to 8 chips, providing up to 1,576 TFLOPS of bf16 compute and 128 GB of combined HBM.^[6]

Multi-host (multiple workers)

A multi-host configuration distributes training across multiple TPU VMs. Each worker runs the same training program (following the SPMD, or Single Program Multiple Data, paradigm), but operates on different portions of data or different shards of the model.^[3] Workers within a multi-host slice are connected by ICI, enabling high-bandwidth collective operations such as all-reduce.^[3]

Multi-host configurations are required when the model or batch size exceeds what a single worker can handle. For example, a TPU v4 slice with a 4x4x4 topology consists of 64 chips spread across 16 workers (4 chips per worker in TPU v4).^[5]

Slices, pods, and Multislice

TPU workers are organized into increasingly large groupings.

Slice

A slice is a collection of TPU chips within the same pod, connected by high-speed ICI links. All workers in a slice can communicate directly through ICI without going through the data center network.^[3] Slice topology is specified as a tuple describing the chip layout:

• 2D topologies (TPU v5e, v6e): Specified as AxB (e.g., 4x4 for 16 chips). Chips are connected in a 2D torus.
• 3D topologies (TPU v4, v5p, Ironwood): Specified as AxBxC (e.g., 4x4x4 for 64 chips). Chips are connected in a 3D torus, where each chip links to its six nearest neighbors.

Some topologies support a twisted torus configuration that increases bisection bandwidth. For example, a 4x4x8_twisted topology provides approximately 70% higher bisection bandwidth compared to a non-twisted 4x4x8 layout.^[10]

Pod

A TPU pod is the largest contiguous grouping of TPU chips connected by ICI. Pod sizes vary by generation:

Scroll sideways for more →

Within a pod, reconfigurable optical circuit switches (OCS) can dynamically rearrange inter-cube connections, improving fault tolerance and scheduling flexibility. This feature was introduced with TPU v4.^[8]

Multislice

Multislice extends TPU capacity beyond a single pod by connecting multiple slices over the data center network (DCN). Within each slice, chips continue to communicate over ICI. Between slices, data is transferred from TPU chips to the CPU host over PCIe, then across the DCN to other hosts.^[4]

Developers do not need to write explicit inter-slice communication code. The XLA compiler detects the hybrid ICI/DCN topology and automatically generates hierarchical collective operations, overlapping DCN communication with computation to hide latency.^[4]

Multislice has been used to run training jobs on over 50,000 TPU v5e chips simultaneously, representing the largest distributed LLM training job on TPUs as of the announcement.^[12]

Arithmetic intensity requirements

For Multislice training to be efficient, the ratio of computation to communication must be high enough to keep TPU chips busy while gradients are synchronized over DCN.^[4] For TPU v4 chips (275 TFLOPS each) with a per-host DCN bandwidth of 50 Gbps, the required arithmetic intensity is approximately 22,000 FLOPS per bit. In practice, this means transformer models trained across two slices need a minimum batch size of roughly 350,000 tokens, with 700,000 or more tokens recommended when using many slices.^[4]

Parallelism strategies across workers

Distributed training across TPU workers uses several parallelism approaches, often combined.

Data parallelism

Each worker holds a complete copy of the model and processes a different subset of the training batch. After computing gradients locally, workers synchronize through an all-reduce operation. Data parallelism is the simplest strategy and works well when the model fits in a single worker's memory.^[10]

Fully sharded data parallelism (FSDP)

FSDP partitions model parameters, gradients, and optimizer states across workers. Each worker stores only a shard of the full model state, reducing per-worker memory requirements. Parameters are gathered on demand for forward and backward passes, then re-sharded afterward. FSDP is commonly used within a slice over ICI.

Tensor parallelism

Large matrix operations are split across multiple chips, with each chip computing a portion of the result. This requires frequent inter-chip communication and works best over ICI within a slice. Tensor parallelism is not recommended over DCN due to the high communication overhead.^[10]

Pipeline parallelism

Different layers of the model are assigned to different workers. Data flows through the pipeline in micro-batches. Pipeline parallelism reduces per-worker memory requirements but introduces pipeline bubbles (idle time between micro-batches).

Combined strategies

Large-scale training typically combines multiple parallelism approaches. A common pattern for LLM training on TPUs uses data parallelism across slices over DCN, where each slice stores a model replica. Within each slice, the model is sharded across chips using FSDP or a combination of FSDP and tensor parallelism over ICI.^[10]

The ICI and DCN parallelism dimensions are configured independently. Within a slice, the product of ici_data_parallelism, ici_fsdp_parallelism, and ici_tensor_parallelism must equal the number of chips per slice. Across slices, the corresponding DCN values must multiply to equal the number of slices.^[3] Google's documentation recommends always setting dcn_tensor_parallelism to 1, since DCN bandwidth is too low for the frequent communication that tensor parallelism requires.^[3]

XLA compilation and execution on workers

The XLA (Accelerated Linear Algebra) compiler is central to how TPU workers execute computation.

1. Graph tracing: The ML framework (JAX, PyTorch/XLA, or TensorFlow) traces the computation into an intermediate representation (IR) graph.
2. HLO generation: The IR graph is converted into HLO (High Level Operations) format.
3. Optimization: XLA performs operator fusion, memory layout optimization, and parallelization passes on the HLO graph.
4. LLO compilation: The optimized HLO is further compiled into LLO (Low Level Optimized) format, producing TPU-specific machine instructions.
5. Execution: The compiled program runs on the TPU chips. Compiled graphs are cached so that subsequent executions with the same computation graph and input shapes can reuse the binary.^[1]

For distributed execution, XLA uses the GSPMD (General-purpose SPMD) model. Developers annotate how tensors should be sharded across devices using high-level APIs such as JAX's jax.sharding or shard_map. The XLA compiler then automatically inserts the necessary collective communication operations (all-reduce, all-gather, reduce-scatter) and maps the single program onto all TPU chips in the configuration.^[10]

TPU workers in Kubernetes (GKE)

Google Kubernetes Engine (GKE) provides managed orchestration of TPU workers through TPU slice node pools.

Single-host node pools

Each node in a single-host TPU node pool is an independent TPU VM. The TPUs attached to different VMs are not interconnected via ICI. Nodes can be added or removed individually, and standard Kubernetes scaling behavior applies.^[7]

Multi-host node pools

Multi-host TPU slice node pools contain two or more interconnected TPU VMs that form a single slice. GKE treats these as atomic units:

• If one node in the slice fails to deploy, the entire slice deployment fails.
• During scaling operations, GKE scales the entire set of nodes to zero and creates new nodes as a unit.
• Nodes cannot be individually added to or removed from a multi-host node pool after creation.
• During upgrades, GKE atomically recreates the entire node pool rather than performing a rolling update.^[7]

A container requesting TPU resources in GKE must consume all TPU chips on the node; partial consumption is not allowed. TPU slice nodes carry a google.com/tpu taint that prevents non-TPU workloads from being scheduled on them.^[7]

Fault tolerance and recovery

At the scale of thousands of TPU chips, hardware failures are expected. TPU workers and their infrastructure include several resilience mechanisms.

ICI resiliency

For TPU v4, v5p, and Ironwood (v7), ICI resiliency is enabled by default for slices of one cube (4x4x4, or 64 chips) or larger. When an optical ICI link or optical circuit switch fails, the system routes traffic around the fault. This improves scheduling availability at the cost of a temporary performance reduction.^[5]

Multislice auto-recovery

In a Multislice configuration, if one slice experiences a failure, Cloud TPU automatically creates a replacement slice. All other slices in the environment are restarted to re-establish the distributed training job. Proper checkpoint management is required to resume training from the last saved state.^[4]

Checkpointing best practices

Because multi-host node pools are recreated atomically on failure, TPU training jobs should save checkpoints frequently to durable storage such as Google Cloud Storage. Frameworks like Orbax (for JAX) and standard PyTorch checkpointing utilities handle distributed checkpoint saving across workers.^[3]

TPU worker performance across generations

The following table summarizes peak per-chip performance for each TPU generation:

Scroll sideways for more →

TPU v1 was designed for inference only.^[9] All subsequent generations (v2 onward) support both training and inference.^[13]

How does a TPU worker differ from a GPU worker?

TPU workers differ from GPU-based workers in several fundamental ways.

Scroll sideways for more →

The architectural difference in interconnect topology is notable. GPUs use hierarchical switching that approximates point-to-point connections, offering high bandwidth between any two GPUs in the same node but requiring network switches for cross-node communication. TPUs use a torus topology with constant per-device link bandwidth, which provides more predictable communication patterns and lower cost per link, but limits each chip to communicating directly only with its nearest neighbors.^[10]

Which frameworks does a TPU worker support?

TPU workers support three major ML frameworks:

• JAX: The most native TPU framework. JAX's functional programming model maps naturally to XLA's compilation approach and the SPMD execution model. JAX is used extensively within Google for large-scale training.^[10]
• PyTorch: Supported through the PyTorch/XLA library, which traces PyTorch operations into XLA graphs for TPU execution. PyTorch/XLA supports both lazy tensor evaluation and eager mode (with some limitations).
• TensorFlow: The original framework for TPU development. Some TensorFlow operations are not available on TPU, and users must verify compatibility.^[2]

What are TPU workers used for?

TPU workers are well-suited for workloads with specific characteristics:

• Large-scale model training: Training large language models (LLMs), vision transformers, and other large architectures that benefit from high-bandwidth matrix multiplication.
• Recommendation and ranking: Models with large embedding tables benefit from SparseCores.^[11]
• Extended training runs: Jobs that run for weeks or months, where TPU energy efficiency translates to meaningful cost savings.
• Large batch training: Workloads that can use large effective batch sizes to fully utilize TPU compute capacity.
• Inference at scale: Serving trained models with low latency and high throughput, particularly for transformer-based models.

History

Google began internal TPU development in 2013 under Amir Salek, who was recruited to establish a custom silicon team.^[13] Norman P. Jouppi served as the technical lead and principal architect, directing the design and deployment of the first TPU in approximately 15 months. The first TPU (v1) was publicly unveiled at Google I/O in 2016 and had been used internally since 2015 to power services including Google Search, Google Photos, and Google Translate.^[9] The TPU v1 was also used in the AlphaGo system that defeated world Go champion Lee Sedol in March 2016.^[13]

Google made TPUs commercially available through Google Cloud Platform in 2018. Broadcom has served as a co-developer across all TPU generations, providing technologies such as SerDes high-speed interfaces and managing fabrication through TSMC.^[13]

Jonathan Ross, one of the original TPU engineers, later founded Groq, a company building its own ML accelerator chips.^[13]

See also

• Tensor Processing Unit
• TPU master
• TPU pod
• TPU node
• TPU chip
• Cloud TPU
• Data parallelism
• Model parallelism
• XLA
• JAX
• GPU computing

References

1. Google Cloud. "TPU architecture." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/system-architecture-tpu-vm ↩
2. Google Cloud. "Introduction to Cloud TPU." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/intro-to-tpu ↩
3. Google Cloud. "Training on TPU slices." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/training-on-tpu-pods ↩
4. Google Cloud. "Cloud TPU Multislice overview." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/multislice-introduction ↩
5. Google Cloud. "TPU v4." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/v4 ↩
6. Google Cloud. "TPU v5e." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/v5e ↩
7. Google Cloud. "About TPUs in GKE." Google Kubernetes Engine Documentation. https://docs.cloud.google.com/kubernetes-engine/docs/concepts/tpus ↩
8. Jouppi, N. P., et al. "TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings." Proceedings of the 50th Annual International Symposium on Computer Architecture (ISCA '23), 2023. https://arxiv.org/abs/2304.01433 ↩
9. Jouppi, N. P., et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit." Proceedings of the 44th Annual International Symposium on Computer Architecture (ISCA '17), 2017. ↩
10. JAX Team. "How to Think About TPUs." Scaling Book. https://jax-ml.github.io/scaling-book/tpus/ ↩
11. Google Cloud. "Cloud TPU VMs with Ranking and Recommendation are generally available." Google Cloud Blog. https://cloud.google.com/blog/products/compute/cloud-tpu-vms-are-generally-available ↩
12. Google Cloud. "The world's largest distributed LLM training job on TPU v5e." Google Cloud Blog. https://cloud.google.com/blog/products/compute/the-worlds-largest-distributed-llm-training-job-on-tpu-v5e ↩
13. Wikipedia contributors. "Tensor Processing Unit." Wikipedia. https://en.wikipedia.org/wiki/Tensor_Processing_Unit ↩
14. Google Cloud. "Inside the Ironwood TPU codesigned AI stack." Google Cloud Blog. https://cloud.google.com/blog/products/compute/inside-the-ironwood-tpu-codesigned-ai-stack ↩
15. Google Cloud. "TPU architecture (TPU VM and worker definitions)." Google Cloud Documentation. https://docs.cloud.google.com/tpu/docs/system-architecture-tpu-vm ↩
16. AI Wiki. "TPU master." https://aiwiki.ai/wiki/tpu_master ↩