人工智能时代已经到来，必须认识到 AI 驱动软件的成本结构与传统软件差异巨大。芯片微架构和系统架构在创新新形式软件的开发和可扩展性中扮演着关键角色。

与早期软件相比，AI 软件运行的基础硬件设施对资本支出和运营支出，进而对毛利率的影响更为显著，而开发成本相对较低。因此，更加有必要投入大量精力优化 AI 基础设施，

以部署 AI 软件。在基础设施方面具有优势的企业，在部署和扩展 AI 应用方面也将占据优势。

1. 谷歌TPU芯片发展简史
谷歌早在 2006 年就推广了构建 AI 专用基础设施的想法，但问题在 2013 年达到了沸点。他们意识到，如果想要大规模部署 AI，必须将数据中心数量翻倍。
因此他们开始为他们的 TPU 芯片打下基础，这些芯片于 2016 年投入生产。将这一点与亚马逊进行比较很有趣，同年，亚马逊也意识到他们需要构建定制硅片。
2013 年，他们启动了 Nitro 计划，该计划专注于开发硅片以优化通用 CPU 计算和存储。两家非常不同的公司为不同计算时代和软件范式的基础设施优化了他们的努力。
自 2016 年以来，谷歌已研发了 6 种不同的 AI 专用芯片，包括 TPU、TPUv2、TPUv3、TPUv4i、TPUv4 和 TPUv5。谷歌主导设计了这些芯片，
Broadcom 则提供了不同程度的中间层和后端协作。这些芯片均由台积电代工制造。自 TPUv2 开始，这些芯片还采用了三星和 SK 海力士提供的 HBM 
内存。尽管谷歌的芯片架构颇具特色，也是本报告后续将深入探讨的内容，但还有一个更为重要的话题值得关注。谷歌在以低成本、高性能可靠大规模部署 AI 方面具有近乎无与伦比的能力。
由于谷歌从微架构到系统架构的整体化方法，其在 AI 工作负载方面相对于微软和亚马逊具有性能/总拥有成本（perf/TCO）优势。将生成式 AI 商业化给企业和消费者是另一个话题。
技术领域是一场永无止境的军备竞赛，人工智能是发展最快的战场。经过训练和部署的模型架构随着时间的推移发生了显著变化。以谷歌内部数据为例。2016 年至 2019 年，
CNN 模型迅速崛起，但随后又再次下降。CNN 的计算特性、内存访问、网络等与 DLRMs、Transformers 和 RNNs 有着非常不同的特点。RNNs 也经历了同样的情况，它们被 Transformer 完全取代。
因此，硬件必须灵活适应行业的发展并支持它们。底层硬件不能过度专注于任何特定的模型架构，否则随着模型架构的变化，它将面临过时的风险。
芯片开发到大规模量产部署通常需要 4 年时间，因此硬件可能会落后于软件想要在其上执行的任务。这一点已经在某些初创公司使用的 AI 加速器架构中显现出来，
这些架构以特定模型类型为优化目标。这是导致大多数 AI 硬件初创公司失败/将会失败的原因之一。
这一点在谷歌自家的 TPUv4i 芯片上尤为明显，该芯片是为推理设计的，却无法在谷歌最佳模型如 PaLM 上运行推理。
上一代谷歌 TPUv4 和 Nvidia A100 显然并非为大型语言模型而设计。同样，最近部署的谷歌 TPUv5 和 Nvidia H100 也并非为 AI 壁垒所设计，
也不是为了应对该壁垒而开发的新模型架构策略。这些策略是 GPT-4 模型架构的核心部分。
此外，芯片微架构只是人工智能基础设施真实成本的一小部分。系统级架构和部署灵活性才是更为重要的因素。
2. 谷歌的系统基础设施优势
谷歌在基础设施方面的优势之一在于他们始终从系统层面设计 TPU。这意味着单个芯片很重要，但在现实世界中如何将其在系统中使用更为重要。
因此，我们的分析将从系统架构到部署使用再到芯片层面逐层展开。
虽然英伟达也从系统角度思考，但他们的系统规模比谷歌的小而窄。此外，直到最近，英伟达在云部署方面没有经验。
谷歌在 AI 基础设施方面的最大创新之一是在 TPU 之间使用自定义网络堆栈 ICI。这条链路相对于昂贵的以太网和 InfiniBand 部署具有低延迟和高性能。它更类似于英伟达的 NVLink。
谷歌的 TPUv2 可以扩展到 256 个 TPU 芯片，与英伟达当前一代的 H100 GPU 数量相同。他们通过 TPUv3 将这一数字增加到 1024，通过 TPUv4 增加到 4096。根据趋势线，
我们预计当前一代的 TPUv5 可以扩展到 16384 个芯片，而无需经过基于以太网的低效方式。虽然从大规模模型训练的性能角度来看，这一点很重要，但更重要的是他们能够将其分配用于实际用途。
谷歌的 TPUv4 系统每台服务器配备 8 块 TPUv4 芯片和 2 个 CPU。这种配置与英伟达的 GPU 相同，英伟达的服务器包含 8 块 A100 或 H100 芯片和 2 个 CPU。
单个服务器通常是 GPU 部署的计算单元，但对于 TPU，部署单元是一个更大的“切片”，包含 64 块 TPU 芯片和 16 个 CPU。
这 64 块芯片通过直接连接的铜缆，在内部以 4^3 立方体形式通过 ICI 网络连接。
在这个由 64 块芯片组成的单元之外，通信将转移到光域。这些光收发器比无源铜缆贵 10 倍以上，因此谷歌优化了它们的切片大小为 64 这个数字，以从网络的角度最小化系统层面的成本。
3. 通过拓扑结构最小化网络成本
尽管谷歌大力推广这一观点，但必须认识到英伟达和英伟达网络的拓扑结构完全不同。英伟达系统部署的是“Clos 网络”，这种网络是“非阻塞”的。这意味着它们可以在没有
任何冲突或阻塞的情况下，同时建立所有输入和输出对之间的全带宽连接。这种设计为在数据中心连接大量设备提供了一种可扩展的方法，最小化了延迟，并增加了冗余。
谷歌的 TPU 网络则放弃了这一点。他们使用 3D 环面拓扑结构，将节点连接在一个三维网格状结构中。每个节点在网格中与其六个相邻节点相连（上、下、左、右、前、后），
在三个维度（X、Y 和 Z）中形成闭合环路。这创造了一个高度互联的结构，其中节点在所有三个维度中形成连续的环路。
第一张图更合理，但如果你仔细想想，又有点饿了，这个网络拓扑字面意思就是一个甜甜圈！
环状拓扑相对于英伟达使用的 Clos 拓扑有几个优势：

另一方面，3D 环形网络也有许多缺点。

总的来说，虽然 Clos 具有优势，但 Google 的 OCS 缓解了许多这些问题。OCS 能够实现多个切片和多个 Pod 之间的简单扩展。
3D 环形拓扑面临的最大问题是错误可能成为更严重的问题。错误可能会出现，而且确实会出现。即使主机可用性达到 99%，2,048 个 TPU 的滑动也会接近 0 的运行能力。即使在 99.9%的情况下，没有 Google 的 OCS，使用 2,000 个 TPU 的训练运行也会有 50%的好输出率。
4. OCS 的美丽之处在于它能够动态重新配置路由。
需要备用资源来允许在部分节点失败的情况下调度任务。操作员无法在不冒失败风险的情况下，从 4k 节点的 Pod 中调度两个 2k 节点的切片。
基于 Nvidia的训练运行通常需要大量额外的开销用于检查点、拉取失败节点并重启它们。谷歌通过简单地绕过失败节点在一定程度上简化了这一过程。
OCS 的另一个好处是，切片可以在部署后立即使用，而无需等待整个网络。
5. 部署基础设施——用户视角
从成本和功耗角度来看，基础设施效率的提升很棒，这让谷歌每美元能部署更多 TPU，而其他公司每美元只能部署 GPU，但这对于用户来说毫无意义。
谷歌内部用户能体验到的最大优势之一是，他们可以根据自己的模型定制基础设施需求。
没有任何芯片或系统能够匹配所有用户想要的内存、网络和计算类型配置。芯片必须具备通用性，但与此同时，用户希望获得灵活性，他们不希望使用一刀切的解决方案。
英伟达通过提供多种不同的 SKU 变体来解决这一问题。此外，他们还提供不同内存容量层级，以及更紧密的集成选项，如 Grace + Hopper 和用于 SuperPods 的 NVLink 网络。
谷歌无法负担这种奢侈。每个额外的 SKU 意味着每个 SKU 的总体部署量会降低。这反过来又降低了他们在基础设施上的利用率。
更多的 SKU 也意味着用户在需要的时候更难获得他们想要的计算类型，因为某些选项不可避免地会被过度分配。这些用户随后将被迫使用次优配置。
因此，谷歌面临着一个难题，既要满足研究人员对特定产品的需求，又要最小化 SKU 的多样性。与英伟达必须为他们的更大、更多样化的客户群支持数百种不同规模部署和 SKU 不同，
谷歌正好有 1 种 TPUv4 部署配置，包含 4,096 个 TPU。尽管如此，谷歌仍然能够以独特的方式切割和分配这些资源，使内部用户能够获得他们所期望的基础设施灵活性。
谷歌的 OCS 还支持创建自定义网络拓扑，例如扭曲的环形网络。这些是三维环形网络，其中某些维度被扭曲，这意味着网络边缘的节点以一种非平凡、非线性的方式连接，
从而在节点之间创建额外的捷径。这进一步提高了网络直径、负载均衡和性能。
谷歌团队充分利用这一优势来协助特定模型架构。以下是 2022 年 11 月某一天中各种 TPU 配置的芯片数量和网络拓扑的流行度快照。
尽管许多配置在系统中的芯片数量相同，但仍有超过 30 种不同的配置，以适应正在开发的多种模型架构。这是谷歌在使用 TPU 和灵活性方面的巨大洞察力。
此外他们还有许多使用率较低且未在图中展示的拓扑结构。
要充分利用可用的带宽，用户将数据并行映射到 3D 环面的一维上，并在其他维度上映射两个模型并行参数。谷歌声称最佳拓扑选择能够使性能提升 1.2 倍至 2.3 倍。
6. 规模最大的 AI 模型架构：DLRM
DLRM 旨在通过建模分类和数值特征来学习用户-物品交互的有意义表示。该架构由两个主要组件组成：嵌入组件（处理分类特征）和多层感知器（MLP）组件（处理数值特征）。
用最简单的话来说，多层感知器组件是密集的。特征被输入到一系列全连接层中。这类似于旧的、GPT 4 之前的 Transformer 架构，它们也是密集的。
密集层与硬件上的大规模矩阵乘法单元映射得非常好。
嵌入组件是 DLRMs 高度独特的部分，也是使其计算特征如此独特的原因。DLRM 的输入是表示为离散、稀疏向量的分类特征。简单的谷歌搜索仅包含整个语言中的几个词。
由于这些稀疏输入本质上更类似于哈希表而非张量，因此它们与硬件中发现的巨大矩阵乘法单元映射不佳。由于神经网络通常在密集向量上表现更好，因此采用嵌入将分类特征转换为密集向量。

密集向量：
嵌入函数将分类空间（英语中的单词、与社交媒体帖子的互动、对某种类型帖子的行为）映射到更小、密集的空间（代表每个单词的 100 维向量）。这些函数使用查找表实现，
查找表是 DLRMs 的重要组成部分，并且通常是 DLRM 模型的第一层。嵌入表的大小差异很大，从几十兆字节到几百吉字节甚至太字节不等。
Meta 的 2 岁 DLRM 模型参数超过 12 万亿，运行推理需要 128 个 GPU。如今，最大的生产 DLRM 模型至少要大数倍，仅模型嵌入就需要超过 30TB 的内存。
预计明年嵌入量将增加到超过 70TB！因此，这些表需要跨多芯片的内存进行分区。主要有三种分区方法：列分片、行分片和表分片。
DLRM 的性能主要受内存带宽、内存容量、向量处理性能以及芯片间网络/互连的限制。嵌入查找操作主要包含小规模的聚合或分散内存访问，其算术强度较低（FLOPS 完全不重要）。
对嵌入表的访问本质上是无结构的稀疏性。每个查询都必须从分布在数百或数千个芯片上的 30TB+嵌入数据中获取数据。这可能导致超级计算机在 DLRM 推理时计算、内存和通信负载的不平衡。
对于 MLP 和类似 GPT-3 的 Transformer 中的密集操作，这一点差异很大。芯片 FLOPS/sec 仍然是主要性能驱动因素之一。当然，除了 FLOPs 之外，
还有多种因素制约着性能，但 GPU 在 Chinchilla 风格的 LLMs 中仍能实现超过 71%的硬件 FLOPs 利用率。
7. 谷歌的 TPU 架构
谷歌的 TPU 在架构上引入了一些关键创新，使其区别于其他处理器。与传统处理器不同，TPU v4 没有专门的指令缓存。相反，它采用了一种直接内存访问（DMA）机制，类似于 Cell 处理器。
TPU v4 中的向量缓存不是标准缓存层次结构的一部分，而是被用作临时存储区。临时存储区与标准缓存的不同之处在于它们需要手动写入，而标准缓存则自动处理数据。
TPU v4 配备了160MB 的 SRAM 用于scratchpad，并且拥有 2 个 TensorCore，每个 TensorCore 包含 1 个 Vector Unit，
该 Vector Unit 有 4 个 Matrix Multiply Units (MXUs)和 16MB 的 Vector Memory (VMEM)。这两个 TensorCore 共享 128MB 的内存。它们支持 275 TFLOPS 的 BF16，
同时也支持 INT8 数据类型。TPU v4 的内存带宽为 1200GB/s。芯片间互连（ICI）通过六个 50GB/s 的链路提供 300GB/s 的数据传输速率。
A 322b 非常长指令字（VLIW）标量计算单元包含在 TPU v4 中。在 VLIW 架构中，指令被组合成一个长指令字，然后被派发给处理器执行。这些组合的指令，也称为包，
在程序编译期间由编译器显式定义。VLIW 包包含多达 2 条标量指令、2 条向量 ALU 指令、1 条向量加载和 1 条向量存储指令，以及 2 个用于向 MXU 传输数据的槽位。
向量处理单元（VPU）配备了 32 个 2D 寄存器，包含 128x 8 32b 元素，使其成为一个 2D 向量 ALU。矩阵乘法单元（MXU）在 v2、v3 和 v4 上是 128x128，而 v1 版本具有 256x256 的配置。
这种变化的原因是谷歌模拟发现四个 128x128 的 MXU 利用率比一个 256x256 的 MXU 高 60%，但四个 128x128 的 MXU 占用的面积与 256x256 的 MXU 相同。
MXU 输入使用 16b 浮点（FP）输入，并以 32b 浮点（FP）进行累积。
这些较大的单元允许更高效的数据重用以突破内存墙。
8. 谷歌 DLRM 优化
谷歌是率先大规模在其搜索产品中使用 DLRM 的公司之一。这种独特的需求催生了一种非常独特的解决方案。上述架构存在一个主要缺陷，即无法有效处理 DLRM 的嵌入。
谷歌的主 TensorCore 非常大，与这些嵌入的计算特征不匹配。谷歌不得不在其 TPU 中开发一种全新的“SparseCore”，这与上面描述的用于密集层的“TensorCore”不同。
SparseCore（SC）为 Google TPU 中的嵌入提供了硬件支持。从 TPU v2 开始，这些特定领域的处理器就直接与每个 HBM 通道/子通道相连。
它们加速了训练深度学习推荐模型（DLRM）时最内存带宽密集的部分，而仅占芯片面积和功耗的约 5%。通过在每个 TPU v4 芯片上使用快速 HBM2 进行嵌入，
而不是使用 CPU，Google 将其内部生产 DLRM 的速度提高了 7 倍，与将嵌入留在主机 CPU 的主内存（TPU v4 SparseCore 与 TPU v4 Skylake-SP 上的嵌入）相比。
SparseCore 能够实现从 HBM 的快速内存访问，配备专门的获取、处理和刷新单元，将数据移动到稀疏向量内存（Spmem）的银行，
并由可编程的 8 位宽 SIMD 向量处理单元（scVPU）进行更新。这些单元的 16 个计算瓦片组成了一个 SparseCore。
额外的跨通道单元执行特定的嵌入操作（DMA、排序、稀疏归约、分支、连接）。每个 TPU v4 芯片有 4 个 SparseCore，每个 SparseCore 配备 2.5MB 的 Spmem。
展望未来，我们推测 SparseCore 的数量将继续增加到 6 个用于 TPUv5，而瓦片数量将增加到 32 个，这是由于 HBM3 上子通道数量的增加。
虽然迁移到 HBM 能带来巨大的性能提升，但性能扩展仍然受限于互连二分带宽。TPU v4 中 ICI 的新 3D 环形结构进一步提升了嵌入查找性能。
然而，当扩展到 1024 个芯片时，性能提升会下降，因为 SparseCore 的开销成为瓶颈。
这个瓶颈很可能导致每块板的 Spmem 也随着 TPUv5 的增加而增加，如果谷歌觉得他们的 DLRMs 需要超过~512 片的大小和容量。