TPU v4: An Optically Reconfigurable Supercomputer for Machine LearningTPU v4: An Optically Reconfigurable Supercomputer for Machine Learning

Reading plan (40 minutes)閱讀計畫（40 分鐘）

Read §2 slowly; it is the network. §2.1–2.2 give the cube and the OCS, §2.3 availability, §2.4–2.6 the other motivations, §2.8–2.9 the twisted torus and slice statistics. §3 (SparseCore) is about embeddings and uses the interconnect in a different way; skim it. §5–§7 compare against TPU v3, the A100 and Graphcore's IPU. The paper is written by computer architects, so it argues from cost, power and utilization rather than from protocols.

慢慢讀 §2；它就是網路。§2.1–2.2 講 cube 與 OCS，§2.3 講可用性，§2.4–2.6 其他動機，§2.8–2.9 扭轉環面與切片統計。§3（SparseCore）講 embedding，用互連的方式不同；略讀即可。§5–§7 與 TPU v3、A100 及 Graphcore 的 IPU 比較。這篇由計算機架構師撰寫，所以論證從成本、功耗與使用率出發，而非從協定出發。

1. The problem Google set itself1. Google 給自己設定的問題

TPU v3 pods had 1,024 chips in a fixed 2-D torus. Google wanted four times as many chips, but a rigid machine of 4,096 chips and 1,024 CPU hosts has a problem that grows with size: if each host is unavailable 0.1–1.0% of the time for repairs and updates, then at any moment several hosts, and the chips attached to them, are down, and in a fixed topology a dead node breaks the torus for the whole slice that contains it. Google also wanted to sell slices of the machine to many customers at once, which requires isolation, and to let each job pick a topology that fits its parallelism. A packet-switched network would give flexibility but at a cost in power, latency and switch hardware that an architect paper weighs carefully. The answer was to keep the torus, whose links are cheap and whose routing is trivial, and to make the links between racks re-pluggable with optical circuit switches.

TPU v3 的 pod 有 1,024 顆晶片組成固定的二維環面。Google 想要四倍的晶片，但一台由 4,096 顆晶片與 1,024 台 CPU 主機組成的剛性機器有一個隨規模增長的問題：若每台主機因維修與更新有 0.1–1.0% 的時間不可用，那麼任何時刻都有數台主機及其晶片處於停機狀態，而在固定拓樸中，一個死掉的節點會讓包含它的整個切片的環面斷裂。Google 也想同時把機器的切片賣給許多客戶，這需要隔離；還想讓每個任務挑選符合其平行化的拓樸。封包交換網路能提供彈性，但在功耗、延遲與交換器硬體上有代價，架構論文會仔細衡量。答案是保留環面——其鏈路便宜、路由簡單——並用光路交換器讓機櫃之間的鏈路可以重新插接。

2. Cubes, faces and mirrors2. 立方體、面與鏡片

The chip and its links. Each TPU v4 chip has six inter-chip interconnect (ICI) links, each 50 GB/s, roughly 400 Gb/s. Six links give a chip a neighbor in each direction of three dimensions: $\pm x$, $\pm y$, $\pm z$. There is no NIC and no switch; a chip forwards traffic for its neighbors in hardware. Four chips share one CPU host.

晶片與它的鏈路。每顆 TPU v4 晶片有六條晶片間互連（ICI）鏈路，每條 50 GB/s，約 400 Gb/s。六條鏈路給一顆晶片在三個維度每個方向各一個鄰居：$\pm x$、$\pm y$、$\pm z$。沒有 NIC、沒有交換器；晶片以硬體為鄰居轉送流量。四顆晶片共用一台 CPU 主機。

The cube. 64 chips in a 4×4×4 arrangement, with their 16 hosts, fit in one rack. Inside the cube the links are electrical. A 4×4×4 block has six faces, and each face has 16 chips on it; each of those chips has one link pointing out of the cube. So each face has 16 outgoing links and the cube has 96. These 96 links are optical, and they are the only links that leave a rack.

立方體。64 顆晶片以 4×4×4 排列，加上 16 台主機，裝進一個機櫃。cube 內部的鏈路是電的。一個 4×4×4 區塊有六個面，每個面上有 16 顆晶片；每顆都有一條指向 cube 外的鏈路。所以每個面有 16 條向外的鏈路，一個 cube 有 96 條。這 96 條鏈路是光的，也是唯一離開機櫃的鏈路。

The optical circuit switch. Google's Palomar OCS has 136 ports: 128 in use and 8 spares for testing and repair. Inside, an array of micro-electro-mechanical mirrors tilts to reflect the light from any input fiber into any output fiber. The OCS never decodes a packet; it makes a physical light path between two ports and holds it until told to change, which takes milliseconds. Optical circulators let one fiber carry light in both directions, halving the fiber and port count. The two opposite faces of a cube (for example $+x$ and $-x$) are wired to the same set of OCSes: each cube sends 2 links to each of 48 OCSes. The link and port counts match exactly, one question per row:

光路交換器。Google 的 Palomar OCS 有 136 個埠：128 個使用中、8 個備用供測試與維修。內部有一排微機電鏡片，傾斜後把任何輸入光纖的光反射進任何輸出光纖。OCS 從不解碼封包；它在兩個埠之間建立一條物理光路，並保持到被告知改變為止，改變需要毫秒。光循環器讓一條光纖同時雙向傳光，光纖與埠數減半。一個 cube 的兩個相對面（例如 $+x$ 與 $-x$）接到同一組 OCS：每個 cube 向 48 台 OCS 各送 2 條鏈路。鏈路數與埠數剛好相等，逐列計算如下：

Every port in use on an OCS carries exactly one cube link. By setting the mirrors, the OCSes decide which cube's $+x$ face connects to which cube's $-x$ face, and so build any torus out of any set of cubes.

OCS 上每個使用中的埠剛好承載一條 cube 鏈路。藉由設定鏡片，OCS 決定哪個 cube 的 $+x$ 面接到哪個 cube 的 $-x$ 面，因此能用任何一組 cube 組出任何環面。

3. What reconfigurability buys3. 可重構性買到什麼

Availability. This is the argument the paper leads with. Figure 4 sweeps host availability from 99.0% to 99.9% and measures goodput, the useful work delivered. Without OCSes, a fixed torus delivers reasonable goodput only at 99.9% host availability; a single bad host in a slice spoils the slice. With OCSes, the scheduler simply picks other cubes and rewires the torus around the failure, and goodput is fair even at 99.0–99.5%. Spare capacity is the limit for very large slices: a slice of one quarter of the machine reaches 75% goodput at those availabilities.

可用性。這是論文領頭的論點。Figure 4 把主機可用性從 99.0% 掃到 99.9% 並量測 goodput（真正交付的有用工作）。沒有 OCS 時，固定環面只有在 99.9% 的主機可用性下才有合理的 goodput；切片中一台壞主機就毀了整個切片。有 OCS 時，排程器只需挑其他 cube 並把環面繞過故障重新接線，即使在 99.0–99.5% 的可用性下 goodput 仍可接受。備用容量是極大切片的限制：占整台機器四分之一的切片在那些可用性下達到 75% 的 goodput。

Deployment and modularity. TPU v3 could not enter production until all 1,024 chips and cables were installed and tested. With TPU v4, each 4×4×4 block went into production as soon as its 64 chips and cables worked, because a slice can be built from any set of working cubes.

部署與模組化。TPU v3 必須等全部 1,024 顆晶片與線纜安裝並測試完畢才能投入生產。TPU v4 的每個 4×4×4 區塊只要 64 顆晶片與線纜正常就能投入生產，因為切片可以由任何一組正常的 cube 組成。

Slices shaped to the job. A slice can be any $4i \times 4j \times 4k$ shape, so sizes need not be powers of two: a 192-chip job gets $4 \times 4 \times 12$. More importantly, a job can pick the shape that matches its parallelism. A pipeline-parallel job with a long chain of stages can ask for a 512-chip "cigar" of $4 \times 4 \times 32$, which gives a long dimension for the pipeline; a job that wants the most bisection bandwidth asks for the $8 \times 8 \times 8$ cube instead. Table 3 shows that choosing the topology and hyperparameters per job gained 1.2× to 2.3×: a 2.3× difference between a novice's choice and the best choice for an LLM on 512 chips, and 1.2× over an expert's pick for GPT-3 pre-training.

依任務塑形的切片。切片可以是任何 $4i \times 4j \times 4k$ 的形狀，所以大小不必是二的冪：192 顆晶片的任務拿到 $4 \times 4 \times 12$。更重要的是，任務可以挑選符合其平行化的形狀。有一長串階段的管線平行任務可以要一個 512 顆晶片、$4 \times 4 \times 32$ 的「雪茄」，給管線一個長維度；想要最大對分頻寬的任務則要 $8 \times 8 \times 8$ 的立方體。Table 3 顯示依任務選擇拓樸與超參數帶來 1.2 到 2.3 倍的增益：在 512 顆晶片上訓練 LLM，新手的選擇與最佳選擇相差 2.3 倍；GPT-3 預訓練則比專家的選擇好 1.2 倍。

Isolation. Two slices never share a switch or a link, because each circuit belongs to one slice. The paper calls this an air gap, and it matters when many customers share one supercomputer. A Clos network gives isolation only by software policy.

隔離。兩個切片永遠不會共用交換器或鏈路，因為每條電路屬於一個切片。論文稱之為 air gap，在許多客戶共用一台超級電腦時很重要。Clos 網路只能靠軟體政策提供隔離。

Cost and power. The entire optical fabric, including transceivers, fiber and the OCSes, is under 5% of the system's capital cost and under 3% of its power. The authors' point is that the expensive part of a training machine is the accelerators, and anything that keeps more of them busy more of the time pays for itself.

成本與功耗。整個光學網路——包含光模組、光纖與 OCS——不到系統資本成本的 5%、功耗的 3%。作者的重點是：訓練機器中昂貴的部分是加速器，任何能讓更多加速器更長時間忙碌的東西都划得來。

4. The twisted torus4. 扭轉的環面

In an ordinary torus the wrap-around link from the last chip in a row goes back to the first chip of the same row. Because in TPU v4 those wrap-around connections pass through an OCS, they can be pointed anywhere, and a twisted torus offsets them so that the last chip of one row connects to the first chip of a different row. The effect is that the longest shortest path through the slice gets shorter and more links are available between any two halves, which raises bisection bandwidth without adding any hardware. The measured all-to-all throughput gain is 1.63× on a $4 \times 4 \times 8$ slice and 1.31× on $4 \times 8 \times 8$. Only slices of shape $n \times n \times 2n$ or $n \times 2n \times 2n$ with $n \ge 4$ can be twisted. A sample of production slices from November 2022 showed that 29% were smaller than one cube and could not twist; of the remainder, about half had a twistable shape, and 86% of those were in fact twisted.

在普通的環面中，一列最後一顆晶片的環繞鏈路接回同一列的第一顆。因為 TPU v4 的環繞連線經過 OCS，它們可以指向任何地方，扭轉環面把它們錯開，讓一列的最後一顆接到另一列的第一顆。效果是切片中最長的最短路徑變短、任意兩半之間可用的鏈路更多，不加任何硬體就提高了對分頻寬。實測的 all-to-all 吞吐增益在 $4 \times 4 \times 8$ 切片是 1.63 倍，在 $4 \times 8 \times 8$ 是 1.31 倍。只有 $n \times n \times 2n$ 或 $n \times 2n \times 2n$（$n \ge 4$）形狀的切片可以扭轉。2022 年 11 月生產切片的抽樣顯示 29% 小於一個 cube、無法扭轉；其餘中約一半形狀可扭轉，而那些之中有 86% 確實被扭轉了。

5. SparseCore: the other use of the interconnect5. SparseCore：互連的另一種用法

Recommendation models spend their time looking up embeddings: small vectors fetched from huge tables by random index. These accesses are tiny, random and spread across the whole machine, the opposite of the large, regular collectives of a dense model. TPU v4 adds SparseCores, small dataflow processors that use about 5% of the die area and 5% of the power, each with 16 compute tiles tied to HBM channels. They operate on a single flat, globally addressable memory space: all the HBM in a 4,096-chip machine, 128 TiB, is one address space reachable over ICI. The torus serves this as a remote-memory fabric rather than as a collective fabric, and the paper reports a 5–7× speedup on embedding-heavy models. For a network reader the lesson is that the same links are used in two completely different patterns, and the torus handles both because routing is deterministic and the links are not shared with other jobs.

推薦模型的時間花在查 embedding：以隨機索引從巨大的表中取出小向量。這些存取很小、隨機、散布整台機器，與稠密模型大而規律的集體操作相反。TPU v4 加入 SparseCore，一種小型資料流處理器，使用約 5% 的晶粒面積與 5% 的功耗，每個有 16 個計算 tile 綁定 HBM 通道。它們在單一、平坦、全域可定址的記憶體空間上運作：4,096 顆晶片機器中所有的 HBM，共 128 TiB，是一個透過 ICI 可達的位址空間。環面在此是遠端記憶體網路而非集體通訊網路，論文報告 embedding 密集模型有 5–7 倍加速。對網路讀者的教訓是：同樣的鏈路被以兩種完全不同的模式使用，而環面兩種都能處理，因為路由是確定的、鏈路也不與其他任務共用。

6. Results6. 結果

Against its predecessor, TPU v4 delivers 2.1× the performance and 2.7× the performance per watt of TPU v3. Against the NVIDIA A100, the GPU of the same generation, on MLPerf-style training workloads it is 1.2–1.7× faster and uses 1.3–1.9× less power; the mean power measured during BERT training was 197 W per TPU v4 against 380 W per A100 (1.93×), and during ResNet 206 W against 273 W. Against Graphcore's IPU Bow it is 4.3–4.5× faster at 256 chips. For large language models the authors report that TPU v4 sustains about 60% of peak FLOP/s, and in a later section note that PaLM held 57.8% over its 50-day run. Compare this with the 38–55% MFU in the GPU papers of this primer; the authors credit part of the difference to a topology that matches the communication pattern and has no congestion.

相對前一代，TPU v4 的效能是 TPU v3 的 2.1 倍、每瓦效能 2.7 倍。相對同世代的 GPU NVIDIA A100，在 MLPerf 類型的訓練工作負載上快 1.2–1.7 倍、功耗低 1.3–1.9 倍；BERT 訓練期間實測平均功耗每顆 TPU v4 為 197 W、每顆 A100 為 380 W（1.93 倍），ResNet 則是 206 W 對 273 W。相對 Graphcore 的 IPU Bow，在 256 顆晶片時快 4.3–4.5 倍。對大型語言模型，作者報告 TPU v4 持續達到約 60% 的峰值 FLOP/s，並在後面章節指出 PaLM 在 50 天的運行中維持 57.8%。把它與本入門中 GPU 論文的 38–55% MFU 比較；作者把部分差異歸功於符合通訊模式、且沒有壅塞的拓樸。

7. Critique and context7. 評論與脈絡

The torus works because training collectives are neighbor-to-neighbor by construction: ring all-reduce, pipeline send/receive and 2-D/3-D decompositions all map onto a mesh. It works less well for all-to-all traffic, which is exactly why the paper spends a section on twisting (to raise bisection bandwidth) and why mixture-of-experts models are a harder fit. The design is also vertically integrated: Google controls the chip, the compiler (XLA chooses the collective algorithms for the slice shape), the scheduler and the OCS, which no GPU-and-Ethernet operator can match. What the paper does not give is a transport story: there is no discussion of lossless delivery, flow control or failure recovery at the link level, because with dedicated circuits and deterministic routing those problems are much smaller. Read it against HPN and the rail-only paper: all three agree that the network should be shaped to the parallelism and that optics are a small share of cost, and they disagree on whether any-to-any reachability is worth keeping. Google later brought OCSes into its general data-center network, Jupiter (SIGCOMM 2022), and TPU v5p and later generations keep the torus-plus-OCS design.

環面可行，是因為訓練的集體操作天生就是鄰居對鄰居：環狀 all-reduce、管線 send/receive 與二維/三維分解都能映射到網格上。它對 all-to-all 流量就沒那麼好，這正是論文花一節講扭轉（提高對分頻寬）的原因，也是混合專家模型較難適配的原因。這個設計也是垂直整合的：Google 掌控晶片、編譯器（XLA 依切片形狀選擇集體演算法）、排程器與 OCS，任何 GPU 加 Ethernet 的營運者都無法比擬。論文沒有給的是傳輸層的故事：沒有討論無損傳遞、流量控制或鏈路層的故障恢復，因為有專用電路與確定性路由，這些問題小得多。把它與 HPN 及 rail-only 論文對照：三者都同意網路應依平行化塑形、光學只占成本一小部分，但對是否值得保留任意對任意的可達性有分歧。Google 後來把 OCS 帶進其通用資料中心網路 Jupiter（SIGCOMM 2022），TPU v5p 與之後的世代也維持環面加 OCS 的設計。

Key numbers關鍵數字

Terms used in this paper本篇用到的術語

Torus

A mesh whose edges wrap around so every node has the same number of neighbors; 3-D here, with six links per chip.

邊緣環繞連接、每個節點鄰居數相同的網格；此處是三維，每顆晶片六條鏈路。

ICI

Inter-chip interconnect, the direct chip-to-chip links of a TPU; no NIC or switch in between.

晶片間互連，TPU 直接晶片對晶片的鏈路；中間沒有 NIC 或交換器。

Cube

A 4×4×4 block of 64 chips with 16 hosts in one rack; the unit of deployment and of slice construction.

一個 4×4×4、64 顆晶片、16 台主機的區塊，裝在一個機櫃；部署與組成切片的單位。

OCS

Optical circuit switch: MEMS mirrors connect an input fiber to an output fiber; no packet processing; reconfigures in milliseconds.

光路交換器：微機電鏡片把輸入光纖接到輸出光纖；不處理封包；毫秒級重新配置。

Slice

The set of cubes allocated to one job, wired by the OCSes into a torus of shape $4i \times 4j \times 4k$.

分配給一個任務的 cube 集合，由 OCS 接成 $4i \times 4j \times 4k$ 形狀的環面。

Twisted torus

A torus whose wrap-around links are offset to a different row, shortening paths and raising bisection bandwidth.

環繞鏈路錯開到不同列的環面，縮短路徑並提高對分頻寬。

Goodput

Useful work delivered per unit time, after losses from unavailable hosts and unusable slices.

扣除不可用主機與無法使用切片的損失後，每單位時間交付的有用工作。

SparseCore

A small dataflow processor on the TPU v4 die for embedding lookups over the global 128 TiB memory space.

TPU v4 晶粒上的小型資料流處理器，在 128 TiB 的全域記憶體空間上做 embedding 查找。