Google TPU ICI
ICI 如何直连形成 Torus 拓扑,以及 OCS 光交换层怎么支持重构
核心要点:
- ICI 是 TPU 专用片间互联,不依赖外部交换机
- 直接形成 2D / 3D Torus 拓扑,TPU 仅与近邻直连
- 单跳延迟百 ns 级,比 IB Switch 低一个数量级
- TPU v4 引入 OCS 光交换层,让 3D Torus 可重构
- AllReduce 按维度分解,多个 Ring 并行执行
- 与 NVLink 的对比:ICI 牺牲拓扑灵活换超低延迟+超大规模
代际规格怎么演进的?
端口数固定在 4-6,带宽随代际翻倍。Pod 规模从 256 chips 扩到 9216 chips,见 @tbl-hw-ici-gen。
| TPU 代际 | 年份 | 每芯片 ICI 带宽 (双向) | ICI 端口数 | 支持拓扑 | Pod 规模 |
|---|---|---|---|---|---|
| TPU v2 | 2017 | ~500 Gbps | 4 | 2D Torus | 256 chips |
| TPU v3 | 2018 | ~656 Gbps | 6 | 2D Torus | 1024 chips |
| TPU v4 | 2022 | ~4.8 Tbps | 6 | 3D Torus + OCS | 4096 chips |
| TPU v5e | 2023 | ~1.6 Tbps | 4 | 2D Torus | 256 chips/slice |
| TPU v5p | 2023 | ~4.8 Tbps | 6 | 3D Torus | 8960 chips |
| Ironwood (TPU v7) | 2025 | ~9.6 Tbps (1200 GB/s) | 4 (OSFP, 3D 邻居 6) | 3D Torus | 9216 chips |
@tbl-hw-ici-gen TPU ICI 代际规格
口径说明:TPU v5e 面向推理,端口少成本低;v5p 和 Ironwood 面向大规模训练。带宽数字为双向聚合,各代因物理链路数和编码效率不同口径略有差异。
为什么 ICI 选择直连 Torus 而非交换机?
牺牲拓扑灵活性换超低延迟、带宽可预测和线性扩展。
基本原理:每个 TPU 的 ICI 端口直接连相邻 TPU。多 TPU 在 X/Y/Z 三个维度排列,首尾相连形成环 (Torus)。
2D Torus (TPU v3, 4×4 示意):
[T00] - [T01] - [T02] - [T03]
| | | |
[T10] - [T11] - [T12] - [T13]
| | | |
[T20] - [T21] - [T22] - [T23]
| | | |
[T30] - [T31] - [T32] - [T33]
行、列首尾相连 (T00-T03, T00-T30 等),构成环形
3D Torus (TPU v4): X 维度同行相邻 + Y 维度同列相邻 + Z 维度不同层相邻,每 TPU 6 个 ICI 端口 (±X, ±Y, ±Z)。
无交换机的优势:
- 延迟极低:单跳百 ns 级 (vs IB Switch ~100-200 ns 端口延迟累积)
- 带宽可预测:ICI 链路独占,无共享竞争,集合通信时间可精确建模
- 线性扩展:Pod 规模扩大时每 TPU 端口数不变,仅增加芯片数
直连的代价:
- 拓扑固定:物理连线决定通信路径,无法像 Fat-tree 那样任意重构
- 非全连接:非相邻 TPU 通信需多跳 (最大跳数 = 维度大小一半),远距离延迟较高
- 布线复杂:3D Torus 需跨机架长距离线缆 → 引入 OCS 的动机
OCS 怎么让 3D Torus 变成可重构的?
Optical Circuit Switch 在机架间引入可重配置光交换层,把 Z 维度连接关系变成可编程。
背景问题:3D Torus 物理布线要求 Z 维度相邻 TPU 物理相邻,但机架空间布局使 Z 连接往往需穿越整个机房 — 普通铜缆撑不住,光纤又固定了拓扑。
OCS 方案:
- 每机架 TPU 通过光纤连 OCS 端口
- OCS 按需在任意两端口建立光路 → 等价于"这两机架在 Torus Z 维度相邻"
- 重配置时间约数十毫秒,适合按任务静态配置
对拓扑的实质影响:
- 不同训练任务可配不同机架-机架映射,每次用最优 3D Torus 形状
- 等价于"可重构 3D Torus",而非固定物理接线
- 机架间不再需固定顺序物理相邻,数据中心机架排布自由度大幅提升
集合通信优化:
- 矩形 Pod (如 4×4×16) 可先在小维度 (4×4) 内 AllReduce,再沿大维度 (16) AllReduce
- OCS 在不同任务间切换映射,避免多任务共享时拓扑争抢
AllReduce 在 ICI Torus 上怎么实现?
按维度分解,每维度内部独立 Ring AllReduce,维度间串行。
DOR (Dimension-Order Routing):数据包先沿 X,再沿 Y,最后沿 Z 路由。保证无死锁,路径长度确定 (最大跳数 = 各维度半径之和)。
AllReduce 维度分解 (3D Torus,X × Y × Z):
步骤 1: X 维度 Ring ReduceScatter → 每个 Z×Y 面内 X 环并行
步骤 2: Y 维度 Ring ReduceScatter → 每个 Z×X 面内 Y 环并行
步骤 3: Z 维度 Ring AllGather → 每个 X×Y 面内 Z 环并行
(实际实现中 ReduceScatter 和 AllGather 维度顺序有多种等效方案)
各维度 Ring 独立执行,维度间串行。总通信量与平坦 Ring AllReduce 相同,但 Torus 并行度更高 (多环同时工作),整体延迟更低。
延迟模型 (3D Torus,X×Y×Z,消息 M bytes):
$$\begin{equation} T_{\text{allreduce}} \approx 2 \cdot \frac{X-1}{X} \cdot \frac{M}{BW_x} + 2 \cdot \frac{Y-1}{Y} \cdot \frac{M}{BW_y} + 2 \cdot \frac{Z-1}{Z} \cdot \frac{M}{BW_z} + T_{\text{overhead}} \label{eq:hw-ici-allreduce} \end{equation}$$$BW_{x/y/z}$ 为各维度 ICI 链路带宽,$T_{\text{overhead}}$ 为跳数相关的基础延迟。
TPU 集群外部互联怎么做?
Pod 内 ICI 直连 + Pod 间标准以太网,分层 AllReduce。
层级 1: Pod 内 (ICI 直连)
技术: ICI (2D/3D Torus 或 Mesh)
带宽: 每芯片 1.6-9.6 Tbps (代际相关)
延迟: 百 ns 级
角色: TP/SP AllReduce, 大规模 DP 内部 AllReduce
层级 2: Pod 间 (数据中心以太网)
技术: 标准 Ethernet
带宽: 数十 GB/s (商用网卡)
延迟: us 级
角色: 跨 Pod DP 梯度同步, 模型 checkpoint 存储
分层 AllReduce:先 Pod 内 ICI 局部 AllReduce (延迟低带宽高),再 Pod 间以太网汇总 (因已局部聚合数据量大幅减小)。逻辑类似 NCCL tree AllReduce,但 Pod 内完全由 ICI 硬件完成。
与 NVLink / IB 怎么对比?
ICI 以无交换、超低延迟、超大规模为目标,代价是拓扑固定。对比见 @tbl-hw-ici-vs-others。
| 指标 | ICI (TPU v4) | ICI (Ironwood) | NVLink 4.0 (H100) | IB NDR |
|---|---|---|---|---|
| 每芯片带宽 (双向) | ~4.8 Tbps | ~9.6 Tbps | 900 GB/s | 400 Gbps (~50 GB/s) |
| 需要交换机 | 否 (OCS 可选) | 否 | 需 NVSwitch | 需 IB Switch |
| 拓扑 | 3D Torus (可重构) | 3D Torus | 全连接 (NVL72) | Fat-tree |
| 单跳延迟 | 百 ns (无交换) | 百 ns | ~1-2 us | ~1-2 us (含 Switch) |
| 单 Pod/Domain 最大规模 | 4096 chips | 9216 chips | 576 chips (NVL72×多机架) | 数万端口 (Fat-tree) |
| 集合通信实现 | XLA/JAX 内置维度分解 | XLA 内置 | NCCL | NCCL/MPI |
| 异构互联支持 | 否 (TPU 专用) | 否 | 是 (NVLink-C2C) | 是 (通用 RDMA) |
@tbl-hw-ici-vs-others ICI vs NVLink / IB
核心差异:
- ICI:无交换 + 超低延迟 + 超大规模直连,代价是拓扑固定、仅 TPU 同构
- NVLink:全连接 + 高带宽 + 生态兼容,借 NVSwitch 提供全连接但引入交换延迟
- InfiniBand:通用 + 超大规模 Fat-tree,适合异构混合,但节点内带宽远低于 ICI 和 NVLink
Takeaway
| 知识点 | 核心结论 |
|---|---|
| 代际演进 | v2 (500 Gbps) → Ironwood (9.6 Tbps),端口数固定 4-6 |
| 拓扑哲学 | 直连 Torus 无交换机,单跳百 ns,但非全连接需多跳 |
| OCS 创新 | TPU v4 起 Z 维度可重构,任务间切换映射 |
| AllReduce | 维度分解,各维度 Ring 并行,维度间串行 |
| 集群分层 | Pod 内 ICI + Pod 间以太网,分层 AllReduce |
| vs NVLink/IB | ICI 牺牲灵活换延迟和规模;仅 TPU 同构可用 |