PaperlayerView exact revision on arXiv

Source-linked AI summary

TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

Norman P. Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, Cliff Young, Xiang Zhou, Zongwei Zhou, David Patterson

arXiv:2304.01433v3cs.ARcs.AIcs.LGcs.PF

TL;DR

The paper addresses how rapidly changing ML workloads, especially large-scale models and embedding-heavy recommendation systems, require scalable and reliable ML supercomputers. It presents TPU v4’s reconfigurable optical interconnect, SparseCore embedding support, and model–hardware co-optimization, reporting system-level performance and sustainability benefits.

  • Problem

    Rapidly changing ML workloads create scale, reliability, topology, and embedding-processing demands for ML supercomputers.

  • Method

    TPU v4 combines OCS-based reconfigurable topologies, SparseCore processors for embeddings, and ML search to co-optimize hardware configurations with DNN models.

  • Results

    TPU v4’s OCSes provide broad system benefits at <5% of overall cost and <3% of overall power, while 3K-chip TPU v4 slices greatly reduce LLM training time over TPU v3.

  • Takeaways & Limitations

    TPU v4 supercomputers support large language model training and can reduce operational energy and CO2e when deployed in energy-optimized warehouse-scale computers.

  • Takeaways & Limitations

    MLPerf’s DLRM benchmark differs substantially from production DLRM workloads in model size, batch size, and feature types, limiting direct realism for production scaling.

Abstract

from arXiv · show

In response to innovations in machine learning (ML) models, production workloads changed radically and rapidly. TPU v4 is the fifth Google domain specific architecture (DSA) and its third supercomputer for such ML models. Optical circuit switches (OCSes) dynamically reconfigure its interconnect topology to improve scale, availability, utilization, modularity, deployment, security, power, and performance; users can pick a twisted 3D torus topology if desired. Much cheaper, lower power, and faster than Infiniband, OCSes and underlying optical components are <5% of system cost and <3% of system power. Each TPU v4 includes SparseCores, dataflow processors that accelerate models that rely on embeddings by 5x-7x yet use only 5% of die area and power. Deployed since 2020, TPU v4 outperforms TPU v3 by 2.1x and improves performance/Watt by 2.7x. The TPU v4 supercomputer is 4x larger at 4096 chips and thus ~10x faster overall, which along with OCS flexibility helps large language models. For similar sized systems, it is ~4.3x-4.5x faster than the Graphcore IPU Bow and is 1.2x-1.7x faster and uses 1.3x-1.9x less power than the Nvidia A100. TPU v4s inside the energy-optimized warehouse scale computers of Google Cloud use ~3x less energy and produce ~20x less CO2e than contemporary DSAs in a typical on-premise data center.

1 INTRODUCTION

TPU v4 responds to rapidly changing ML workloads and supercomputer-scale reliability challenges with reconfigurable optical networking, embedding-specific hardware, and model–hardware co-optimization. These features target scale, topology, reliability, and performance for workloads including LLMs and recommender systems.

  • Motivation: Rapidly evolving ML models, including LLMs, embeddings, Transformers, and BERT, drove Google’s scale from 256 TPU v2 nodes to 4096 TPU v4 nodes.This scale increased reliability challenges for checkpoint/restore DNN training.
  • Optical Circuit Switches: TPU v4 introduces OCSes with optical data links to support a 4K-node supercomputer while tolerating 1K unavailable CPU hosts.The unavailable-host rate is 0.1%–1.0%.
  • SparseCore: SparseCore provides TPU hardware support for embedding workloads in deep learning recommendation models.The paper describes and evaluates this as the first accelerator support for embeddings in a commercial ML system.
  • Topology: Embedding workloads add all-to-all communication that strains bisection bandwidth, while OCSes enable flexible topologies such as twisted 3D tori.This differs from backpropagation’s all-reduce pattern, which maps well to 2D and 3D tori.
  • Impact: OCS topology flexibility and deployment advantages unexpectedly improved LLM training time beyond their original scale and reliability goals.The paper focuses on TPU v4’s three novel features: OCSes, SparseCore support, and ML-based co-optimization.

2 RECONFIGURABLE OPTICAL SWITCH

TPU v4 uses optical circuit switches to connect 4×4×4 building blocks into a 4096-chip, 64-rack supercomputer. Its construction combines electrically connected packages and trays with an optical fabric spanning racks.

  • Building block: 4×4×4 building blocks provide the electrically cabled TPU v4 unit, with 64 chips and 16 CPU hosts fitting in one rack.
  • Optical fabric: 48 OCSes connect 64 4×4×4 blocks, yielding the 4096-chip TPU v4 supercomputer.Each block connects to 48 OCSes through paired opposing-face links.
  • Optical fabric: The Palomar OCS provides 128 ports plus 8 spares for link testing and repairs.
  • System assembly: The 48 OCSes join eight rack rows into the complete 64-rack system.
  • Hardware packaging: Each TPU v4 package contains an ASIC and four HBM stacks, while each board mounts four liquid-cooled packages and provides PCIe and inter-tray ICI connectors.

2.3 OCS Availability Benefits

OCS connectivity improves TPU v4 availability and deployment by routing around host failures and allowing independently deployable rack-scale blocks. Scheduling can also use flexible, non-power-of-two slice geometries.

  • Availability: OCSes route around CPU-host failures, preserving reasonable slice goodput at lower host availability than a statically connected system.Without OCSes, host availability must reach 99.9% for reasonable slice goodput.
  • Production slices: Table 2 samples production slice popularity for a day in November 2022, including every slice used at least 0.1%.Half of the slices have dimensions whose values are either 4 or 8.
  • Deployment: OCSes made each 4×4×4 block independently deployable once its chips and cables were installed and tested.This replaced TPU v3’s requirement that all 1024 chips and cables be ready before use.
  • Scheduling: TPU v4 slices can use geometries 4^i×4^j×4^k, enabling requests such as a 192-chip 4×4×12 slice.

2.6 OCS Modularity and Security Benefits

OCS reconfiguration lets TPU v4 match interconnect topology to each job’s parallelism and communication needs. The system also supports flexible slice placement, isolation, and topology choices beyond a fixed cubic arrangement.

  • Reconfiguration: OCS circuits switch in milliseconds, allowing topology changes for the application, node count, and system running the jobs.
  • Topology and parallelism: Users can select TPU v4 topologies to match data, model, or pipeline parallelism, including cigar-shaped or cube-shaped slices.Embedding-heavy applications often prefer the cube for highest bisection bandwidth.
  • Topology and parallelism: Varying topology and hierarchy produces performance gains of 1.2x to 2.3x in the examples reported in Table 3.
  • Per-job configuration: TPU v4 uses one static topology per training job, which can be co-optimized with that job’s communication requirements.Per-job configuration is not a fundamental limitation of the OCS.

2.8 Twisting the Torus

TPU v4 can reconfigure optical links to use twisted tori, improving communication properties without physical recabling. Twisted topologies are selectively useful across production slice geometries and add little system cost and power.

  • Twisted topology: A twisted torus rewires selected links between 4×4×4 blocks through OCS routing changes rather than physical recabling.
  • Measured performance: Twisted tori improve all-to-all throughput by 1.63x on 4×4×8 slices and 1.31x on 4×8×8 slices over regular tori.
  • Figure interpretation: TPU v4’s 3D network is illustrated in two dimensions, with electrical links fixed and optical links changing between regular and twisted configurations.
  • Production usage: 40% of production topologies with at least 4^3 blocks use twisted tori.Only geometries meeting the required n×n×2n or n×2n×2n forms can twist.
  • System overhead: The complete optical fabric costs <5% of TPU v4 supercomputer capital costs and consumes <3% of total power.The accounting includes optics modules, fiber, and OCS infrastructure.

3. SPARSECORE: EMBEDDINGS SUPPORT

SparseCore is a dedicated dataflow architecture for embedding training, matching embeddings’ sparse memory and communication patterns rather than dense computation. It improves embedding performance through globally addressable HBM/ICI resources and higher-bisection-bandwidth topologies.

  • Embedding workloads: Embeddings convert categorical feature values into dense vectors through lookup tables, forming a key first layer in production DLRMs.Tables may contain many sizes and support univalent or multivalent lookups.
  • Embedding workloads: Embedding tables can reach O(100 GiB) individually and several TiB in aggregate, requiring partitioning across TPU chips.Column, row, and table sharding provide model-parallel distribution; small tables may instead be replicated.
  • Performance demands: Embedding lookups are bottlenecked by memory bandwidth, capacity, vector processing, and interconnect performance rather than chip FLOPS.Model parallelism uses variable-length all-to-all exchange, making network bisection bandwidth especially important.
  • SparseCore design: SparseCore co-design combines supercomputer-scale HBM, dedicated ICI, and fast gather/scatter support for embedding training.Separate SparseCores enable parallel execution across dense computation, sparse computation, and ICI communication.
  • Topology and results: 2D-to-3D torus scaling raises bisection bandwidth from N^1/2 to N^2/3; TPU v4 reaches 2–4x higher bandwidth and 1.1x–2.0x embedding speedups over TPU v3.At 1024 chips, SparseCore overheads begin to dominate, reducing sensitivity to bisection bandwidth.
  • Production result: At 128 chips, TPU v4 beats TPU v3 by 3.1x and CPUs by 30.1x on the internal DLRM0 recommendation model.The comparison includes TPU v4 configurations without SparseCore, where embeddings reside in CPU memory.

4 USING ML TO TAILOR THE DNN TO THE TPU AND THE TPU TOPOLOGY TO THE DNN

The paper uses platform-aware neural architecture search to co-optimize DNN computation and TPU v4 topology. This search improves quality–performance trade-offs and finds topology changes that outperform both novice and expert designs.

  • DNN co-optimization: PA-NAS automatically tailors DNN models to TPU v4 supercomputers while preserving comparable accuracy.A PA-NAS-designed CNN1 achieves approximately 1.6x better QPS and latency than a generic-NAS baseline.
  • DNN co-optimization: PA-NAS shifts computation between SparseCores and TensorCores to optimize DLRM performance and quality.DLRMs use sparse and dense layers, so balancing their execution is central to the optimization.
  • DNN co-optimization: >10% end-to-end performance improvement results when PA-NAS approaches balanced SparseCore–TensorCore utilization on DLRM0.The reported uplift is comparable to historical improvements achieved by more than ten experts over roughly half a year.
  • Topology co-optimization: 2.3x performance improvement comes from changing topology for a 512-chip TPU v4 LLM slice relative to a novice design.The example changes the topology and parallelism configuration for the LLM workload.
  • Topology co-optimization: 1.2x performance improvement over an expert design is demonstrated for GPT-3 pre-training through topology search.The result shows that topology selection remains beneficial even for an expert-designed configuration.

5 PRODUCTION WORKLOAD PERFORMANCE

TPU v4 improves production-workload performance and efficiency over TPU v3 across scale, with especially large gains for embedding-heavy DLRMs and workloads benefiting from CMEM. The reported LLM comparison is bounded because TPU v3 was unoptimized and lacked sufficient capacity.

  • Scalability: Half of eight production workloads scale well to 3K TPU v4 chips.The workloads identified as scaling well are CNN0, RNN0, RNN1, and BERT1.
  • Hardware comparison: 2.2x peak performance and 1.3x HBM bandwidth distinguish TPU v4’s key hardware improvements over TPU v3.The faster clock and increased matrix multipliers drive the peak-performance gain.
  • Workload results: 3.0–3.5x faster DLRM0 and 2.8x faster DLRM1 performance occur at 512 chips versus TPU v3.The DLRM gains are attributed in the passage to TPU v4’s doubled SparseCores and faster clock.
  • Overall results: 2.1x performance and 2.7x performance/Watt are achieved by TPU v4 relative to TPU v3.Approximately 40% of the performance/Watt improvement comes from technology, with the remainder from design improvements.
  • Workload results: 1.2x overall performance gain comes from CMEM, rising to 2x for RNN1.These results are reported with CMEM turned off for the comparison.
  • Scope boundary: TPU v3 LLM performance is omitted because its fixed 2D topology and limited chip capacity prevent a meaningful optimized comparison.The paper states that LLM training will become an MLPerf benchmark in a future release.

6 MLPERF BENCHMARK PERFORMANCE

Against published MLPerf Training results, TPU v4 compares favorably with the Nvidia A100 and Graphcore MK2 IPU on BERT and ResNet. The comparisons use systems and power measurements with differing reporting conditions.

  • Evaluation basis: MLPerf comparisons are based on published results, while vendors may choose the system size used for reporting.The paper notes that equal system size, cost, or power is not required.
  • Performance comparison: 1.15x faster BERT and 1.67x faster ResNet are reported for similar-sized TPU v4 systems versus A100 systems.The comparison uses reported results interpolated by chip count where needed.
  • Performance comparison: Approximately 4.3x faster BERT and 4.5x faster ResNet are reported for TPU v4 versus the Graphcore MK2 IPU.The IPU comparison is limited to the MLPerf workloads and scales shown in the figure.
  • Power comparison: A100 systems use 1.3x–1.9x more power on average than TPU v4 for the measured MLPerf comparisons.The measurements cover 64-chip systems and include DSA plus HBM power.

7 DISCUSSION

The discussion emphasizes that balanced compute, memory, and interconnect design matters more than peak FLOPS alone, while TPU v4’s flexibility and deployment context affect performance, energy, and scalability. It also highlights rapidly changing workloads and the resulting need for flexible architectures and compilers.

  • System balance: 4.3x–4.5x faster on two MLPerf benchmarks, TPU v4 outperforms IPU Bow despite only a 1.10x peak-FLOPS advantage.Against A100, TPU v4 is 1.15x–1.67x faster for the same number of chips despite A100’s 1.13x peak-FLOPS rate.
  • Interconnect trade-offs: OCSes and ICI provide higher bandwidth and lower-power connectivity than Infiniband, but replacing them increases cost, power, and performance penalties.ICI bandwidth is 400 Gbit/s versus 200 Gbit/s for Infiniband; simulated slowdowns depend on collective operation and slice size.
  • Memory and energy: 1.18x performance and 1.24x performance/Watt result when CMEM increases on-chip SRAM from 32 MB to 160 MB.The larger SRAM permits larger DRAM transfers, improving energy efficiency.
  • Energy and carbon: ~18.3x higher operational CO2e is estimated for a contemporary DSA in an average on-premise data center than for TPU v4 in Google Cloud.The corresponding energy comparison is 2.85x more energy for the contemporary DSA.
  • Workload evolution: 4.2x growth in weights and 3.8x growth in embeddings over five years, with 43 DLRM0 releases, underscore the need for compilers that exploit evolving DSA features.The paper frames multi-generational DSA design as a balance between specialization, generality, and flexibility.
  • Benchmark scope: MLPerf DLRM scales only to ≤128 chips, whereas production recommendation workloads reach batch sizes of 2048–4096 and scale usefully to 1024 chips.The paper attributes the difference to model size, batch-size limits, feature counts, and fixed per-batch overheads.

8 RELATED WORK

TPU v4 combines a dedicated 3D torus and reconfigurable optical circuit switching within a broader history of torus, circuit-switching, and embedding-acceleration research. The paper distinguishes its production-scale deployment from prior designs and testbeds.

  • Interconnect architectures: TPU v4 uses a dedicated 3D torus, while Nvidia systems combine NVLink and NVSwitch within GPU groups with Infiniband beyond them.The paper situates TPU v4 among traditional tightly connected 3D-torus supercomputers.
  • Twisted tori: Twisted tori extend earlier 2D and processor-array designs, and TPU v4 follows a k×2k×2k configuration for its 3D topology.The paper identifies prior work on twisted dimensions and doubly twisted tori.
  • Optical circuit switching: TPU v4 is presented as the first commercial supercomputer using OCS and the first with a reconfigurable interconnect that enhances performance.Earlier OCS-based proposals were paper designs or proof-of-concept, testbed-scale demonstrations.
  • Memory systems: TPU v4’s logically shared address space uses software-controlled access and asynchronous DMA writes, differing from the Cray T3E’s richer remote-memory operations.The comparison concerns memory-system semantics and supported remote-access mechanisms.
  • Embedding acceleration: Embedding acceleration has prior work in TPU v2, Neo, and Nvidia MLPerf systems, which use techniques including parallelism, communication overlap, and fused kernels.The paper places TPU v4’s SparseCore in this line of embedding-focused hardware and software research.

9 SUMMARY

TPU v4’s main contributions are reconfigurable optical interconnect and SparseCore embedding acceleration, each delivering broad system benefits at modest hardware cost. Production deployments demonstrate strong performance, scalability, availability, and energy or carbon advantages.

  • SparseCore: 5x–7x embedding acceleration comes from SparseCore, which uses only ~5% of die area and power.SparseCore’s dataflow sea-of-cores architecture can place embeddings anywhere in the supercomputer’s 128 TiB physical memory.
  • Optical circuit switching: <5% of system cost and <3% of system power are consumed by OCSes and optical components that provide eight system-level benefits.These benefits include scalability, availability, modularity, performance, power efficiency, scheduling, deployment, and security.
  • Interconnect comparison: Replacing OCS and ICI with Infiniband increases costs, raises power consumption, and degrades performance.The paper therefore treats optical and integrated interconnect as central to TPU v4’s system advantages.
  • Comparative performance: TPU v4 is faster and lower power than contemporary DSA chips built with similar technologies for similarly sized systems.The paper notes that the power advantage may be larger when interconnects are included.
  • LLM deployment: 57.8% of peak hardware floating-point performance was sustained for 50 days while training the 540B-parameter PaLM model.The paper connects TPU v4’s performance, scalability, and availability with large-language-model training.
  • Energy and carbon: ~20x lower operational CO2e is reported for TPU v4 in energy-efficient Google warehouse-scale computers than for a typical contemporary on-premise ML DSA.The same comparison reports TPU v4 consuming ~⅙–½ of the energy.
Loading 2304.01433v3…