过去十年机器学习软件开发行业概览：英伟达 CUDA 垄断地位下降，PyTorch 超越谷歌 TensorFlow

在过去的十年间，机器学习软件开发的格局翻天覆地。流水的框架铁打的英伟达，这些框架大多都极度依赖英伟达的 CUDA（统一计算架构），且在英伟达 GPU 上性能最好。不过，随着 PyTorch 2.0 和 OpenAI 的 Triton 的到来，英伟达因其软件护城河而在这一领域的制霸地位恐将不保。

本篇报告中将涵盖以下主题：为何谷歌的 TensorFlow 输给了 PyTorch，谷歌为何没能利用其在人工智能领域的早期领导地位，机器学习模型训练时间的主要构成成分，内存容量、带宽、成本墙，模型优化，为何其他人工智能硬件公司至今无法撼动英伟达的统治地位，为何硬件地位逐渐重要，英伟达在 CUDA 上的竞争优势是如何消失的，以及英伟达的竞争对手在大型云的芯片训练上所取得的重大胜利。

问题简述是，英伟达闭源 CUDA 将不再涵盖机器学习模型的默认软件栈。英伟达虽然抢占先机，但却让 OpenAI 和 Meta 后来居上掌控了软件栈，英伟达专有工具的失败致使后者在生态系统中建立了自己的工具，英伟达的护城河也遭到了永久地削弱。

TensorFlow vs. PyTorch

仅仅几年前，在相当分散的框架生态系统中，掌握最常用框架 TensorFlow 的谷歌作为领跑者，设计并部署了唯一成功的人工智能指定应用加速器 TPU，他们看似已经抢占先机、准备好制霸机器学习行业了。

2019 年机器学习框架状态，PyTorch 掌控研究领域，Tensorflow 掌控行业领域

可事实却是 PyTorch 赢了，而谷歌没能将先手优势转换为对新生机器学习行业的主导权。如今的谷歌使用自研软硬件栈而非 PyTorch 和对应 GPU，导致其在机器学习社区的地位颇为尴尬，谷歌甚至还有另一套名为 Jax 的框架，直接和 TensorFlow 相竞争。

不仅如此，关于谷歌在搜索及自然语言处理方面的主导地位会因大型语言模型而衰退之类的言论也是甚嚣尘上，尤其是来自那些 OpenAI 及各种利用 OpenAI 的 API 的、或基于类似基础模型的初创公司。虽说可能是有些杞人忧天，但我们今天不谈这些。目前来说，虽然挑战不断，但谷歌仍是处于机器学习模型的最前沿。谷歌所发明的 Transformer 在 PaLM、LaMBDA、Chinchilla、MUM，以及 TPU 等诸多领域依旧是最先进的。

让我们回到 PyTorch 赢得一筹的话题中去。虽然也有从谷歌手中夺取掌控权的成分在，但 PyTorch 主要还是赢在了其相对 TensorFlow 而言更高的灵活性和可用性，从原则为先的角度来看，PyTorch 与 TensorFlow 的不同在于前者使用的是“动态图模式（Eager Mode）”而非“图模式（Graph Mode）”。

动态图模式可被看作是一种标准的脚本执行方式。深度学习框架会随调用立即逐行执行所有操作，这点和任何的 Python 代码执行都一样。因此，代码的调试和理解也都更加容易，操作间的结果和模型表现也更直观。

相较之下，图模式则有两个阶段，第一阶段定义需要执行操作的计算图，其中计算图是一系列代表操作、变量的相交节点，节点相连的边则代表其间的数据流。第二阶段定义延迟执行计算图的优化版本。因为图的执行过程中我们无从得知到底发生着什么，所以这种分阶段的方式让代码调试更具挑战，也更难理解。这与“解释性”和“编译性”编程语言类似，可解释的 Python 调试要比 C++更容易。

尽管现在 TensorFlow 也默认拥有了动态图模式，但研究社区和多数大型技术公司都已经习惯了 PyTorch 的解决方案。至于 PyTorch 更胜一筹的深层解释，请见这里。总之，主流 AI 大会 NeurIPS 上最优秀的那些非谷歌的生成性人工智能都是用的 PyTorch。

机器学习训练组件

追根究底，影响机器学习的模型训练耗时主要有两个因素:1. 计算（FLOPS），即每层内运行密集的矩阵乘法 2. 内存（带宽），即等待数据或层权重获取计算资源。常见受带宽限制的操作有归一化、点式操作、SoftMax、ReLU。

曾经主要影响机器学习训练时长的计算时间、等待矩阵乘法等因素，随着英伟达 GPU 的不断发展都已不再重要。英伟达的 FLOPS 在摩尔定律下英伟达的 FLOPS 在摩尔定律下提升了多个数量级，但主要架构变化还是集中在张量核心及低精度浮点格式上，内存方面则没有太多变化。

英伟达 GPU 增长趋势

2018 年的 BERT 模型和英伟达的 GPU V100 均是时代顶尖产品，我们不难发现，矩阵乘法已经不再是提高模型性能的主要因素。自此之后，最为先进的模型在参数数量上有了 3 至四个数量级的增长，而最快的 GPU 也有了一个数量级的增长。

PyTorch中操作类占比

即使是在 2018 年，纯粹的计算性质工作负载在占据 99.8%的 FLOPS 同时，仅占用了 61%的运行时间。归一化和逐点操作相较矩阵乘法而言，分别拥有 250 倍和 700 倍 FLOPS 的减少，但同时也消耗了模型近 40%的运行时间。

内存墙

随着模型规模的不断扩张，大型语言模型的权重本身便占据了千亿、乃至太字节。百度和 Meta 所部署的生产型推荐网络其中的大规模嵌入表就可占用几十兆字节内存，大型模型训练或推理中的大部分时间并没有花在矩阵乘法的计算，而是在等待数据到达计算资源。或许你会问，为什么架构师不把更多内存放在靠近计算模块的地方？答案只有一个字，贵。

存储层级

存储层级所遵循的规律是从近且快，到慢但廉价的。共享内存池最近可以在同一块芯片上，通常这种情况是由 SRAM（静态随机存取存储器）组成。有的机器学习 ASIC 尝试利用大型 SRAM 池保存模型权重，但即使是Cerebras价值约5百万美元的晶圆规模芯片上也只有 40G 的 SRAM。我们没有足够的内存来容纳 100B 以上的参数模型权重。

英伟达的架构常常会在芯片上使用更为少量的内存，当前一代 A100 包括了 40MB 内存，而下一代的 H100 上也只有 50MB。台积电 5 纳米工艺节点上 1GB 的 SRAM 需要大约 200 平方毫米的硅，加上相关的控制逻辑和结构的实现后就需要超过 400 平方毫米的硅，或使用英伟达数据中心 GPU 总逻辑面积的 50%左右。考虑到 A100 GPU 超过 1 万美元的价格，H100 的价格很可能也会两万美元起步，经济层面上这条路行不通。就算我们忽略英伟达数据中心 GPU 的 75%毛利率（约为四倍加价），实现产品的完全产出所需要每 GB 的 SRAM 内存，成本仍在一百美元上下。

此外，芯片上 SRAM 内存的花销并不会随着传统摩尔定律所带来的工艺技术缩减而下降太多，在下一代台积电 3 纳米工艺技术下，同样 1GB 的内存成本实际是在增长的。3D SRAM 虽然会在一定程度上降低 SRAM 成本，但这也只是价格曲线的暂时下跌。

存储层次中的下一层是紧密耦合的片外内存 DRAM。DRAM 相较于 SRAM，延迟要高上一个数量级（约 100 纳秒和 10 纳秒的区别），但 DRAM 也要便宜许多（DRAM 每 GB 一美元，SRAM 每 GB 一百美元）

数十年来 DRAM 一直遵循摩尔定律，事实上，在戈登·摩尔创造“摩尔定律”这个词时，英特尔的主要业务就是 DRAM。摩尔关于晶体管密度与成本的经济预测对 2009 年之前的 DRAM 通常都是准确的，但自 2012 年来，DRAM 的成本几乎没有提升。

DRAM 每 GB 价格

我们对内存的需求只增不减，目前 DRAM 已经占据了服务器总成本的50%。内存墙的存在已经开始在产品中显露出来了。相比英伟达 2016 年的 GPU P100，2022 年刚刚发售的 GPU H100 在内存容量上提升了五倍（16GB 到 80GB 的提升），而 FP16 性能却提升了足足 46 倍（21.1 TFLOPS 至 989.5 TFLOPS）。

容量瓶颈与另一同样重要的带宽瓶颈息息相关。并行化是增加内存带宽的主要手段，如今的 DRAM 每 GB 价格区区几美元，而英伟达为了达到机器学习所需要的巨大带宽而用上了 HBM 内存，这是一种由3D堆叠的DRAM层所组成的设备，且需要更为昂贵的包装。HBM 的价格区间为 10 至 20 美元每 GB，其中包含有包装与产量成本。

内存带宽与容量限制在英伟达 A100 GPU 中被反复提及。没有大量优化过的 A100 常常会有极低的 FLOPS 利用率，FLOPS 利用率是通过计算（模型训练时总计算 FLOPS）/(GPU 在模型训练时间内理论可计算 FLOPS）而得出。

即使是在顶尖研究者所做的大量优化下，60%的 FLOPS 利用率对于大型语言模型训练而言也算是非常高的了。剩下的时间都是开销，包括等待其他计算或内存数据的空闲期，或为减少内存瓶颈而进行即时重新计算结果。

FLOPS 在 A100 至 H100 两代间增长了 6 倍有余，但内存带宽却只有 1.65 倍的增长，从而导致了许多对 H100 低利用率问题的担忧。人们为让 A100 绕过内存墙搞出了许多变通方案，而这种努力在 H100 上恐怕会只多不少。

H100为Hopper架构带来了分布式共享内存和二级组播，其中不同 SM（可看作内核）可直接写入其他 SM 的 SRAM（共享内存/L1 缓存）。此举在有效增加缓存大小的同时，缩减了DRAM读写所需的带宽。后续架构也将通过减少向内存传输的操作缓解内存墙的影响。值得注意的是，因为对 FLOPS 的需求会随参数数量增加而立方扩展，对内存带宽及容量的需求则常呈二次曲线发展，所以较大型的模型也更倾向于实现更高的利用率。

算子融合 - 治标不治本

同机器学习模型训练一样，明白自己所处状态才能更精确地进行重要的优化。举例来说，如果我们处于内存带宽约束的状态，时间全花在了内存传输上，那么增加 GPU 的 FLOPS 并不能解决问题。而如果是处于计算约束的状态，笨重的矩阵乘法非常耗时的话，那么试图通过将模型逻辑改写成 C++来削减开销也是没效果的。

虽然 PyTorch 是通过动态图模式增加了灵活性和可用性而赢得的比赛，但动态图模式也不是完美的。在动态图模式中执行的每个操作都需要从内存中读取、计算，再发送到内存之后，才能处理下一个操作。在缺乏大量优化的情况下，这种模式将大大增加对内存带宽的需求。

因此，动态图模式中执行模型的主要优化手段之一是算子融合。用融合操作符的方式取代将每次的中间结果写入内存，在一次计算中计算多个函数，从而尽可能减少对内存的读写。算子融合改善了操作符的调度，也削减了内存带宽和容量的成本。

算子融合简图

这种优化方式常常需要编写自定义 CUDA 内核，这可比简单使用 Python 脚本要难多了。PyTorch 的内置变通方案长期以来在 PyTorch 内部实现了越来越多的操作符，这其中的许多操作符都只是将多个常用操作融合到一个更为复杂的函数中。

操作符的增加让 PyTorch 中的模型创建更加轻松，随着内存读写次数的减少，动态图模式的性能也更快。但随之而来的代价是 PyTorch 中运算符在短短几年内就膨胀至两千有余。

人们常说软件开发者是懒人，但老实说又有谁不是呢。在熟悉了 PyTorch 中某个新增的操作符后，开发者们通常只会觉得这个新操作符能让自己少些点代码，而完全没有意识到其中的性能提高。

此外，并不是所有的操作都能融合。大部分时间我们都在决定要融合哪些操作，又要将哪些操作分配给芯片和集群层面特定的计算资源。虽然一般来说算子融合的策略都多少有些类似，但根据架构的不同也会有所区别。

英伟达称王

操作符数量的发展和其默认的首选地位让英伟达受益很多，每个被迅速优化的操作符都是针对英伟达架构的，且对其他硬件并不适用。任何想要完整实现 PyTorch 的人工智能硬件初创公司，都得靠高性能才能原生支持两千多个且还在不断增长的操作符。

在 GPU 上训练高 FLOPS 利用率的大型模型所需的技术水平越来越高，为达到最佳性能所需花样也越发繁多。动态图模式的执行和算子融合，意味着软件、技术、模型的开发都要被迫适应最新一代 GPU 的计算和内存比例范围。

内存墙是任何机器学习芯片的开发者都逃不开的命运。ASIC 必须要能支持常用框架的同时，也要能支持混合使用英伟达及外部库的、基于 GPU 优化的 PyTorch 代码这一默认开发策略。因此，为图更高 FLOPS 及更严格的编程模型，而主动放弃 GPU 的各种非计算包袱这种行为是非常没有意义的。

易用性是王道。

要打破这种恶性循环的唯一方式是将英伟达 GPU 上运行模型的软件，尽可能地无缝转移至其他硬件上。随着 PyTorch 2.0、OpenAI Triton，以及诸如MosaicML的MLOps公司所提供的模型架构稳定性和抽象逐渐得到主流承认，芯片解决方案的架构和性价比逐渐取代了英伟达卓越的软件所带来的易用性，成为驱动购买力的主要因素。

PyTorch 2.0

数月之前刚刚成立的PyTorch基金会正式脱离了Meta的掌握。在向开放式开发和管理模式转变的同时，2.0 的早期测试版本已经发布，并预计于 2023 年三月全年上市。PyTorch 2.0 与前代最主要的区别在于，新增的一个支持图执行模型的编译解决方案，让各种硬件资源的利用更加轻松。

PyTorch 2.0 在英伟达 A100 上的训练性能有了86%的提升，CPU上的推理则有26%的提升，极大地缩减了模型训练所需的计算时间和成本。而这种性能提升也可以类推至包括AMD、英特尔、Tenstorrent、Luminous Computing、特斯拉、谷歌、亚马逊、微软、Marvell、Meta、Graphcore、Cerebras、SambaNova 在内的多个 GPU 和加速器上。

PyTorch 2.0 的性能提升在未经优化的硬件上更为明显。Meta 及其他公司对 PyTorch 的大量贡献背后，是希望能在他们数十亿美元的训练集群上，以最小的努力实现更高的 FLOPS 利用率，让他们的软件栈更易于移植到其他硬件上，从而为机器学习领域引入新的竞争力。

分布式训练也受益于 PyTorch 2.0，数据并行、分片、管道并行及张量并行均得到了更优秀的 API 支持。除此之外，PyTorch 2.0 也通过全栈提供了对动态图形的原生支持，让LLM不同序列长度等更易于支持。这也是主流编译器第一次支持从训练到推理的动态形状。

PrimTorch

对任何非英伟达 GPU 之外的任何机器学习 ASIC 来说，想要编写一个完整支持全部两千余个操作符的高性能后端是非常具有挑战性的。而 PrimTorch 却可以在保障 PyTorch 终端用户可用性不变的前提下，将操作符数量减少至约 250 个原始操作符，让非英伟达的 PyTorch 后端实现更简单容易，定制硬件和操作系统的供应商也更容易提出自己的软件栈。

TorchDynamo

稳健的图定义是向图模式转变的必需品，而过去五年间 Meta 和 PyTorch 在这方面解决方案的尝试都有着明显的缺陷，直到 TorchDynamo 的出现。TorchDynamo 可接收任何 PyTorch 用户脚本并生成FX图，甚至是调用三方外部库的脚本也可以。

Dynamo 将所有复杂操作都压缩为 PrimTorch 中约 250 个原始操作符。图成型后所有未使用的操作会被弃置，成型的图决定了有哪些中间操作需要被存储或写入内存，有哪些可以被融合。这种方式极大地削减了模型内开销，对用户而言也是无感的。

目前在不修改任何源码的前提下，TorchDynamo已在超过七千个PyTorch模型上通过了可行性测试，其中不乏来自 OpenAI、HuggingFace、Meta、英伟达、Stability.AI 的模型。这七千多个模型是从 GitHub 上热度最高的 PyTorch 项目中直接选取的。

谷歌的 TensorFlow、Jax 及其他图模式的执行管道，通常需要用户自行保障模型对编译器架构的兼容性，才能确保图可以被捕获。而 Dynamo 通过启用部分图捕获、受保护的图捕获及即时重新捕获进行改善。

• 部分图捕获允许模型包括不支持或非 Python 的结构。在无法生成图的模型构造部分插入图断点，并在部分图之间以动态图模式执行。

• 受保护图捕获校验被捕获的图是否可有效执行。保护是指需要重新编译的代码变更，毕竟多次重复执行的同一段代码并不会重新编译。

• 即时重新捕获允许无效执行的图重新被捕获。

PyTorch 意图创建一个依赖 Dynamo 图生成的、统一且流畅的 UX。这项解决方案在不改变用户体验的同时显著提高性能，而图捕获意味着在大型计算资源的基础上执行可以更高效地并行进行。

Dynamo 及AOT自动求导会在之后将优化后的 FX 图传入 PyTorch 本地编译器层，即 TorchInductor。其他硬件企业也可直接取用此时的图并输入至他们自己的后端编译器中。

TorchInductor

作为原生的 Python 深度学习编译器，TorchInductor 可为多个加速器和后端生成快速代码。Inductor（电感器）可接收包含约 250 个操作符的 FX 图，并进一步将其操作符数量削减至 50 左右。在这之后，Inductor 会进入调度阶段，融合算子并确定内存规划。

之后，Inductor 会进入“代码封装”阶段，生成可在 CPU、GPU 及其他人工智能加速器上运行的代码。封装后的代码可调用内核并分配内存，取代了编译器堆栈中解释器的部分。其中，后端代码的生成部分借助 OpenAI 的 GPU 语言 Triton，输出 PTX 代码。对 CPU 而言，英特尔编译器所生成的 C++代码也可以在非英特尔的 CPU 上运行。

后续还会新增更多对硬件的支持，但 Inductor 确实显著降低了在编写 AI 硬件加速器的编译器时所需的工作量。此外，代码性能也得到了优化，对内存带宽和容量的要求也大大地降低了。

我们写的编译器不能只支持 GPU，而要能扩展到对各类硬件后端的支持。C++及（OpenAI）Triton 迫使着我们一定要具备这种通用性。——Jason Ansel – Meta AI

OpenAI Triton

OpenAI 的 Triton 语言对英伟达机器学习闭源软件的护城河有着毁灭性的打击。Triton 可直接接收 Python 脚本，或者更常见地接收通过PyTorch的Inductor堆栈的信息流。随后 Triton 会将输入转换为 LLVM 中间表示并生成代码，使用 cutlass 等开源库取代英伟达的闭源 CUDA 库（如 cuBLAS）。

CUDA 在专职于加速计算的开发者中更为常用，在机器学习研究者或数据科学家之间则没什么知名度。高效地使用 CUDA 并不容易，需要使用者对硬件架构有深入理解，并可能会拖慢开发过程。因此，机器学习专家们常常会依赖 CUDA 专家对他们的代码进行修改、优化，以及并行化。

而 Triton 则弥补了这一差距，让高层语言达到与底层语言相媲美的性能水平。Triton 的内核本身对一般的机器学习研究者而言具备可读性，这一点对语言可用性非常重用。Triton 将内存凝聚、共享内存管理，以及 SM 内部的调度全部自动化，但对元素层面的矩阵乘法没什么太大帮助，后者本身已经足够高效了。此外，Triton 在昂贵的逐点操作方面很有效，对涉及矩阵乘法的大型算子融合而言，也可明显削减复杂如Flash注意力等操作的开销。

时至今日，OpenAI 的 Triton 才正式支持英伟达的 GPU，不过很快就要不同了。数个其他硬件供应商都会在后续对其提供支持，这项开源项目的前途一片光明。其他硬件加速器能够直接集成至 Triton 中 LLVM IR，意味着在新硬件上建立人工智能编译器堆栈的时间将大幅缩短。

英伟达庞大的软件组织缺乏远见，没能利用自己在机器学习软硬件方面的巨大优势，一举成为机器学习的默认编译器。英伟达对可用性关注的缺失让外界中 OpenAI 及 Meta 得以开发出向其他硬件方向移植的软件栈。他们为什么没能为机器学习研究者们开发出一个像是 Triton 之类的*简化版*CUDA？为什么像是Flash注意力一类的技术是出自一个博士生而不是英伟达本身？

本篇报告中的剩余部分将会列出能让微软拿下一城的具体硬件加速器，以及目前正被快速集成至 PyTorch 2.0 或 OpenAI Trion 软件栈中多家公司的硬件产品。此外，报告中也将列出相反观点，为英伟达在人工智能培训领域的护城河或实力提供辩护。

查看英文原文：How Nvidia’s CUDA Monopoly In Machine Learning Is Breaking - OpenAI Triton And PyTorch 2.0