Hotdry.
归档
ai-systems

TPU 脉动阵列与 HBM 带宽对比 GPU 张量核：稠密矩阵乘法与动态负载下的能效扩展

剖析 TPU 脉动阵列在稠密矩阵乘法中的 HBM 带宽利用与能效优势，对比 GPU 张量核在动态负载下的灵活性与长期训练热功耗挑战，提供集群参数与监控清单。

在 AI 训练中，Google TPU 的脉动阵列（systolic array）架构针对稠密矩阵乘法（dense matmul）表现出色，其核心在于数据 “流动式” 计算，极大减少内存访问开销，从而在 HBM 带宽利用和功耗控制上领先 NVIDIA GPU 的张量核（Tensor Cores）和流式多处理器（SMs）。这一设计特别适合 Transformer 模型的注意力机制和前馈层，这些层占训练计算的 90% 以上。相比之下，GPU 虽在动态形状工作负载中更灵活，但其通用架构带来的缓存和线程管理 overhead 导致带宽浪费和热密度升高，长周期训练易受电源与散热限制。

TPU 的脉动阵列由网格状处理单元组成，数据仅加载一次权重，即可通过相邻单元逐级传递输入，实现零额外内存读写。以 Ironwood（TPU v7）为例，其单芯片 BF16 算力达 4614 TFLOPS，配备 192GB HBM3e 内存，带宽高达 7370 GB/s，比前代 v5p 的 2765 GB/s 提升近 3 倍。这一带宽直接支撑大规模 matmul，避免 Transformer 注意力层的 “数据饥饿”。证据显示，在 64 芯片集群中，TPU v4 的功耗仅为同规模 A100 GPU 的 52%-77%，每瓦性能高出 2-3 倍。“TPU 采用脉动阵列架构，数据流入阵列后几乎无需额外内存访问，大幅降低能耗。” 谷歌的 OCS（光路交换机）进一步消除光电转换损耗，支持 9216 芯片集群，总算力 42.5 ExaFLOPS。

GPU 的张量核虽优化 matmul，支持混合精度（如 BF16/FP8），但 SMs 的分支预测和多级缓存（L1/L2 + HBM）在动态负载（如变长序列或自定义层）中效率低下。以 Blackwell B200 为例，192GB HBM 带宽约 8TB/s（Ultra 版更高），但实际利用率受非规整访问拖累，仅 60%-70%。在 MLPerf 测试中，A100 集群性能落后 TPU v4 40%，大规模训练中 GPU 能耗高出 30%。动态工作负载下，GPU 的 NVLink（900GB/s/ 链路）互联虽快，但热密度高：H100 单卡 TDP 700W，集群易达 PUE 1.3+，长期训练（>1 周）需复杂液冷。

针对稠密 matmul 主导的 AI 训练，TPU 在 power/heat scaling 上占据长期优势：v7 每瓦性能较 Trillium 提升 2 倍，3D 环形拓扑（ICI 1.2Tbps / 链路）确保低延迟扩展，避免 GPU 的 NVSwitch 瓶颈。实际部署中，TPU Pod（如 v5p 8960 芯片）MFU（Model FLOPs Utilization）稳定 50%-60%，GPU 常降至 30%-40%。

可落地工程参数与清单：

  1. 集群配置阈值：

    • TPU：优先 4x4x4 或倍数 Pod（带 wraparound），最小 256 芯片起规模化；HBM 利用率监控 >85%，低于阈值调大 batch size（推荐 1024-4096）。
    • GPU：DGX H100 8 - 卡节点起，NVLink 拓扑；动态负载下 batch size 512-2048，避免 >16 节点无 InfiniBand。

  2. 带宽优化参数：

    • TPU：XLA 编译启用 --tpu_periodic_copy=off，VMEM 预加载权重（arithmetic intensity >20 FLOPs/byte）；HBM 带宽峰值监控，目标 90% 饱和。
    • GPU：cuBLAS GEMM 调优 CUBLAS_MATH_DISALLOW_REDUCED_PRECISION=0，Tensor Core 启用 TF32；动态 shape 用 Triton 自定义内核，带宽利用 >70%。

  3. 功耗 / 热管理监控点：

    指标 TPU 阈值 GPU 阈值 回滚策略
    TDP / 芯片 <450W <700W 降频 10%
    集群 PUE <1.15 <1.25 加液冷节点
    热密度 (W/cm²) <200 <300 分散 Pod
    能效 (TFLOPS/W) >10 >5 切换低精度

  4. 长周期训练清单：

    • 预热阶段（1-2h）：基准 matmul 测试，验证 systolic vs Tensor Core 利用率。
    • 扩展阶段：TPU 用 3D torus 渐增至 4096 芯片，监控 ICI 延迟 <5μs；GPU 限 1024 卡，NVSwitch 饱和 <80%。
    • 稳态监控：Prometheus + Grafana，每 5min 采样 HBM 带宽、功耗、MFU；异常阈值：带宽降 20%、热升 15°C → 暂停 checkpoint。
    • 回滚策略：TPU 降级 v5p Pod，GPU 切换 A100 备用；成本计算：TPU TCO 低 40%，但生态锁定风险高。

在动态负载占比 <20% 的稠密训练场景，TPU 的 systolic + HBM 组合确保 dominance：Gemini 2.5 推理延迟降 67%，远超 GPU。转向 TPU 可将月训成本降 30%-50%，但需 JAX/TensorFlow 重构。

资料来源：

  • Hacker News 讨论：TPUs vs. GPUs 长远 AI 竞赛（链接于 uncoveralpha.com）。
  • UncoverAlpha：Google TPU 深度剖析（2025-11-24）。
  • MLPerf 基准与 SemiAnalysis 数据（2025）。

（正文字数：1268）

查看归档