TPU 不出售，但为什么？

原文towardsdatascience.com/tpus-are-not-for-sale-but-why-5964f87f7a15?sourcecollection_archive---------4-----------------------#2024-04-30观点Google 在 AI 硬件方面的独特做法分析https://haifeng-jin.medium.com/?sourcepost_page---byline--5964f87f7a15--------------------------------https://towardsdatascience.com/?sourcepost_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 UnsplashNvidia 的股价飙升因为其 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 核心非常不同。所有的动态性和通用性都被去除了。它有一个网格状的节点每个节点仅按预定义的方式进行乘法和加法。结果会直接传送到下一个节点进行下一个乘法和加法。因此一切都是预定义和固定的。通过这种方式我们节省了时间因为不需要进行指令解码它只会进行接收到的乘法和加法。没有用于读写的寄存器因为我们已经知道结果应该去哪里也无需为随后的任意操作存储结果。除了 MXUTPU 还被设计为更具可扩展性。它具有用于高带宽芯片间互联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 硬件将会更加优化以适应特定的应用比如变压器模型。无论是硬件还是模型都难以轻易跳出这种模式崩溃的局面。