一文搞懂 Google不同型号TPU显存

起因：

      最近在做具身智能相关的工作，看的多是google的文章，里面频繁出现类似 We train our policies on a TPUv3 8x8 pod (64 TPUv3 chips)，网上没有找到很清晰地关于不同型号TPU显存大小的说明，于是自己参考前辈的文章和官方文档总结了一份，有错误希望大家及时指出。

结论：

      每个TPU v3设备通常包括8个核心（cores），每个核心16GB显存，因此一个完整的TPU v3设备共有128GB内存。TPU的计算单元是core，一张物理卡上面有两个或一个Core，而具体的logical Core随着版本变化而不同。在TPU v3中，一个物理Core对应于一个logical Core，而在TPUv4中，两个物理Core对应于一个logical Core。上层框架（例如TensorFlow、PyTorch、JAX等）将每个logical Core视为一个计算设备。

---

常用的计算方式为TPUv3-8/TPUv4-8，v后面表示版本（3或者4），横线后面则表示物理核的个数。

• 每个 v2 TPU 芯片包含两个 TensorCore，总的显存为16GB。
• 每个 v3 TPU 芯片包含两个 TensorCore，每个TensorCore包含16GB显存，所以一个v3 TPU 芯片共32GB显存。对于单个 v3 TPU，使用 --accelerator-type=v3-8（四个芯片，8个TensorCore）。对于具有 128 个 TensorCore 的 v3 Pod 切片，使用 --accelerator-type=v3-128。
• 每个 TPU v4 芯片包含两个 TensorCore，合并成了一个logical Core，因此在整个芯片上具有统一的 32 GiB HBM 内存空间，可让两个片上 TensorCore 之间更好地协调。
• 每个 v5e 芯片包含一个 TensorCore，总的显存为16GB。
• 每个 v5p 芯片包含两个 TensorCore，总的显存为95GB。

所以对于64 TPUv3 chips，对应的显存是 $64\times 2\times 16=2048$GB。

参考：

1. 知乎文章
2. Google 官方文档