Prepare and test your skills
Prepare and test your skills
Data parallelism strategies scale training by replicating the same model on every worker and splitting the training data among the replicas. On Google Cloud, Cloud TPU (Tensor Processing Unit) (opens in a new tab) provides the hardware substrate for this pattern, with single-host and multi-slice TPU7x (opens in a new tab) topologies acting as worker pools. Effective design matches the worker topology to the model size and keeps every replica fed with data so collective communication such as All-Reduce scales efficiently.
The TPU execution foundation is a matrix processor specialized for neural-network multiply-and-accumulate operations, unlike general-purpose CPUs and GPUs. A CPU based on von Neumann architecture loads values from memory, performs one calculation, and stores the result back for every operation, which creates the von Neumann bottleneck where slow memory access limits throughput. A GPU increases throughput with 2,500–5,000 Arithmetic Logic Units (ALUs) that execute thousands of multiplications and additions at once, but each ALU calculation must still read operands from registers or shared memory and store intermediate results.
Cloud TPU removes that per-calculation memory access with a systolic array where thousands of multiply-accumulators connect directly to each other in a large physical matrix. Cloud TPU v3 contains two systolic arrays of 128 x 128 ALUs on a single processor, and as each multiplication executes the result passes directly to the next multiply-accumulator with no memory access during matrix multiplication. The output is the summation of all multiplication results between data and parameters, which gives TPUs high computational throughput on large neural networks.
The following comparison explains why data-parallel replicas scale better on TPUs for large matrix work, because each worker spends cycles on chained accumulation rather than repeated memory reads and writes.
| Processor | Execution model | Memory behavior per calculation |
|---|---|---|
| CPU | General-purpose von Neumann processor, flexible for any software | Loads operands from memory and stores result back for every calculation, limited by von Neumann bottleneck |
| GPU | General-purpose processor with 2,500–5,000 ALUs for massive parallelism | Must access registers or shared memory to read operands and store intermediate results for every ALU calculation |
| Cloud TPU | Matrix processor with systolic arrays, for example two 128 x 128 arrays on TPU v3 | Passes results directly between connected multiply-accumulators with no memory access during matrix multiplication |
Choose TPUs when the workload is massive matrix operations and choose CPUs or GPUs when general-purpose software flexibility is required, since TPUs cannot run word processors, control systems, or transaction processing.
Data flow through the TPU worker is the fixed host-to-TPU-to-host path that moves training data from the host CPU through on-chip memory and compute and back to the host. The TPU host streams training data into an infeed queue, the TPU stages data and parameters in High Bandwidth Memory (HBM), and the Matrix Multiplication Unit (MXU) performs the systolic computation. This separation keeps parameters and data staged in HBM before compute begins, so parallel workers do not stall waiting for host memory during the matrix operation.
Execution on the TPU worker follows an ordered sequence where each stage depends on the previous stage:
Throughput in this pipeline depends on the host keeping the infeed queue full and draining the outfeed queue promptly, because the MXU depends on HBM staging and HBM staging depends on the host. Frameworks such as PyTorch and JAX run workloads on TPUs through Compute Engine, Google Kubernetes Engine, and Vertex AI using this same path.
Exam tip: Data enters through the infeed queue to HBM to the MXU, and results leave through the outfeed queue to host memory — parameters load from HBM into the MXU before data streams through the systolic array.
The TPU7x worker pool building block is a full-host, 4-chip virtual machine (VM) that scales from single-host to multi-slice topologies. Each TPU7x VM contains 4 chips with 224 vCPUs, 960 GB of RAM, and access to 2 Non-Uniform Memory Access (NUMA) nodes, and all TPU7x slices use this 4-chip VM shape. The VM boot disk is Hyperdisk Balanced by default, with Hyperdisk Balanced and Hyperdisk ML supported for additional attached storage. Because the unit of scale is the 4-chip VM connected by high-bandwidth interconnects (ICI), adding workers means adding VMs or slices rather than subdividing a VM.
Training-optimized models on TPU7x are tuned and validated to take advantage of the dual-chiplet architecture, ICI, HBM, and mixed-precision compute such as FP8 and BF16 without code changes. These recipes configure optimal parallelism strategies, including data, tensor, pipeline, and context parallelism, to maximize throughput and scaling efficiency across single-host and multi-slice topologies. Supported resources for configuring distributed training on TPU7x include:
Choose a single-host topology when the model and data-parallel replicas fit within one 4-chip VM, and choose a multi-slice topology when scaling requires distributing replicas and model shards across multiple VMs joined by ICI.
Model parallelism partitions a large model that exceeds single-device memory across multiple accelerators so each device holds and computes only part of the model. On TPU Multislice (opens in a new tab), which connects multiple TPU slices over a data center network (DCN), model parallelism means deciding how weights, activations, and pipeline stages are placed within a slice versus across slices. The choice controls memory pressure, DCN communication volume, required batch size, and accelerator utilization.
Sharding options across TPU Multislice determine how model state and data are distributed to achieve maximum performance. The same sharding strategy that worked within a single slice can be kept within each slice, while cross-slice traffic flows over the slower DCN boundary. The documentation describes three practical choices for cross-slice sharding:
Choose among these options based on whether memory or batch size is the binding constraint, because data-based strategies reduce communication efficiency at small batches while pipeline strategies introduce idle time.
Tensor parallelism shards individual layers or activations across the model dimensions so multiple devices cooperate on the same operation. Tensor parallelism across TPU slices is not recommended because sharding activations across model dimensions requires too much inter-slice bandwidth. The DCN between slices cannot sustain the fine-grained, high-frequency communication that tensor-parallel matrix multiplications need, so the workload becomes communication-bound. Tensor parallelism remains useful within a slice where interconnect bandwidth is high, but it should not be used as the cross-slice strategy.
Pipeline parallelism places consecutive groups of layers on different slices and streams micro-batches through those stages in order. A global batch flows from the first stage slice to the next stage slice, then to the final stage slice, with forward activations moving forward and gradients moving backward across the DCN boundary. Pipeline parallelism requires micro-batching to keep multiple stages busy, and micro-batching decreases the effective batch size, arithmetic intensity, and ultimately model FLOPS utilization. Pipeline parallelism also incurs the pipeline bubble, where chips sit idle because they are waiting for data from an upstream or downstream stage, so minimizing bubbles depends on balancing stage sizes and scheduling enough micro-batches to fill the pipeline.
Device placement in Multislice separates high-bandwidth intra-slice work from DCN-limited inter-slice work. Computation inside a slice uses the fast intra-slice interconnect, while synchronization between slices depends on the DCN network and is the limiting factor for throughput. In data parallelism, every step performs all-reduce over DCN to synchronize gradients, using 32 bits of DCN bandwidth per parameter for two slices and up to 64 bits per parameter when more than two slices trigger a full-shuffle all-reduce. In FSDP, each operation all-gathers the needed weights so each slice has the weights it needs, then reduce-scatters gradients as they are produced so each slice receives only the gradients for the weights it owns. FSDP therefore decreases weights and optimizer-state memory per slice as slice count grows, but it increases network traffic and the risk of blocking on a delayed collective.
Choosing pipeline parallelism over data-based strategies is a batch-size decision driven by convergence limits and DCN costs. Pure data parallelism and FSDP both require scaling the global batch size linearly with the number of slices, and strong scaling requires the backward-pass time to exceed the DCN all-reduce time. For Transformer language models, operational intensity equals 4 * B / 32 for two slices and 4 * B / 64 for many slices, where B is batch size in tokens, which implies a theoretical minimum of 176k to 352k tokens and a practical margin of at least 350k tokens per Pod for two Pods to 700k tokens for many Pods. Pipeline parallelism allows pipeline slices to share a batch instead of scaling it, so it becomes relevant when that required batch exceeds the preferred maximum batch for convergence.
The tradeoffs across memory, communication, and batch size determine the selection order:
| Strategy | Memory effect | Communication and batch tradeoff |
|---|---|---|
| Data parallelism | Each slice holds full weights and optimizer state | All-reduce over DCN each step; needs very large batch for strong scaling |
| Fully sharded data parallelism (FSDP) | Weights and optimizer state per slice shrink as slices grow | All-gather plus reduce-scatter over DCN; still needs linearly scaled batch but stores more activations or parameters |
| Pipeline parallelism | Each slice holds only its stages | Stage-to-stage activation and gradient transfers over DCN plus pipeline bubble; shares one batch but lowers utilization through micro-batching |
Start with pure data parallelism, use FSDP to relieve memory pressure or grow parameters and activations, and use pipeline parallelism only when those approaches require too large a global batch. FSDP tends to achieve higher model FLOPS utilization, so pipeline parallelism is the better choice only when convergence per sample slows at the batch size needed for efficient FSDP.
Exam tip: Do not use tensor parallelism across TPU slices; use pipeline parallelism across slices only when data parallelism or FSDP would require a global batch larger than the convergence limit.
Hardware-specific optimizations and trade-offs compare specialization versus flexibility when training on Cloud TPU Pods or multi-GPU configurations. Cloud TPUs (opens in a new tab) are application-specific integrated circuits designed as matrix processors for neural-network workloads, while graphics processing units (GPUs) and central processing units (CPUs) remain general-purpose processors that support many kinds of software. Choosing between them depends on model size, dominant operations, memory access patterns, interconnect needs, and framework support for data and model parallelism.
Core specialization determines how efficiently hardware converts parallel math into training throughput. A CPU follows von Neumann architecture where values are loaded from memory, calculated, and stored back for every operation, which creates the von Neumann bottleneck as memory access limits throughput despite full software flexibility. A GPU increases throughput with thousands of Arithmetic Logic Units (ALUs), typically 2,500–5,000 per processor, that execute thousands of multiplications and additions simultaneously for matrix operations. A Cloud TPU removes per-operation memory access by connecting thousands of multiply-accumulators directly into a systolic array, so once data and parameters are loaded, results pass from one accumulator to the next as a summation without intermediate memory reads.
Cloud TPU v3 illustrates the scale of specialization with two systolic arrays of 128 x 128 ALUs on a single processor. The comparison below summarizes the design and memory tradeoff that drives hardware selection.
| Processor | Primary design | Memory behavior per calculation |
|---|---|---|
| CPU | General-purpose von Neumann processor | Load value, calculate, store result for every operation |
| GPU | General-purpose parallel processor with thousands of ALUs | Read operands from registers or shared memory and store intermediates for each ALU operation |
| Cloud TPU | Matrix processor with systolic arrays | Stream operands once, then pass partial sums between multiply-accumulators with no memory access during multiplication |
Use Cloud TPUs when training is dominated by large multiply-and-accumulate operations, and use CPUs or GPUs when the workload needs general-purpose software or diverse non-matrix logic.
The Cloud TPU memory hierarchy moves data from the host through on-chip memory and back in a fixed order that keeps the matrix unit fed. The hierarchy depends on HBM to hold both data and parameters close to the MXU and depends on the host for feeding and draining the queues. The trade-off is high computational throughput for neural-network calculations at the cost of inability to run general software such as word processors, control systems, or transaction processing.
The flow proceeds through an ordered sequence where compute depends on staging and staging depends on the host:
If the host fails to keep the infeed full or to drain the outfeed promptly, the MXU stalls even though the matrix hardware itself is ready.
Interconnect bandwidth determines whether distributed training scales linearly or stalls on communication between chips. TPU slices connect many chips with reconfigurable high-speed links, and large slices use Inter-Chip Interconnect (ICI) to form a mesh across hosts. A TPU v2 (opens in a new tab) slice can scale to 512 chips interconnected in one fabric and is requested by accelerator type such as v2-8 for a single device with 8 cores or v2-128 for 128 TensorCores. When the ICI mesh cannot be discovered, frameworks such as JAX fail initialization with link errors between processing elements, which shows that the training job depends on correct physical topology as well as software configuration.
TPU7x (opens in a new tab) continues this design with high-bandwidth ICI and High Bandwidth Memory (HBM) plus mixed-precision compute in FP8 and BF16 to raise throughput per chip. Supported models on TPU7x are tuned for dual-chiplet architecture, ICI, HBM, and mixed precision without code changes. Choose Pod-scale TPU topologies when gradient synchronization, tensor sharding, and collective operations need dense, low-latency chip-to-chip paths across tens to hundreds of chips.
The TPU7x virtual machine (VM) is the unit of host compute, memory, and locality for Ironwood (TPU7x) slices. Each TPU7x VM uses the tpu7x-standard-4t machine type in a full-host configuration with 4 chips per VM and access to two Non-Uniform Memory Access (NUMA) nodes. The technical specification per VM is 224 vCPUs, 960 GB of RAM, and 2 NUMA nodes. Because all TPU7x slices use 4-chip VMs, input pipelines, data loading threads, and host-side preprocessing scale with large vCPU and memory capacity on the same host that feeds the 4 local chips.
Parallelism strategy maps model size and sequence length to hardware by splitting data, tensors, pipeline stages, or context across devices. Training-optimized models on TPU7x provide pre-configured recipes and hyperparameters for distributed training frameworks that combine data parallelism, tensor parallelism, pipeline parallelism, and context parallelism to maximize throughput and scaling efficiency across single-host and multi-slice topologies. The open-source MaxText framework is a high-performance large language model (LLM) training framework written in JAX and optimized for Cloud TPU, with official training recipes that include precision configurations, step-time benchmarks, and throughput metrics for popular open models. Inference-optimized models use vLLM with the tpu-inference plugin, which provides hardware-optimized kernels and tensor parallelism with optimized PagedAttention kernels validated for correctness and serving throughput.
Use data parallelism when the full model fits on each device and the goal is higher global batch size, and add tensor, pipeline, or context parallelism when weights, activations, or long context exceed single-device HBM. Framework choice follows the same boundary: JAX with MaxText fits large-scale TPU training with explicit sharding, PyTorch fits portable training across TPU Pods, and vLLM with tpu-inference fits high-performance serving on TPU7x.
Capacity mode controls visibility, isolation, and responsibility for failures at large scale. Managed TPU capacity holds back part of the reservation so Google can automatically replace faulty machines and restart slices, while All Capacity mode (opens in a new tab) exposes the full reservation with no holdbacks and makes the user responsible for managing failures and planned maintenance. All Capacity mode provides dedicated, physically co-located capacity with ultra-low-latency networking, precise control over maintenance by VM, cube, Pod, or reservation, and topology-aware scheduling based on full views of health and utilization. Its physical topology has four levels where a cluster is a deployment unit of Pod multiples in one zone, and blocks, sub-blocks, and hosts map to physical components below the cluster.
Exam tip: All Capacity mode is supported on Trillium (TPU v6e) and Ironwood (TPU7x) and future generations, not on older TPU generations.
Distributed training selection moves from hardware capability to sharding scope to batch feasibility. Data parallelism on TPU systolic arrays and ICI fabrics maximizes throughput while the full model fits per device and the host sustains infeed and outfeed, FSDP extends that data-based path by sharding weights and optimizer state at the cost of all-gather and reduce-scatter over DCN, and pipeline parallelism across Multislice becomes the fallback when linear batch scaling would exceed convergence limits. TPU7x 4-chip VMs, Pod-scale interconnects, and All Capacity mode then provide the locality, bandwidth, and placement control to implement the chosen strategy.
Cloud TPU (Tensor Processing Unit) passes partial sums directly between connected multiply-accumulators with no memory access during matrix multiplication, while CPUs and GPUs access memory for every calculation. A CPU loads operands from memory and stores the result back for every operation which creates the von Neumann bottleneck, a GPU uses 2,500 to 5,000 Arithmetic Logic Units that must still read registers or shared memory for each calculation, and Cloud TPU v3 contains two systolic arrays of 128 x 128 Arithmetic Logic Units that accumulate results across the array.
Data flows from the TPU host through the infeed queue to High Bandwidth Memory to the Matrix Multiplication Unit, and results return through the outfeed queue to host memory. The TPU host streams input data into the infeed queue, the TPU loads data into High Bandwidth Memory and loads parameters from High Bandwidth Memory into the Matrix Multiplication Unit for systolic execution, then loads completed results into the outfeed queue for the host to read.
Use pipeline parallelism across TPU Multislice only when data parallelism or Fully Sharded Data Parallelism would require a global batch larger than the convergence limit. Pure data parallelism and Fully Sharded Data Parallelism both require scaling the global batch linearly with the number of slices and need very large batches for strong scaling, while pipeline slices share one global batch but lower utilization through micro-batching and incur pipeline bubbles when stages wait for upstream or downstream data.
Last updated