""" ============================================================================= CS349E · 第 8 讲 · 02_tpu_rooflines.py —— 复现 TPU v1 论文的 roofline 分析（Jouppi et al., ISCA 2017） ============================================================================= 【这个脚本要回答的问题】 TPU v1 有 92 TOPS 的峰值，比同时期的 K80 GPU 高 30 多倍，可它的内存带宽（34 GB/s）只有 K80 的五分之一。 Google 数据中心里 95% 的推理负载（MLP、LSTM、CNN 各两个应用）在它上面跑得怎样？瓶颈在哪？ 演示 1：三块芯片的规格与 roofline 拐点（论文表 2、图 5–7）。 演示 2：六个应用在 TPU 上：roofline 上界 vs 实测（论文表 1、表 3）；性能计数器把时间花在哪。 演示 3：和 CPU、GPU 比（论文表 4、表 6）：为什么 K80 只比 Haswell 快一点。 演示 4：如果改设计（论文图 11、§7 的 TPU'）：内存带宽 ×4、时钟 ×4、换 GDDR5——各快多少。 【运行】python3 02_tpu_rooflines.py （< 1 秒） 【依赖】本目录的 hw_engines；../lecture01 的 console。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from hw_engines import HASWELL, K80_DIE, TPU_APPS, TPU_V1, roofline_tops # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 # 论文表 3 的性能计数器（每列是一个应用）：阵列在算、等权重、移权重、其他（非矩阵）各占的周期比例 TABLE3 = { "MLP0": (12.7, 53.9, 15.9, 17.5), "MLP1": (10.6, 44.2, 13.4, 31.9), "LSTM0": (8.2, 58.1, 15.8, 17.9), "LSTM1": (10.5, 62.1, 17.1, 10.3), "CNN0": (78.2, 0.0, 0.0, 21.8), "CNN1": (46.2, 28.1, 7.0, 18.7), } # 论文表 6 的最后一行（TPU / K80，由未四舍五入的数据算出，所以和用上面两列相除略有出入） TABLE6_RATIO = {"MLP0": 16.7, "MLP1": 60.0, "LSTM0": 8.0, "LSTM1": 1.0, "CNN0": 25.4, "CNN1": 26.3} # ----------------------------------------------------------------------------- def demo1_chips(): print_header("演示 1：三块芯片（论文表 2，每个 die）与 roofline 的拐点") rows = [] for c, paper in ((HASWELL, 13), (K80_DIE, 9), (TPU_V1, 1350)): rows.append([c.name, f"{c.nm} nm", f"{c.mhz} MHz", f"{c.tdp_w:.0f} W", f"{c.peak_tops:g}", f"{c.mem_gbs:g}", f"{c.onchip_mib:g} MiB", f"{c.ridge():.0f}", paper]) print_table(["芯片", "工艺", "时钟", "TDP", "峰值 TOPS", "内存 GB/s", "片上存储", "拐点（乘加/字节）", "论文"], rows) print("（TOPS 把乘和加各算一次；拐点的单位是\"每读 1 字节权重做几次乘加\"，所以 = 峰值 / 2 / 带宽。" "K80 的峰值、带宽是论文表 2 关掉 Boost 模式、按一个 die 的数字。）") print_observations([ "TPU 的峰值是 K80 的 33 倍，内存带宽却只有它的约五分之一：拐点在 1350，是 K80（9）的 150 倍。" "也就是说，只有\"每个权重字节被用 1350 次以上\"的应用才能在 TPU 上跑到峰值。", "TPU 的面积主要给了数据通路：24 MiB 的统一缓冲区占芯片的近三分之一，矩阵单元占四分之一，控制逻辑只占 2%（论文图 2）——" "CPU / GPU 的控制逻辑（缓存、乱序执行、多线程）大得多。", ]) # ----------------------------------------------------------------------------- def demo2_apps_on_tpu(): print_header("演示 2：六个应用在 TPU 上——roofline 上界 vs 实测（论文表 1、表 3）") rows = [] for name, (_, ops_b, batch, nw, _, meas, _, _) in TPU_APPS.items(): bound = roofline_tops(TPU_V1, ops_b) act, stall, shift, _ = TABLE3[name] rows.append([name, batch, f"{nw / 1e6:.0f} M", ops_b, "内存" if ops_b < TPU_V1.ridge() else "计算", f"{bound:.1f}", f"{meas:.1f}", f"{meas / bound:.0%}", f"{act:.0f}%", f"{stall + shift:.0f}%"]) print_table(["应用", "批大小", "权重数", "乘加/字节", "roofline 判断", "上界 TOPS", "实测", "实测/上界", "阵列在算", "等权重 + 移权重"], rows) print_observations([ "MLP 和 LSTM 的强度只有 64–200（批大小就是每个权重的复用次数），在拐点 1350 的左边：受限于内存带宽。" "实测达到上界的 43%–90%，性能计数器显示 58%–79% 的周期在等权重或移权重，阵列真正在算的只有 8%–13%。", "CNN 的权重在空间上被反复使用（强度 1750–2888），受限于算力：CNN0 跑到 86 TOPS（峰值的 93%）。", "CNN1 是例外：强度很高，实测却只有 14.1 TOPS。论文表 3 的解释：一半的乘加单元没有有用的权重（有些层的特征深度太浅，" "填不满 256 宽的阵列），而且它最后的 4 个全连接层强度只有 32，等权重占了 35% 的周期。" "roofline 只按整个应用的平均强度算，看不到这种\"一部分层受限于内存\"的情况。", ]) # ----------------------------------------------------------------------------- def demo3_vs_cpu_gpu(): print_header("演示 3：和 CPU、GPU 比（论文表 4、表 6）") print("MLP0 在 99 百分位响应时间 ≤ 7 ms 的要求下（论文表 4，每个 die 的吞吐）：") print_table(["芯片", "批大小", "99% 响应时间", "每秒推理数", "占无限制时最高吞吐"], [["Haswell", 16, "7.2 ms", "5,482", "42%"], ["K80", 16, "6.7 ms", "13,461", "37%"], ["TPU", 200, "7.0 ms", "225,000", "80%"]]) rows = [] for name, (*_, tpu_rel, gpu_rel) in TPU_APPS.items(): rows.append([name, f"{gpu_rel:.1f}", f"{tpu_rel:.1f}", f"{TABLE6_RATIO[name]:.1f}"]) rows.append(["几何平均", "1.1", "14.5", "13.2"]) rows.append(["按部署比例加权", "1.9", "29.2", "15.3"]) print("\n相对 Haswell 的推理性能（每个 die，含主机开销；论文表 6）：") print_table(["应用", "K80 / CPU", "TPU / CPU", "TPU / K80"], rows) print_observations([ "响应时间才是推理的约束：CPU 和 GPU 要满足 7 ms，只能用 16 的批，吞吐只有它们自己最高值的 37%–42%；" "TPU 没有缓存、乱序执行、多线程这些\"提高平均、不管尾部\"的机制，执行时间是确定的，能用 200 的批、跑到最高吞吐的 80%。", "所以 K80 的峰值是 Haswell 的两倍多，实际推理只快 1.1 倍（几何平均）。论文的结论：GPU 当时的设计（高吞吐、靠大量线程掩盖延迟）" "不适合有严格延迟要求的推理——这一点后来的 GPU（Volta 引入 Tensor Core）也在改。", "TPU 平均快 14.5 倍（几何平均）到 29.2 倍（按部署比例加权）；按功耗算，整机性能/瓦是 CPU 的 17–34 倍、GPU 的 14–16 倍。", ]) # ----------------------------------------------------------------------------- def what_if(name, mem_x=1.0, clock_x=1.0): """ 两种估计"改设计以后快多少"的办法： roofline：只看上界：min(峰值 × clock_x, 2 × 强度 × 带宽 × mem_x)，按"实测 / 上界"的比例不变来缩放； 计数器（Amdahl）：用论文表 3 的周期分解：阵列在算 + 其他（非矩阵）的周期随时钟缩短，等权重 + 移权重的周期随带宽缩短。 【返回】(roofline 估计的加速比, 计数器估计的加速比) """ _, ops_b, *_ = TPU_APPS[name] b0 = roofline_tops(TPU_V1, ops_b) b1 = min(TPU_V1.peak_tops * clock_x, 2 * ops_b * TPU_V1.mem_gbs * mem_x * 1e9 / 1e12) act, stall, shift, other = TABLE3[name] t1 = (act + other) / clock_x + (stall + shift) / mem_x return b1 / b0, (act + stall + shift + other) / t1 def demo4_design_space(): print_header("演示 4：如果改设计（论文图 11、§7）——按部署比例加权的平均加速比") shares = {n: v[4] for n, v in TPU_APPS.items()} tot = sum(shares.values()) rows = [] for label, mx, cx, paper in (("内存带宽 ×4", 4, 1, "约 3×"), ("时钟 ×4", 1, 4, "几乎不变（只有 CNN 约 2×）"), ("换成 GDDR5（带宽 ×5.3，即 K80 的 180 GB/s）", 180 / 34, 1, "3.9×（不含主机开销）"), ("GDDR5 + 时钟 1.5×", 180 / 34, 1.5, "和只换 GDDR5 差不多")): rf = sum(shares[n] * what_if(n, mx, cx)[0] for n in TPU_APPS) / tot ct = sum(shares[n] * what_if(n, mx, cx)[1] for n in TPU_APPS) / tot per = " / ".join(f"{what_if(n, mx, cx)[1]:.1f}" for n in TPU_APPS) rows.append([label, f"{rf:.2f}×", f"{ct:.2f}×", paper, per]) print_table(["改动", "roofline 估计", "计数器（Amdahl）估计", "论文的性能模型", "计数器估计：MLP0/MLP1/LSTM0/LSTM1/CNN0/CNN1"], rows) print(f"\nGDDR5 让拐点从 {TPU_V1.ridge():.0f} 移到 {TPU_V1.peak_tops * 1e12 / 2 / 180e9:.0f}（论文：\"from 1350 to 250\"）。") print_observations([ "内存带宽是最大的杠杆：90% 的负载（MLP、LSTM）受限于内存，带宽 ×4 时 roofline 估计快近 4 倍，" "按性能计数器的周期分解（只有等权重、移权重的周期变短）估计快约 2 倍；论文用更细的性能模型得到约 3 倍，在两者之间。", "时钟 ×4 对 MLP、LSTM 帮助很小（roofline 说完全没用，计数器估计 1.2–1.5 倍——它们大部分时间在等内存），对 CNN0 有用。" "\"GDDR5 + 时钟 1.5×\"一行：论文的模型说加权平均不再提高，我们的计数器估计却从 2.2 倍涨到 2.9 倍" "（它把\"其他\"周期全算成随时钟缩短）——两个粗模型在这里和论文不一致，说明它们都只能看趋势。", "所以论文说：如果当初用 K80 的 GDDR5 内存，TPU 还能再快约 3 倍（TPU'）。后来的 TPU（v2 起）都用了 HBM。" "这篇 2017 年的论文已经预告了今天 LLM 推理的核心矛盾：算力涨得快，带宽跟不上（第 1 讲 §5.3 的\"内存墙\"）。", ]) if __name__ == "__main__": demo1_chips() demo2_apps_on_tpu() demo3_vs_cpu_gpu() demo4_design_space()