最近学习了论文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》。

谷歌的TPU v4（张量处理单元第四代）是专门为机器学习（ML）尤其是深度学习模型训练设计的加速器，主要在大规模计算中应用，如大规模语言模型（LLM）、深度学习推荐模型（DLRM）以及其他需要高吞吐量、低延迟的计算任务。TPU v4针对前代TPU的不足，结合了新的硬件架构和先进的互联技术，显著提高了性能、可扩展性和可靠性。

1. TPU v4的三大核心创新

TPU v4引入了三大创新特性，以应对大规模机器学习计算中的关键挑战：

(1) 光学电路交换（OCS）与光纤数据链路

TPU v4采用了光学电路交换（Optical Circuit Switches，OCS）技术，用光纤连接来实现超大规模系统的高效通信。通过这种光学交换方式，TPU v4能够构建一个最大支持4000个TPU节点的超级计算机，同时保持良好的可扩展性和可靠性。与传统电子包交换（如InfiniBand）相比，光学交换不仅能提高系统的吞吐量，还能减少功耗，且由于其高带宽和低延迟的特性，能有效支撑深度神经网络（DNN）训练中的全对全通信模式（All-to-All communication），这对于处理大规模数据特别重要。

(2) SparseCore硬件支持

SparseCore是TPU v4的一个新特性，专门优化了稀疏嵌入（embedding）操作，特别是对于深度学习推荐模型（DLRM）等需要大量稀疏矩阵计算的任务。SparseCore提供了一种数据流海量核心架构，允许将嵌入数据灵活地分配到TPU v4超级计算机的128TiB物理内存中，这种架构极大提升了稀疏计算的效率，能够比传统方法快5到7倍。

(3) 灵活的拓扑重配置

TPU v4在硬件拓扑上引入了光学交换的灵活性，能够根据实际应用需求调整拓扑结构（如3D环形拓扑或扭曲环形拓扑），优化互连带宽的使用。这种拓扑重配置能力使得TPU v4能根据不同计算任务的需求，选择最适合的拓扑结构，从而大幅提高了性能并减少了通信瓶颈。

2. TPU v4的关键优势

(1) 可扩展性和可靠性

通过光学电路交换（OCS），TPU v4能够支持超过4000个TPU节点的超级计算机规模，而系统能容忍1K个CPU主机不可用的情况。传统的分布式系统通常在硬件故障时需要复杂的恢复机制，但TPU v4通过OCS提高了系统的可靠性，使得训练过程更加稳定。

(2) 性能与能效的提升

TPU v4相较于TPU v3，提升了多个方面的性能，尤其在大规模语言模型（LLM）的训练中，使用3K个TPU v4切片和3D环形拓扑显著减少了训练时间。OCS的光学交换技术相较于传统的电子包交换，在减少功耗的同时，还优化了性能，因此TPU v4在能效方面优于同类的其他加速器，例如Nvidia的A100和Graphcore的MK2 IPU。

(3) 稀疏矩阵加速

SparseCore提供的稀疏计算加速功能是TPU v4的另一个亮点。这种硬件支持将稀疏矩阵操作效率提高了5到7倍，对于大规模推荐系统和其他稀疏计算密集型任务，具有极大的优势。并且，这种加速仅占TPU v4总芯片面积和功耗的5%左右，成本效益非常高。

(4) 模块化设计和拓扑灵活性

TPU v4的设计支持模块化扩展，可以通过灵活配置拓扑来优化性能。用户可以根据具体需求，选择不同的拓扑结构（如3D环形、扭曲环形等），这样不仅提升了互连带宽，还能够减少系统在不同应用场景下的性能瓶颈。

3. TPU v4与前代产品及其他竞争产品的比较

与TPU v3相比，TPU v4在多个方面的性能提升使得其成为训练大规模模型（如PaLM模型）的理想选择。PaLM是一个拥有540亿参数的大规模语言模型，在TPU v4超级计算机上训练时，能够保持接近硬件浮点性能的57.8%。同时，TPU v4在功耗、性价比和计算密度方面，较Nvidia A100和Graphcore MK2 IPU等同类产品具备明显优势。

4. TPU v4的能效和环境影响

TPU v4的另一大优势是其能效。在集成到高效能的仓库级计算机（WSC）中时，TPU v4能够显著减少碳足迹。与传统的深度学习加速器相比，TPU v4在相似规模的数据中心中，能够消耗20倍更少的能源，且二氧化碳排放量降低约95%。这一点对于支持可持续的人工智能发展具有重要意义，尤其是在大规模机器学习模型训练过程中。

5. 总结

TPU v4代表了谷歌在深度学习硬件加速领域的最新突破。通过光学电路交换技术、SparseCore加速和灵活的拓扑配置，TPU v4不仅提高了大规模机器学习模型训练的性能，还在可靠性、能效和可扩展性方面做出了显著改进。TPU v4的设计满足了当今AI领域对于更强计算力、低功耗和高可靠性的需求，特别是在训练大规模语言模型（LLM）和深度学习推荐模型（DLRM）等任务时，展现出无与伦比的优势。

更多「老虎说芯」精彩内容，欢迎点击「阅读原文」

谷歌TPU v4处理器论文学习笔记