TPU：为什么谷歌不排队等 NVIDIA GPU（2/2）
Source: Dev.to
深度学习的精度格式
传统科学计算使用 FP64（双精度）或 FP32（单精度），精度极高。
然而深度学习并不需要如此高的精度。因此 Google 创建了 bfloat16（Brain Floating Point），一种 16 位格式，保留了 FP32 的宽动态范围（≈ 1e‑38 到 1e38），但牺牲了一些小数位精度。
- FP16 的范围有限（约 6e‑5 到 6e4），可能导致训练不稳定。
- bfloat16 保持了 FP32 的范围，因而适合 AI 工作负载。
NVIDIA 随后在其 A100 和 H100 GPU 上采用了 bfloat16。
TPU Pods 与芯片间互连
单个 TPU 芯片在矩阵乘法上表现出色，但无法单独处理当今的大模型。Google 将芯片组织成层级化的 TPU Pod：
- TPU 芯片 → TPU 板卡
- 板卡 → TPU 机架
- 机架 → TPU Pod
一个 Pod 最多可容纳 4,296 芯片，软件将其呈现为单一的、极度并行的处理器。
芯片间互连 (ICI)
标准以太网对于训练期间持续、低延迟的数据交换来说太慢。TPU Pods 使用专用的 芯片间互连 (Inter‑Chip Interconnect, ICI)，直接绕过 CPU。芯片以 3‑D 环形拓扑（类似“甜甜圈”形状）连接，使任意芯片只需几跳即可到达最远的芯片。
光学电路交换机 (OCS) – TPU v4
TPU v4 引入了 光学电路交换机 (Optical Circuit Switch, OCS)，省去了电‑光转换步骤：
- MEMS 镜面 倾斜以引导携带光信号的数据束，几乎零延迟。
- 弹性恢复：若部分芯片失效，镜面可重新排列，实现即时流量重路由。
大规模散热
成千上万的芯片聚在一起会产生巨大的热量。Google 采用 直达芯片液冷，将冷却管道直接铺设在芯片上——实际上把数据中心变成了巨大的水族箱。这一做法早于 NVIDIA 最近在 H100 上采用的液冷方案。
软件栈：从 TensorFlow 到 JAX
| 组件 | 角色 | 输入 | 输出 |
|---|---|---|---|
| JAX (前端) | 用户接口 | Python 代码 | 中间表示 (HLO) |
| XLA (后端) | 编译器引擎 | HLO | 计算图 / TPU/GPU 可执行二进制 |
- 自动微分 (grad) 与 向量化 (vmap) 由 JAX 处理。
- XLA 执行 kernel 融合，减少内存访问次数，充分利用 TPU 的 systolic array。
Ironwood TPU (TPU v7)
2025 年 Google 宣布了 Ironwood，第七代 TPU，旨在兼顾大语言模型推理效率和大规模训练。
- FP8 支持（首款原生 8 位浮点的 TPU）。
- 计算性能：4,614 TFLOPS（FP8）。
- 内存：每芯片 192 GB HBM3E，带宽 7.37 TB/s（受内存带宽限制的工作负载）。
- Pod 扩展：每个 Pod 最多 9,216 芯片。
- ICI 带宽：1.2 TB/s 双向。
- 能效：比 TPU v6 提升约 2 倍，仍使用直达芯片液冷。
来源：Google Cloud Blog – Ironwood TPU
为什么 GPU 仍然占主导
- 软件生态：NVIDIA 的 CUDA 自 2006 年起不断演进；超过 90 % 的 AI 研究者编写 CUDA 代码。
- 框架优化：PyTorch 对 CUDA 进行了大量优化。虽然 PyTorch 能在 TPU 上运行，但体验不够成熟。
- 可获取性：GPU 可以购买并在本地部署；TPU 只能通过 Google Cloud Platform 使用。已经在 AWS 或 Azure 上的组织很难采用 TPU。
因此，尽管 TPU 在硬件上具备优势，成熟的软件堆栈和 GPU 的易获取性仍使其在 AI 开发中保持领先地位。