TPU 不出售，但为什么？

原文：towardsdatascience.com/tpus-are-not-for-sale-but-why-5964f87f7a15?source=collection_archive---------4-----------------------#2024-04-30

观点

Google 在 AI 硬件方面的独特做法分析

https://haifeng-jin.medium.com/?source=post_page---byline--5964f87f7a15--------------------------------https://towardsdatascience.com/?source=post_page---byline--5964f87f7a15-------------------------------- Haifeng Jin

·发表于 Towards Data Science ·10 分钟阅读·2024 年 4 月 30 日

–

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/26ebc18d628a6a347ee233a287cca807.png

图片来源：Dollar Gill via Unsplash

Nvidia 的股价飙升，因为其 GPU 在 AI 硬件市场的主导地位。然而，与此同时，Google 的 TPU 作为知名的 AI 硬件，并不出售。你只能在 Google Cloud 上租用虚拟机来使用它们。那么，为什么 Google 不加入销售 AI 硬件的游戏呢？

免责声明：本文所表达的观点仅代表作者个人意见，并不一定反映 Google 或其附属公司的意见或观点。本文中所提供的所有信息完全来源于公开材料。

一个流行的理论

我听到的一个流行理论是，Google 希望吸引更多客户使用其云服务。如果他们将 TPU 卖给其他云服务提供商，他们在云服务市场中的竞争力将减弱。

根据云服务客户的说法，这一理论并不十分合理。没有企业级客户愿意被锁定在某一特定的云服务提供商。他们希望具有足够的灵活性，能够在需要时迁移到其他提供商。否则，如果服务提供商提高价格，他们将无能为力。

如果 TPU 只能在 Google Cloud 上使用，许多企业客户宁愿不使用它。这也是为什么很多客户不愿意使用 TPU 的原因。直到最近，当 OpenXLA（一个中介软件）支持更多框架，如 PyTorch 时，他们才开始感觉到不那么受限。

所以，使用 TPU 吸引客户到 Google Cloud 并不是不出售它们的合理理由。那么，真正的原因是什么呢？为了更好地回答这个问题，我们必须探讨 Google 是如何启动 TPU 项目的。

为什么 Google 启动了 TPU 项目？

简短的回答是为了专有用途。曾经，GPU 无法满足 AI 硬件的计算需求。

让我们试着估算 TPU 项目大约什么时候开始的。考虑到它在2016 年首次对外公布，我们可以合理推测它大约在 2011 年左右开始。如果这个推测正确，那么他们的项目启动得相当早，因为直到 2012 年 AlexNet 的出现，计算机视觉才取得了显著的进展。

通过这个时间线，我们知道在项目启动时，GPU 的性能不如今天。Google 很早就看到了 AI 革命，并希望为大规模计算提供更快的硬件。他们唯一的选择是为此构建一个新的解决方案。

这就是 Google 启动该项目的原因，但还有更多的问题。为什么当时 GPU 不够好？Google 看到了哪些潜在的改进，足够重要以至于启动了他们的新硬件项目？

答案在于 GPU 和 TPU 的微架构。让我们来看看 GPU 和 TPU 核心的设计。

GPU 的设计理念

首先，让我们快速回顾一下 CPU 的背景知识。当一个指令到来时，它会被指令解码器解码，并与寄存器中的数据一起送入算术逻辑单元（ALU）。ALU 负责所有计算并将结果返回到其中一个寄存器。如果 CPU 有多个核心，它们可以并行工作。

什么是 GPU？它是图形处理单元（Graphics Processing Unit）的缩写。它最初是为图形计算设计的，后来发现它适用于机器学习。GPU 中的大多数操作是矩阵运算，可以并行执行。这也意味着，与 CPU 相比，GPU 需要支持的操作较少。

芯片针对特定任务越专业，完成任务的速度就越快。

GPU 最初设计的核心理念是将 CPU 的功能进行简化，采用更小但更多的核心，以实现更快的并行计算。GPU 上支持的指令数量远少于 CPU，这使得单个核心在芯片上所占的面积更小。这样，它们可以将更多的核心集成到芯片上，进行大规模并行计算。

为什么特性越少意味着芯片上占用的面积越小？在软件中，更多的特性意味着更多的代码；在硬件中，所有特性都是通过逻辑电路实现的，而不是代码。更多特性意味着电路更复杂。例如，CPU 必须在芯片上实现更多的指令。

更小也意味着更快。更简单的逻辑门设计导致了更短的周期时间。

TPU 的设计理念

TPU 进一步发展了这种专门用于深度学习的芯片理念。TPU 的定义性特征是其矩阵乘法单元（MXU）。由于矩阵乘法是深度学习中最常见的操作，TPU 为此构建了一个专门的核心，即 MXU。

这比 GPU 核心更为专门化，能够执行多种矩阵操作，而 MXU 只做一件事：矩阵乘法。

它的工作方式与传统的 CPU/GPU 核心非常不同。所有的动态性和通用性都被去除了。它有一个网格状的节点，每个节点仅按预定义的方式进行乘法和加法。结果会直接传送到下一个节点，进行下一个乘法和加法。因此，一切都是预定义和固定的。

通过这种方式，我们节省了时间，因为不需要进行指令解码，它只会进行接收到的乘法和加法。没有用于读写的寄存器，因为我们已经知道结果应该去哪里，也无需为随后的任意操作存储结果。

除了 MXU，TPU 还被设计为更具可扩展性。它具有用于高带宽芯片间互联（ICI）的专用端口。它被设计为能够安装在 Google 数据中心的机架上并用于集群。由于它仅供专有使用，因此不需要担心销售单个芯片或安装芯片到机架上的复杂性。

TPUs 今天仍然更快吗？

其他人没有想到构建专用的张量操作核心（矩阵乘法）的简单想法是没有道理的。即使他们没有想到这一点，后来也不应该不去模仿。

从时间线来看，似乎 Nvidia 在大约相同的时间提出了相同的想法。Nvidia 的类似产品——张量核心（Tensor Cores）——首次公开宣布是在 2017 年，也就是 Google 发布 TPU 的一年后。

目前尚不清楚 TPU 是否仍然比 GPU 更快。我无法找到最新一代 TPU 和 GPU 的公开基准测试，而且我不清楚应该使用哪个代数和度量标准来进行基准测试。

然而，我们可以使用一个通用的面向应用的度量标准：每轮的花费（dollars per epoch）。我发现一个有趣的基准测试来自 Google Cloud，它将不同的硬件对齐到相同的轴心：金钱。如果你使用相同的模型、数据和轮次，在 Google Cloud 上 TPU 看起来更便宜。

大型模型，如 Midjourney、Claude 和 Gemini，都对训练成本非常敏感，因为它们消耗大量计算资源。因此，它们中的许多都在 Google Cloud 上使用 TPU。

为什么 TPU 更便宜？

一个重要的原因是软件堆栈。你不仅使用硬件，还使用与之相关的软件堆栈。Google 在其软件堆栈和 AI 硬件的垂直集成上，比 GPU 做得更好。

Google 拥有专门的工程团队来为其构建完整的软件堆栈，并实现强大的垂直集成，从模型实现（Vertex Model Garden）到深度学习框架（Keras、JAX 和 TensorFlow），再到为 TPU 优化的编译器（XLA）。

GPUs 的软件堆栈非常不同。PyTorch 是与 Nvidia GPU 一起使用的最流行的深度学习框架，主要由 Meta 开发。与 PyTorch 一起使用的最广泛应用的模型库是 HuggingFace 开发的transformers和diffusers。在这些公司之间实现完美的垂直集成的软件堆栈要困难得多。

一个警告是，使用 JAX 和 TensorFlow 实现的模型较少。有时，你可能需要自己实现模型，或者在 TPU 上使用 PyTorch。根据实现的不同，在 TPU 上使用 PyTorch 时，可能会遇到一些摩擦。因此，除了硬件成本外，可能还会有额外的工程成本。

为什么不开始直接销售 TPU 呢？

我们了解到，该项目最初是为专有用途启动的，并由于其较低的价格，在 Google Cloud 上获得了相当不错的用户基础。为什么 Google 不直接开始将其卖给客户，就像 Nvidia 的 GPU 一样？

简短的回答是保持专注。Google 与 OpenAI 在生成式 AI 领域的竞争非常激烈。同时，Google 正处于多轮科技裁员中，以降低成本。当前的明智策略是将有限的资源集中在最重要的项目上。

如果 Google 想要开始销售 TPU，它将同时与两个强大的对手 Nvidia 和 OpenAI 竞争，这在目前可能不是一个明智的举动。

销售硬件的巨大开销

直接向客户销售硬件为公司带来巨大的开销。相反，租赁云服务上的 TPU 要更容易管理得多。

当 TPU 仅在云端提供时，他们可以集中安装所有 TPU 和相关软件。无需处理各种安装环境或部署 TPU 集群的难题。

他们准确知道需要生产多少 TPU。需求完全是内部需求，因此没有不确定性。因此，供应链管理要容易得多。

销售也变得更加简便，因为这仅仅是在销售云服务。无需组建一个有经验销售硬件的新团队。

TPU 方法的优势

在没有直接向客户销售硬件的所有额外开销下，Google 获得了几项优势。

首先，他们可以采用更具攻击性的 TPU 架构设计。TPU 具有一种独特的芯片连接方式。与多个 GPU 连接到同一主板不同，TPU 被组织成立方体的形式。他们将 64 个 TPU 排列成一个 4x4x4 的立方体，将它们彼此连接，以实现更快的芯片间通信。单个 v5p Pod 中有 8960 个芯片。它们可以轻松地一起使用。这是完全控制硬件安装环境的优势。

第二，他们可以更快地进行迭代，推出新一代硬件。由于他们只需支持一小部分专有用途的应用场景，这大大缩短了每一代芯片的研发周期。我想知道英伟达是否早于谷歌想到了 TensorCore 的概念，但由于向外部客户销售硬件的额外开销，它们只能比谷歌晚一年才发布。

从竞争 GenAI 这个最重要目标的角度来看，这些优势使谷歌处于非常有利的地位。最重要的是，通过这种内部硬件解决方案，谷歌避免了以垄断价格购买英伟达 GPU，从而节省了大量资金。

TPU 方法的缺点

到目前为止，我们讨论了谷歌 AI 硬件方法的许多优点，但它是否也有缺点？确实有一个大缺点。谷歌变成了一个技术孤岛。

每一个技术先锋都将在一段时间内变成一个与世界其他地方隔离的孤岛。这是因为他们在相应的基础设施尚未就绪时就开始了。他们需要从头开始建立一切。由于迁移成本，他们会坚持自己的解决方案，即使其他人都在使用其他方案。

这正是谷歌目前所面临的情况。外界正在通过 HuggingFace 和 PyTorch 的模型进行创新，大家迅速互相调整模型，以开发出更好的模型。然而，由于谷歌的基础设施主要围绕 TensorFlow 和 JAX 构建，它不能轻松地加入这个过程。当将外部模型投入生产时，必须重新用谷歌的框架进行实现。

这种“技术孤岛”问题使得谷歌在从外部世界采纳优秀解决方案时变得缓慢，进一步将其与他人隔离开来。谷歌要么开始引入更多的外部解决方案，如 HuggingFace、PyTorch 和 GPU，要么始终确保其内部解决方案在全球范围内是最优秀的。

未来的 AI 硬件会是什么样子？

最后，让我们展望一下 AI 硬件的未来。未来的 AI 硬件会是什么样子？简短的回答是：随着硬件变得更加专业化，出现模式崩溃现象。

硬件将进一步与应用耦合。例如，支持更多精度格式以更好地为语言模型服务。像bfloat16、TF32一样，它们可能会更好地支持int8和int4。Nvidia 宣布了其第二代变压器引擎，该引擎与 Blackwell GPU 兼容。这使得在不改变用户代码的情况下，更容易优化硬件以适应变压器模型。很多共同设计工作正在进行中。

另一方面，软件也难以跳出变压器领域。如果它们做到了，由于缺乏硬件支持，它们的速度将会很慢。相反，它们是在考虑硬件的前提下实施自己的模型。例如，FlashAttention算法就是为了利用 GPU 的内存层次结构来提高性能。

我们预计很快会出现一次大的模式崩溃。当前的硬件和软件在现有模型中彼此优化得非常好。它们都无法轻易脱离现有的设计或算法。如果有一种全新、完全不同于变压器模型的模型，它需要比现有模型好 10 倍才能广泛采用。它必须激励人们像对变压器那样快速且廉价地制造新硬件。

总结

总结来说，TPU 项目最初是为专有用途而启动的，当时 GPU 的计算能力不足。谷歌希望专注于 GenAI，而不是在 AI 硬件市场上竞争，以避免减缓迭代速度并牺牲其创新设计。更快的计算和更低的成本帮助谷歌在 AI 研究和应用开发上取得了显著进展。然而，这也使谷歌成为了一个技术孤岛。

展望未来，AI 硬件将会更加优化，以适应特定的应用，比如变压器模型。无论是硬件还是模型，都难以轻易跳出这种模式崩溃的局面。