TPU vs GPU：硬件 roofline 的三条分水岭
TPU vs GPU：硬件 roofline 的三条分水岭
前情说明：技术深度 D2。《读 JAX Scaling Book》系列第 5 篇（也是最后一篇）。
前四篇讲了算法侧的一切——Roofline 怎么算、Sharded Matmul 怎么切、Transformer Math 怎么核账、LLaMA 3 70B 训练实操 怎么选并行。这些结论都带着一个隐含假设:硬件参数是给定的。
这最后一篇反过来——假设你能选硬件。为什么同样 989 TFLOPs 的 H100 和 918 TFLOPs 的 TPU v6e,跑 LLaMA 70B 的 TPO 可能差 2 倍?为什么 Meta 用 H100 必须 TP≤8 而 Google 用 TPU v5p 可以 TP=128?GB200 NVL72 到底是营销噱头还是真的改变了游戏规则?
答案全在三条硬件分水岭上:matmul 密度 / 内存层级 / 互联拓扑。学完这篇你能在面试被问"选哪种硬件"时,直接把需求转成硬件参数的具体数字对比。
一、分水岭 1:matmul 密度与"模块化程度"
1.1 两种设计哲学
TPU 哲学: 极少数巨型 TensorCore
─────────────────────────────────────
v5p: 每片 2 个 TensorCore
每个 TC 有 1 个大 MXU (128×128 systolic array)
每周期 bf16[8, 128] · bf16[128, 128] → f32[8, 128]
加 4 组独立 ALU 的 VPU (8×128 布局)
↓
总共 2 TC × 4 MXU × 大 systolic array
= 极少数巨型算子,控制简单
GPU 哲学: 大量小型 SM
─────────────────────────────────────
H100: 132 个 SM
每个 SM 分 4 个 subpartition
每 subpartition 一个 Tensor Core + 32 CUDA 核
Tensor Core 可以做 ~8×8×8 matmul/cycle
↓
总共 132 × 4 = 528 个独立 Tensor Core
+ 132 × 4 × 32 ≈ 17000 CUDA 核
= 高度并发,但靠编译器/runtime 调度
类比:TPU 像一头拉货的大象(1-2 个 TC,极重的 systolic array);GPU 像一群蚂蚁(132 个 SM 独立运作)。单位 FLOPs 两者差不多,但做事方式完全不同。
1.2 modularity 的代价
TPU:
✓ 控制简单: 单线程 scalar core 统一调度
✓ 编译器主导:XLA 把一切展平成 vector pipe
✗ 控制流不灵活:数据相关的 branch 很难做
✗ 不规则 workload (稀疏/图/RL) 适配差
GPU:
✓ SIMT (Single Instruction Multiple Threads)
✓ 每个 CUDA core 有独立 instruction pointer
✓ "warp divergence" 允许分支,只是掉速
✓ 可以同时跑几十个独立 kernel
✗ L2 cache 共享,SM 之间"spooky action"
✗ 想达到 roofline 需要调 kernel 技巧多
关键结论:TPU 更容易接近 roofline,但 workload 必须工整。GPU 什么都能跑,但要达到 roofline 需要很多调优。
1.3 单芯片 FLOPs 密度对比
硬件 bf16 TFLOPs fp8/int8 TFLOPs
───────────────────────────────────────────────
TPU v5p 459 918
TPU v6e 920 1840
TPU v7 (Trillium 下代) 2300 4600
───────────────────────────────────────────────
A100 312 624
H100 989 1980
H200 989 1980
B200 2250 4500
单从 FLOPs 看,同代产品 H100 比 TPU v5p 多 2x,但 TPU v6e 又反超 H100。所以不能只看算力——真正决定胜负的是内存和互联。
二、分水岭 2:内存层级与"VMEM/SMEM 的差异"
这是面试被问"TPU 和 GPU 到底哪里不一样"时最容易答不好的点。内存层级决定你的 kernel 能不能被 overlap。
2.1 对标表
层级 TPU v5p H100 作用
─────────────────────────────────────────────────────────
寄存器 256 kB (VREGs) 33 MB (per SM) 最快
一级 scratch 128 MB (VMEM) 33 MB (SMEM) 片上快速
二级 cache -- 50 MB (L2) SM 共享
主存 96 GB (HBM3) 80 GB (HBM3) 大容量
跨芯片 ICI 180 GB/s NVLink 450 GB/s
跨节点 DCN 12.5 GB/s IB 50 GB/s
三个值得注意的差异:
2.2 VMEM vs SMEM:TPU 的最大优势
TPU v5p VMEM = 128 MB (per chip)
H100 SMEM × 132 SM ≈ 33 MB (总和)
↑
还是分散在 132 个 SM 里!
单 SM 只有 256 kB
→ TPU 可以用 VMEM 做"手工 cache"
→ H100 每个 SM 的 SMEM 只能存小块激活
实战意义:TPU 可以把整层权重放进 VMEM,用 VMEM 带宽(~40 TB/s)而不是 HBM 带宽(2.8 TB/s)喂 MXU。等效 critical intensity 从 164 降到 22,小 batch 也能 compute-bound。
这就是为什么 Google 内部服务同样模型可以用极小 batch 做低延迟推理——VMEM prefetching 不是魔法,是物理上多了一层快内存。
2.3 控制粒度:TPU 单线程 vs GPU 多 warp
TPU scalar core:
1 个核控 1 个 VPU(4096 ALU)+ 2 个 MXU + 2 个 XLU
每 cycle 只能发出 1 个 DMA 请求
↓
编译器(XLA)必须手动流水化一切
GPU SM:
每 SM 可以调度 64 个 "resident warp"
Warp Scheduler 自动在 warp 间切换掩盖 I/O
↓
类似多线程 CPU,硬件自己做 overlap
所以 GPU 对新 workload 鲁棒——写个基础 kernel 就能跑起来。TPU 对老 workload 极致——XLA 编译后能逼近理论峰值。这也是 Transformer 时代 TPU 卷土重来的根本原因:Transformer 结构工整,正好对 TPU 胃口。
2.4 B200 的新东西:TMEM
Blackwell 做了一个有意思的改动:
B200 新增 TMEM = 256 kB per SM × 148 SM = 38 MB
原因:
Tensor Core 做得太大(估计 ~2048 FLOPs/cycle)
accumulator 大到 SMEM 存不下
→ 需要专门给 TC 的 memory
趋势解读:NVIDIA 在模仿 TPU。Volta 时代 TC 只有 128 FLOPs/cycle,SMEM 够用;到 Blackwell TC 做大了十几倍,被迫加 TMEM——本质上是在抄 TPU 的 VMEM 设计。
2.5 Arithmetic Intensity 对账
critical intensity = C / W_HBM
TPU v5p : 459e12 / 2.8e12 = 164
TPU v6e : 920e12 / 1.6e12 = 575 ← HBM BW 没跟上!
H100 : 989e12 / 3.35e12 = 295
H200 : 989e12 / 4.8e12 = 206 ← HBM3e 加大 BW
B200 : 2250e12 / 8e12 = 281
反直觉结论:TPU v6e 的 critical intensity 高达 575——因为 FLOPs 翻倍但 HBM 没加,极容易 memory-bound。推理 decode 场景 TPU v6e 反而不如 v5p 划算。H200 相当于 H100 的"推理版"——同算力加大 HBM BW,critical intensity 从 295 降到 206,decode 更友好。
三、分水岭 3:互联拓扑——最关键的分水岭
这才是本文的真正主角。两种拓扑决定了你能跑什么 scale 的模型。
3.1 TPU:3D torus,带宽均匀
TPU v5p pod = 16 × 20 × 28 = 8960 chips 组成一个 3D 环面
┌─ 每片 TPU 连接 6 个邻居(x±,y±,z±)
├─ 每条 link 单向 90 GB/s,双向 180 GB/s
├─ 整片 TPU 的 egress = 540 GB/s(3 轴总和)
└─ pod 内 ANY two TPUs 的带宽都几乎相同
关键特征:
✓ 规整 → compiler 可以静态规划
✓ 带宽均匀 → 公式简单
✗ 只能规则通信(不能任意点到点)
✗ 跨 pod 必须走 DCN(12.5 GB/s/chip,断崖)
3.2 GPU:三层交换树,带宽阶梯
H100 cluster = 8 卡/node + 32 node/SU + 多 SU/pod
Pod (1024+ GPU)
│ Spine 层
│ (IB 400 GB/s/node)
▼
SU (256 GPU)
│ Leaf 层
│ (IB 400 GB/s/node)
▼
Node (8 GPU)
│ NVSwitch
│ (NVLink 450 GB/s/chip,full all-to-all)
▼
GPU
关键特征:
✓ Node 内任意点到点,带宽巨大
✗ 跨 node 带宽骤降 9x (450 → 50 GB/s per chip)
✓ IB fat tree 保证"全二分带宽"但 per-chip 很低
3.3 对账:带宽悬崖
场景:每 GPU 可用带宽
TPU v5p (pod 内 ANY 两点): 540 GB/s (3 轴总和,均匀)
H100 (NVSwitch 域内,机内): 450 GB/s
H100 (跨 node,IB 400Gbps): 50 GB/s ← 掉 9 倍!
TPU v5p (跨 pod,DCN): 12.5 GB/s ← 掉 43 倍!
H100 (跨 pod): 50 GB/s (fat tree 守住)
解读:
- TPU 的带宽悬崖在 pod 边界(8960 卡)——pod 内全部均匀 540 GB/s,出 pod 掉 43 倍
- GPU 的带宽悬崖在 node 边界(8 卡)——每过 8 卡就掉 9 倍
这决定了并行策略的选择。
3.4 对并行策略的决定性影响
回忆 Llama 3 实操 里的 TP 临界条件:F > Y × C/W。
TP 极限(F = 28672, Llama 70B):
alpha = C/W Y_max 公式 物理上限
─────────────────────────────────────────────────────────────
TPU v5p (单轴): 2550 11 ∞ (轴可伸)
TPU v5p (3 轴): 850 33 ∞
H100 NVLink: 1099 26 8 (NVSwitch 域)
H100 IB (跨 node): 19780 <2(不可能) 0
B200 NVLink: 2500 11 8
GB200 NVL72: 2500 11 72 (!)
重大发现:
- TPU 的 TP 没有物理上限——理论上可以 TP 到任意数(只要 F 够大)
- H100 TP 永远 ≤8——跨 node 带宽太低,
F > 316480才能跑 TP=16,现实中不可能 - GB200 NVL72 彻底改变规则——72 卡 NVLink 域,TP 可做到 72,和 TPU 的 pod 概念趋同
3.5 为什么 GB200 NVL72 是 architecture 突破
NVIDIA 2025 推出的 GB200 NVL72 不是简单升级:
H100 node: 8 GPU, 每 GPU 450 GB/s, node egress 400 GB/s
B200 node: 8 GPU, 每 GPU 900 GB/s, node egress 400 GB/s
GB200 NVL72: 72 GPU, 每 GPU 900 GB/s, node egress 3600 GB/s (!)
↑ ↑
全 72 卡 full-mesh 9x H100 node
→ NVLink 域从 8 卡扩到 72 卡
→ TP 可以做 TP=16/32/64,不再被 8 卡限死
→ 本质上在模仿 TPU pod 概念
这对 MoE 训练尤其重要——MoE 的 AllToAll 需要大带宽域,8 卡 node 太小塞不下 256 个 expert,但 72 卡域就能把常见 MoE 塞进单域。
3.6 从拓扑看 Collective 代价
两种拓扑的 collective 代价不一样。前面 Sharded Matmul 讲过通用公式,这里补硬件实测:
Collective TPU v5p (pod 内) H100 (node 内) H100 (跨 node)
──────────────────────────────────────────────────────────────────────────
AllGather B / 540 GB/s B / 450 GB/s B / 400 GB/s
ReduceScatter 同上 同上 同上
AllReduce 2x (无 SHARP) 2x (无 SHARP) 2x
AllReduce+SHARP -- ~1.3x (非 2x) ~1.3x
AllToAll B / (4×540) B / (8×450) B / (M×400)
↑ M = node 数
面试锚点:
- AllToAll in-node 比 AllGather 便宜 8 倍(1/8 的广播代价)——这就是 MoE 推理用 AllToAll 的硬件根据
- AllToAll 跨 node 代价暴涨——如果 M=32 nodes,从 450 GB/s 掉到 50 GB/s,掉 72 倍。这就是为什么大 MoE 必须 NVL72
3.7 NVIDIA SHARP:实际 vs 理论
AllReduce 本来是 2 × AG 代价。NVIDIA 从 Hopper 起支持 SHARP(在网络交换机里做 reduction),理论上能把 AllReduce 从 2B/W 降到 B/W——节省 50%。
实际测量:AllReduce 吞吐
不开 SHARP: ~370 GB/s (vs 理论 450)
开 SHARP: ~480 GB/s (vs 理论 900)
─────────────────────────────────
提升只有 30%,远低于理论 75%
原书原话:"由于 LLM 中纯 AllReduce 并不频繁(都被 FSDP 拆成 AG+RS 了),SHARP 对 LLM 场景帮助不大"。面试冷知识:SHARP 更多是营销卖点。
四、三条分水岭汇总
┌────────────────────────────────────────────────────────────────┐
│ 分水岭 │ TPU 做法 │ GPU 做法 │
├──────────────────────┼───────────────────┼─────────────────────┤
│ 1. matmul 密度 │ 极少数巨型 TC │ 大量小型 SM │
│ 控制粒度 │ 单线程 scalar │ 多 warp 并发 │
│ 对编译器依赖 │ 极高 (XLA) │ 低 (手写 kernel) │
├──────────────────────┼───────────────────┼─────────────────────┤
│ 2. 内存层级 │ VMEM 128 MB │ SMEM 33 MB total │
│ 快内存容量 │ 大 → 层级权重缓存 │ 小 → 单 SM 碎片 │
│ HBM 带宽 │ 2.8-7.4 TB/s │ 3.35-9 TB/s │
├──────────────────────┼───────────────────┼─────────────────────┤
│ 3. 互联拓扑 │ 3D torus,均匀 │ 分层 switch 树 │
│ 内部域大小 │ 8960 (full pod) │ 8 (72 NVL72) │
│ 带宽悬崖位置 │ pod 边界 │ node 边界(更小!) │
│ TP 物理上限 │ 无(受公式) │ 8 / 72(NVL72) │
└────────────────────────────────────────────────────────────────┘
一句话:TPU 为了规整牺牲灵活性,GPU 为了灵活性牺牲规整。Transformer 极度规整 → TPU 占便宜;新 workload 探索 → GPU 占便宜。
五、具体场景:谁更划算?
5.1 场景 1:LLaMA 70B 训练
硬件 方案 MFU 成本/day/卡
────────────────────────────────────────────────────────────
TPU v5p (8960) FSDP=2048 + TP=4 ~45% Google 内部
H100 (16000) FSDP=4000 + TP=8 ~40% 云端 $2/hr
GB200 NVL72 FSDP + TP=16/32 ~50% 云端 $4/hr
选哪个:
- Google 内部:TPU v5p——云端出租 $4/hr,内部实际成本更低
- 外部训练:H100——生态成熟,PyTorch 友好
- 2026+ 新训练:GB200 NVL72——TP 域更大,MoE 友好
5.2 场景 2:LLaMA 70B 推理(decode TPOT)
硬件 TP 每卡 HBM BW 决定性瓶颈
────────────────────────────────────────────────
TPU v5p (8 卡) TP=8 2.8 TB/s HBM-bound
TPU v6e TP=8 1.6 TB/s HBM 掉得更厉害!
H100 (8 卡) TP=8 3.35 TB/s ← 快
H200 (8 卡) TP=8 4.8 TB/s ← 更快
B200 TP=8 8 TB/s ← 最快
TPOT 下界 (70B/8 TP, 单卡 17.5 GB):
v5p : 17.5/2.8 = 6.3 ms
v6e : 17.5/1.6 = 11 ms ← 更慢!
H100: 17.5/3.35 = 5.2 ms
H200: 17.5/4.8 = 3.6 ms
B200: 17.5/8 = 2.2 ms
反直觉:TPU v6e 推理比 v5p 慢!因为算力翻倍但 HBM BW 没涨,critical intensity 高达 575,大多数 decode 场景更 memory-bound。
结论:
- 要 decode 延迟:H200 / B200 > H100 > TPU v5p > TPU v6e
- 要 batch 吞吐:TPU v6e > H100(算力翻倍后 compute-bound 区间整体更强)
5.3 场景 3:大 MoE(DeepSeek V3 规模)
架构 E k 总参数 典型 TP AllToAll 代价
──────────────────────────────────────────────────
DeepSeek V3: 256 + 1 shared, k=8, 总 671B,激活 37B
硬件对比:
─────────────────────────────────────────────
TPU v5p pod (8960):
pod 内全部均匀 → EP 可以做 256
AllToAll 代价均匀,设计简单
H100 (多节点):
EP 只能做到 8 (node 内) 否则跨 IB 性能悬崖
DeepSeek V3 实际:node 内 EP=8,跨 node DP
GB200 NVL72:
EP 可做到 72!能把更多 expert 塞进单域
AllToAll 代价不悬崖
↓
MoE 的黄金硬件
所以 DeepSeek V3 架构选择 E=256, top-8 不是巧合——8 是 H100 node 大小,刚好把 expert group 塞进单 NVLink 域。换 NVL72 后他们可以做 E=72 甚至更大。
5.4 场景 4:Context Parallelism 长上下文
Ring Attention 需要大量 KV 跨卡 AllGather/环传
关键因素:低延迟 + 高带宽
TPU v5p 3D torus:
环形拓扑天然适合 ring attention
延迟均匀,无节点边界问题
H100 cluster:
node 内 ring 快,跨 node ring 慢 9x
1M context 跨 100+ 卡时会撞 node 边界
→ 长上下文训练 TPU 更友好
六、决策流程
实际项目选硬件的思考顺序:
START: 任务是什么?
│
▼
┌─────────────────────────────────────┐
│ 1. 这是 Transformer 吗? │
└─────────┬───────────────────────────┘
│
Yes ─────┼───── No
│ └─→ GPU (灵活性高)
▼
┌─────────────────────────────────────┐
│ 2. 要训还是要推? │
└─────────┬───────────────────────────┘
│
Train │ Infer
│
┌───────┴────────┐
▼ ▼
┌──────────┐ ┌─────────────────┐
│训大模型? │ │ 延迟要求低? │
│>70B? │ │ (TPOT < 10 ms) │
└──┬───────┘ └──┬──────────────┘
│ │
Yes│No Yes│No
│ │
▼ ▼
NVL72/ H200/B200 H100/
TPU v5p (HBM BW 大) TPU v5p
│ (够用)
▼
MoE?
├── Yes → NVL72 (大 EP 域) 或 TPU pod
└── No → H100 pod 够用
6.1 真实选型决策示例
需求: 训 70B dense,15T tokens,< 2 月交付,< $20M 预算
─────────────────────────────────────────────────────
方案 A: 16K H100,FSDP=4000 + TP=8
- 时间: 12 天算力(wall-clock ~45 天含准备)
- 成本: 云端 $2/hr × 16000 × 24 × 12 = $9.2M (算力纯值)
- 加上 bootstrap/data: 约 $15M
✓ 符合预算
方案 B: 8960 TPU v5p pod,FSDP=2048 + TP=4
- 时间: 20 天算力
- 成本: Google Cloud $4/hr × 8960 × 24 × 20 = $17.2M
✗ 偏贵(除非 Google 内部打折)
方案 C: 4-pod GB200 NVL72 (2304 卡)
- 单卡 2.25x 算力,但更贵
- 时间: 8 天算力
- 成本: 云端 $4/hr × 2304 × 24 × 8 = $1.8M (!)
✓ 便宜,但 2026 中前供货紧张
决策: 方案 A (H100,已知量产)或 等 C 开放
七、面试锚点
Q1: TPU 和 GPU 最大的架构差异是什么?
A: 三条分水岭:
- matmul 密度:TPU 极少数巨型 TC,GPU 大量小型 SM
- 内存层级:TPU VMEM 128 MB(大且集中),GPU SMEM 33 MB(分散在 132 SM)
- 互联拓扑:TPU 3D torus 均匀带宽,GPU 分层 switch(node 边界有带宽悬崖)
本质:TPU 为规整牺牲灵活性,GPU 为灵活性牺牲规整。Transformer 规整 → TPU 占便宜。
Q2: 为什么 H100 的 TP 永远 ≤8?
A: 物理决定。NVSwitch 域只有 8 卡(每卡 450 GB/s)。跨 node 走 IB,带宽掉到 50 GB/s/卡,C/W = 989e12/50e9 = 19780。TP 临界条件 F > Y·C/W,TP=16 要 F > 316K,现实模型 F 都远小于此。所以 TP 永远不能跨 node。
Q3: GB200 NVL72 为什么是架构突破?
A: NVLink 域从 8 卡扩到 72 卡(每卡 900 GB/s)。TP 可做到 72,本质上在模仿 TPU pod 概念。对 MoE 影响巨大——EP 域从 8 扩到 72,DeepSeek 这种 E=256 的架构可以完全放进一个域内,AllToAll 不再走 IB 悬崖。
Q4: TPU v6e 推理为什么反而比 v5p 慢?
A: v6e FLOPs 翻倍(920 TFLOPs)但 HBM BW 只有 1.6 TB/s(v5p 是 2.8 TB/s)。critical intensity 从 164 涨到 575——意味着compute 和 memory 严重失衡。decode 场景单 token memory-bound,TPOT = 权重大小 / HBM BW,v6e 的 HBM BW 更低,所以更慢。
Q5: VMEM 和 SMEM 本质区别?
A: 三点:
- 容量:VMEM 128 MB(集中),SMEM 33 MB(分散到 132 SM,每 SM 256 kB)
- 可见性:VMEM 对 scalar core 全局可见,SMEM 仅对 SM 可见
- 程序员控制:VMEM 完全手动管理,SMEM 可选 L1 cache 或手动 shared
实战意义:TPU 可以把整层权重放 VMEM 用 40 TB/s 带宽喂 MXU(等效 critical intensity 从 164 降到 22),H100 做不到——单 SM 的 256 kB 装不下大块权重。
Q6: AllToAll 跨 node 为什么代价暴涨?
A: AllToAll 本质是"每张卡给其他每张卡发一份数据"。在单 node(8 卡)可以通过 NVSwitch 一跳完成;跨 node 则所有 traffic 挤 node 的 IB 400 GB/s 出口。假如 M 节点全互发,每卡实际带宽 ≈ 400 GB/s / M = 50 GB/s / (M/8)——M 越大越惨。相比 node 内 450 GB/s,跨 32 node 的 AllToAll 代价涨 36 倍。这就是 MoE 必须 node 内 EP 的根本原因。
Q7: 为什么 NVIDIA SHARP 的实际收益远低于理论?
A: 理论上 SHARP(in-switch reduction)把 AllReduce 从 2B/W 降到 B/W,省 50%。实测只省 30%(370→480 GB/s)。原因:实际网络有协议开销、延迟、SHARP 本身的 scheduling 复杂度等。而且 LLM 场景纯 AllReduce 本来就不多(都被 FSDP 拆成 AG+RS 了),所以 SHARP 对 LLM 几乎没用——这是 NVIDIA 不愿多谈的事实。
Q8: 为什么 TPU 必须靠编译器,GPU 可以手写 kernel?
A: TPU 是单线程 scalar core 控制所有算子——VPU/MXU/XLU/DMA 必须由编译器(XLA)静态流水化,人类写不来。GPU 是多 warp 并发——每 SM 可同时跑 64 个 resident warp,Warp Scheduler 自动在 warp 间切换掩盖 I/O。所以 CUDA 手写 kernel 可以靠硬件自己做 overlap,这是 CUDA 生态能繁荣的硬件根基。
Q9(进阶): 同样 989 TFLOPs 的 H100 和 TPU v5p 哪个更快?
A: 看任务:
- 大 batch 训练(compute-bound):H100 稍快(算力 2.15x)。但 TPU v5p 的 pod 均匀通信省了很多并行开销,实际 MFU 可能相近
- 小 batch decode(memory-bound):TPU v5p HBM BW 2.8 TB/s,H100 3.35 TB/s,H100 快 20%
- MoE(通信密集):TPU 的 pod 均匀带宽碾压 H100 的 node-boundary 悬崖
- 非 Transformer workload(RL, 稀疏图):GPU 的 SIMT 灵活性碾压
Q10(进阶): B200 加 TMEM 说明了什么?
A: NVIDIA 在抄 TPU 的内存设计。TC 从 Volta 的 128 FLOPs/cycle 涨到 B200 的 ~2048——10 几倍。但 SMEM 还是 256 kB/SM,根本喂不饱 TC(accumulator 都存不下)。于是加 TMEM = 256 kB/SM。这是GPU 从"大量小 SM"向"少量大算子"的架构靠拢,本质上在变得像 TPU。配合 NVLink 域从 8 扩到 72(NVL72),可以说 NVIDIA 正在设计上向 TPU 收敛。
八、总结速查卡
┌──────────────────────────────────────────────────────────────┐
│ 硬件 roofline 速查 │
├──────────────────────────────────────────────────────────────┤
│ TPU v5p │
│ 芯片: 459 TFLOPs bf16, 2.8 TB/s HBM, intensity=164 │
│ 内存: 96 GB HBM, 128 MB VMEM, 256 kB VREG │
│ 拓扑: 3D torus 8960 chips, ICI 180 GB/s/axis bidi │
│ 每片 egress 540 GB/s (3 轴总和) │
│ │
│ H100 (NVLink 4.0) │
│ 芯片: 989 TFLOPs bf16, 3.35 TB/s HBM, intensity=295 │
│ 内存: 80 GB HBM, 50 MB L2, 33 MB SMEM total │
│ 拓扑: 8-GPU node (NVSwitch), IB 400 GB/s/node │
│ 每片 NVLink 450 GB/s, IB per GPU 50 GB/s │
│ │
│ B200 / GB200 NVL72 │
│ 芯片: 2250 TFLOPs bf16, 8 TB/s HBM, intensity=281 │
│ 内存: 192 GB HBM, 126 MB L2, +TMEM │
│ 拓扑: NVL72 = 72-GPU full-mesh (NVLink 5.0) │
│ 每片 NVLink 900 GB/s, node egress 3.6 TB/s │
│ │
│ 关键临界(bf16 FSDP B/N) │
│ TPU v5p (3 轴): B/N > 850 │
│ H100 node 内: B/N > 1099 │
│ H100 跨 node (IB): B/N > 19780 (基本做不到) │
│ B200 NVL72: B/N > 35 (极低!) │
│ │
│ 关键临界(TP F > Y·C/W) │
│ TPU v5p: TP 无物理上限,靠 F ≫ 2550·Y │
│ H100: TP ≤ 8 (NVSwitch 域) │
│ NVL72: TP ≤ 72 │
│ │
│ 决策 │
│ 训大 dense: TPU v5p pod / H100 cluster │
│ 训大 MoE: NVL72 / TPU v5p pod │
│ 低延迟推理: H200 / B200 (HBM BW 大) │
│ 高吞吐推理: TPU v6e / B200 │
│ 长上下文: TPU v5p (3D torus 对 ring attention 友好) │
│ 非 Transformer: GPU (SIMT 灵活) │
└──────────────────────────────────────────────────────────────┘
一句话记住这一篇
TPU 和 GPU 差在三个地方:TPU 用"少量巨型 TC + 128 MB VMEM + 3D torus 均匀带宽"换规整和编译器优化,GPU 用"大量小型 SM + 多 warp 并发 + 分层交换树"换灵活性。Transformer 结构工整 → TPU 占便宜,新 workload 探索 → GPU 占便宜。GB200 NVL72 把 NVLink 域从 8 扩到 72 卡,本质上是 NVIDIA 在向 TPU 的 pod 概念靠拢——两种架构正在趋同,但出发点截然不同。
系列总结
《读 JAX Scaling Book》5 篇到此为止:
- Roofline 分析:算一个 kernel 该多快
- Sharded Matmul:分布式 matmul 的记号系统
- Transformer Math:参数/FLOPs/KV Cache 精算
- LLaMA 3 70B 训练实操:选并行、算预算
- 本文:TPU vs GPU 硬件 roofline
学完这 5 篇你应该能做到:
- 对任何 Transformer 模型秒算参数量、FLOPs、KV Cache
- 判断一个 kernel 到底被什么 bound 住(compute / memory / comms)
- 选对并行方案,算出每种配置的临界 batch
- 估算训练预算(几天、几张卡、几美元)
- 根据任务选硬件,知道每种选择的 trade-off
配套阅读:JAX Scaling Book Ch2 + Ch12 原书章节。
相关文章:分布式 Transformer 深度解析、MoE 推理完整机制、Flash Attention 原理(如有)。
Written by
Zoe
AI Infra Engineer · LLM Serving · GPU/RDMA · 造工具的偏执狂