编者按：在机器学习领域，无论是硬件还是软件，英伟达无疑均拥有巨大优势，后者用 CUDA 建立起了一道软件的护城河。可惜的是，这家公司缺乏远见，未能利用其在机器学习硬软件方面的巨大优势，让自己成为机器学习默认的编译器。而它对可用性与易用性的忽视，让 OpenAI 与 Meta 得以趁虚而入，其主导地位正在被打破。文章来自编译。
过去十年，机器学习软件开发的格局发生了重大变化。许多框架经历了兴衰沉浮，但大多数都严重依赖于英伟达的 CUDA，而且在 Nvidia GPU 上表现最佳。不过，随着 PyTorch 2.0 与 OpenAI 的 Triton 的出现，英伟达在这个领域主要靠软件护城河保证的主导地位正在被打破。
这份报告将讨论各种问题，比如为什么谷歌的 TensorFlow 会输给 PyTorch 、为什么谷歌未能大大方方地利用其早期在 AI 领域的领导地位、机器学习模型训练时间的主要组成部分、内存容量/带宽/成本墙、模型优化、为什么别的 AI 硬件公司至今无法撼动英伟达的主导地位，为什么硬件的作用会逐渐凸显、英伟达在 CUDA 方面的竞争优势如何被抹去，以及英伟达的竞争对手之一如何在一个大型的云服务上训练硅片并取得了重大胜利。
一个高度概括的结论是，机器学习模型的默认软件栈将不再是英伟达闭源的 CUDA。球本来在英伟达的球场上，但他们让 OpenAI 和 Meta 控制了软件栈。由于英伟达专有的工具失败，这个生态体系已经建立起自己的工具，现在英伟达的护城河已经露出破绽，而且永远也补不回来了。
TensorFlow Vs. PyTorch
几年前，这个框架生态体系还十分的碎片化，但 TensorFlow 是领跑者。谷歌看起来本来就要控制机器学习行业了。他们凭借着最常用的框架 TensorFlow，并通过设计/部署唯一成功的 AI 应用加速器 TPU 获得了先发优势。
PyTorch迅速在各大会议上占据口碑
但最后获胜的是 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++ 之别。调试 Python 更容易，主要是因为它是解释型的。
虽然 TensorFlow 现在默认也支持动态图模式，但研究社区和大多数大型科技公司的讨论基本上是围绕着 PyTorch 展开的。几乎所有登上了新闻的生成式 AI 模型都是基于 PyTorch 的，这个事实就是例证。 谷歌自己的生成式 AI 模型也是基于 Jax ，而非 TensorFlow 。
当然，使用类似 TensorFlow 和 Keras 这样的其他框架的图像神经网络还有一长串，但新模型开发的计算预算都流向了 PyTorch 模型。一般来说，如果你在 NeurIPS （主要的人工智能会议）会议大厅里面打听一下的话，所有的生成式人工智能，只要不是谷歌做的，基本上都是使用 PyTorch。
机器学习训练的组成部分
如果我们把机器学习模型训练简化为最简单的形式，那么机器学习模型的训练时间可以有有两个主要的时间组件。
计算（FLOPS）：在每一层跑密集矩阵乘法
内存（带宽）：等待数据或神经网络层的权重送抵计算资源。像规格化、点态运算、SoftMax、ReLU 都是带宽受限操作的常见例子。
在过去，机器学习训练时间的主导因素是计算时间，需要等待矩阵乘法完成。随着英伟达的 GPU 不断发展，这很快就不再是主要问题。
通过利用摩尔定律，英伟达的 FLOPS 提升了好几个数量级，但改变主要是架构性的，比方说 tensor 内核，以及较低精度的浮点格式。相比之下，记忆并没有遵循相同的路径。
内存的提升赶不上FLOPS
如果回到 2018 年，在 BERT 模型是最先进的模型，而 Nvidia V100 是最先进的 GPU 的时候，可以看到矩阵乘法不再是提高模型性能的主要因素。从那时起，最先进的模型在参数数量上以及增长了 3 到 4 个数量级，而最快的 GPU 在 FLOPS 上只增长了一个数量级。
PyTorch不同算子类别在FLOPS与运行时的占比
即便在 2018 年，纯计算密集型的工作负载占了 FLOPS 的 99.8%的时候，在运行时的占比也只有 61%。与矩阵乘法相比，归一化和逐点运算消耗的 FLOPS 分别少了 250 倍和 700 倍，但在模型运行时间的占比却近 40%。
内存墙
随着模型规模的不断飙升，大型语言模型光是用于模型权重就要 100 GB（如果不是 TB的话）。百度和 Meta部署用于生产的推荐神经网络要数十 TB 的内存来存储海量的嵌入表。大型模型训练/推理的大部分时间都没有花在计算矩阵乘法上，而是花在了等待数据到达计算资源上。一个显而易见的问题是，为什么架构师不把更多的内存放在更靠近计算的位置。答案出在钱身上。
离计算越近的内存越贵
内存遵循的是越近越快、越远越慢越便宜的层次结构。最靠近（计算）的共享内存池位于同一块芯片上，一般由 SRAM 构成。一些机器学习 ASIC 试图利用庞大的 SRAM 池来存储模型权重，但这种方法存在问题。即便是 Cerebras 价值约 2500000 美元的晶圆级芯片也只有 40GB 的 SRAM。这样的内存容量不足以容纳参数高达 100B+ 的模型的权重。
英伟达的体系结构在裸片上使用的内存量一直要少得多。目前这代 A100 有 40MB，下一代 H100 有 50MB。台积电 5 纳米工艺节点上 1GB 的 SRAM 需要约 200mm^2 的硅。一旦实现了相关的控制逻辑/结构，将需要超过 400mm^2 的硅，或占据了英伟达数据中心级 GPU 总逻辑面积的约 50%。鉴于 A100 GPU 的成本为 1 万美元以上，而 H100 更是接近 2 万美元以上，从经济角度来看，这种做法是不可行的。即便不考虑英伟达数据中心 GPU约 75% 的毛利率（大概4 倍的加价），成熟产品每 GB SRAM 内存的成本仍将在 100 美元左右。
此外，片上 SRAM 存储器的成本不会遵循传统的摩尔定律，即便工艺技术缩小了成本也不会降低太多。同样是 1GB 内存，采用台积电下一代 3nm 制程工艺的成本反而更高。虽然 3D SRAM 将在一定程度上帮助降低 SRAM 成本，但这也只是暂时压低了一下曲线。
内存层次结构的下一级是紧耦合的片外内存 DRAM。 DRAM 的延迟要比 SRAM 高一个数量级（~>100 纳秒 vs ~10 纳秒），但价格也便宜得多（$1sa GB vs $100s GB。）
几十年来，DRAM 一直遵循着摩尔定律。戈登·摩尔创造出这个词时，英特尔的主要业务就是 DRAM。他对晶体管密度与成本关系的经济预测在 2009 年之前对 DRAM 普遍适用。不过自 2012 年以来，DRAM 的成本几乎都没有改善。
DRAM的成本遵循摩尔定律
对内存的需求只会增加。 DRAM 现在占服务器总成本的 50%。这就是内存墙，这道墙已经出现在产品中。对比一下英伟达的2016 P100 GPU 与刚刚开始出货的 2022 H100 GPU，前者内存容量增加了 5 倍（16GB -> 80GB），但 FP16 性能增加了 46 倍（21.2 TFLOPS -> 989.5 TFLOPS）。
虽然容量是个重要瓶颈，但这个瓶颈其实与另一个主要瓶颈——带宽密切相关。内存带宽增加通常是靠并行性获得的。虽然当今标准 DRAM 的价格仅为每 GB 几美元，但要想获得机器学习所需的海量带宽， 英伟达必须使用 HBM 内存，这是一种由 3D 堆叠的多层 DRAM 层组成的设备，需要更昂贵的封装。 HBM 每 GB大概在 10 到 20 美元之间，其中包括了包装与生产成本。
内存带宽与容量的成本限制频繁出现在英伟达的 A100 GPU 身上。如果不进行大量优化的话，A100 的 FLOPS 利用率往往非常低。FLOPS 利用率衡量的是训练模型所需的 FLOPS 总算量 与 GPU 在模型训练时间内可以计算的理论 FLOPS 之比。
即便经过首席研究人员进行了大量优化，FLOPS 利用率能做到 60% 也被认为是大型语言模型训练很高的利用率了。其余时间都是开销，花在等待来自另一个计算/内存的数据的空闲时间，或者为减少内存瓶颈而对结果进行的即时重新计算。
从目前这代的 A100 到下一代的 H100，FLOPS 增长了 6 倍以上，但内存带宽仅增长了 1.65 倍。这导致许多人担心 H100 的利用率会很低。 A100 需要很多技巧才能绕过内存墙，H100 还需要更多技巧才能绕开。
H100 给 Hopper 架构引入了分布式共享内存与 L2 多播。其想法是不同的 SM（think cores）可以直接写入另一个 SM 的 SRAM（共享内存/L1 缓存）。这可以有效增加缓存的大小并减少了 DRAM 读/写所需的带宽。未来的架构有赖于减少给内存发送的操作，从而最大限度地减少内存墙的影响。应该要注意的是，较大规模的模型往往能实现更高的利用率，因为对 FLOPS 的需求呈指数级增长，而内存带宽和容量需求往往呈线性增长。
算子融合——规避方法
就像训练机器学习模型一样，知道自己所处状态可以缩小有关系的优化的范围。比方说，如果所有时间都花在内存传输上（比方说处在内存带宽受限状态），那么增加 GPU 的 FLOPS 是没有用的。另一方面，如果你将所有时间都花在执行消耗算力的 matmuls函数运算上，那么将模型逻辑用 C++ 重写来减少开销也无济于事。
回顾 PyTorch 获胜的原因，这是由于动态图模式提高了灵活性和可用性，但转向动态图模式并不全是阳光大道。在动态图模式下执行时，每个操作都要从内存读取、计算，然后在处理下一个操作之前再回传给内存。如果不进行大量优化的话，是会显著增加内存带宽需求的。
算子融合（operator fusion）是在动态图模式下执行的模型的主要优化方法之一。算子最后进行了融合，而不是将每个中间结果写入内存，这样一来，一次传递的过程中会计算多个函数，从而将内存读/写的数量最小化。算子融合改善了算子调度、内存带宽以及内存大小成本。
算子融合简化版的解释
这种优化往往涉及到要编写自定义 CUDA 内核，但这比利用简单的 python 脚本要困难得多。作为一种内置的妥协方案， PyTorch正在内部逐步实现了越来越多的算子（operator)。其中很多的算子只是简单地将多个常用运算融合进一个更复杂的函数之中。
算子的增加使得用 PyTorch 创建模型变得更加容易，而且由于内存读/写更少，动态图模式的性能也更快。但缺点是 PyTorch 的算子在几年内激增到 2000 多个。
算子的家族树
有时候我们会说软件开发者很懒，但说实话，差不多每个人都有惰性。如果他们习惯了 PyTorch 内置的新算子，就会继续使用它。开发者可能甚至都没有意识到性能的提高，他们之所以用新算子，是因为这意味着要写的代码变少了。
此外，并非所有算子都可以融合。往往要花费大量时间来决定要融合哪些操作以及将哪些操作分配给芯片与集群级别的特定计算资源。哪些算子在什么地方融合的策略虽然大体相似，但根据架构的不同会有很大差异。
英伟达为王
增加算子以及在业界的默认地位对英伟达是有帮助的，因为每个算子都针对英伟达的架构进行了快速优化，但对任何其他硬件就没有做优化。如果某家 AI 硬件初创企业想要全面实现 PyTorch，就意味着必须高性能地原生支持 2000 多个算子，而且这份清单还在不断增长。
在 GPU 上训练具有高 FLOPS 利用率的大型模型所需的人才水平越来越高，因为实现性能最大化需要运用各种技巧。动态图模式执行，再加上算子融合，意味着开发的软件、技术和模型都得适应当前一代 GPU 的计算/内存比。
所有开发机器学习芯片的都受制于同一道内存墙。ASIC 有责任支持最常用的框架。ASIC 受制于默认的开发方法，针对GPU 优化过的 PyTorch 代码（是英伟达库与外部库的混合）。在这种情况下，避开 GPU 的各种非计算的包袱，去支持更高 FLOPS 以及更严格的编程模型的架构意义不大。
易用性为王。
打破恶性循环的唯一方法是让在 Nvidia GPU 上运行模型的软件尽可能轻松地无缝转移到其他硬件。随着模型架构的稳定和来自 PyTorch 2.0、OpenAI Triton 和 MLOps公司（如MosaicML ）的抽象成为默认，芯片解决方案的架构和经济性开始成为购买的最大驱动力，而英伟达出色的软件带来的易用性退居其次。
PyTorch 2.0
几个月前，PyTorch 刚刚从Meta独立出来，成立了自己的 PyTorch Foundation。除了开发变成开放式，以及采取治理模式以外，2.0 还推出了早期测试，并在 3 月全面上市。PyTorch 2.0 带来了许多变化，但主要区别在于增加了一个支持静态图执行模型的编译解决方案。这一改变可以让正确利用各种硬件资源变得更加容易。
PyTorch 2.0 在 NvidiaA100 上的训练性能提升了 86% ，在 CPU 上的推理性能提升了 26%！这大大降低了训练模型所需的计算时间和成本。这些好处还可以扩展到 AMD 、英特尔、Tenstorrent、Luminous Computing、特斯拉、谷歌、亚马逊、微软、Marvell、Meta、Graphcore、Cerebras、SambaNova 等其他 GPU 和加速器身上。
对于目前未经优化的硬件，PyTorch 2.0 带来的性能改进甚至更大。Meta 等公司对 PyTorch 的巨大贡献源自这样一个事实，即他们希望事半功倍，在自家价值数十亿美元的 GPU 训练集群上少花功夫就能更容易地实现更高的 FLOPS 利用率。他们也有理由让自己的软件栈移植到其他硬件变得更加容易，从而将竞争引入到机器学习领域。
通过为数据并行、分片、流水线并行（ipeline parallelism）以及张量并行（tensor parallelism）提供更好的 API 支持，PyTorch 2.0 还为分布式训练带来了进展。此外，它为整个技术栈提供了对 dynamic shapes（在推理阶段指定部分或者所有输入数据的维度）的原生支持，这使得对 LLM 的各种序列长度的支持容易了许多。这是主要编译器第一次支持从训练到推理的 Dynamic Shapes。
PyTorch 2.0抽象了对硬件资源的利用
PrimTorch
给 PyTorch 写一个表现良好的后端一直都很困难。要想完全支持所有 2000 多个算子，除了 Nvidia GPU 以外，其他机器学习 ASIC 都很难做到。但 PrimTorch 把算子的数量减少到约 250 个原始算子，同时对 PyTorch 最终用户的可用性仍维持不变。PrimTorch 让实现 PyTorch 非英伟达后端变得更加简单，更可达。定制硬件与系统供应商可以更轻松地推出自己的软件栈。
TorchDynamo
转向静态图模式需要给出可靠的静态图定义。Meta 与 PyTorch 大概尝试了 5 年的时间来实现这一点，但是他们提出的解决方案均存在明显缺点。现在，他们终于用 TorchDynamo 破解了这一难题。TorchDynamo 会摄取任意 PyTorch 用户脚本，其中包括调用外部第三方库的脚本，然后生成 FX graph。
Dynamo 把各种复杂算子减少到 PrimTorch 的约 250 个原始算子。一旦graph形成之后，未使用的算子将被丢弃，graph会确定哪些中间算子需要保存或写入内存，哪些可以被融合。这可以极大减少模型的开销，同时对用户来说也是无缝的。
在 7000 个经过测试的 PyTorch 模型当中，TorchDynamo 对其中的 99% 以上都是适用的，包括来自 OpenAI、HuggingFace、Meta、Nvidia、Stability.AI 等的模型，不需要对原始代码进行任何修改。这受测的 7000 个模型是从 GitHub 上使用了 PyTorch 的最流行项目当中任意挑选出来的。
谷歌的 TensorFlow/Jax 等其他的静态图模式执行流水线往往要用户确保自己的模型适配编译器架构，因为这样才能捕捉到静态图。Dynamo 通过支持部分图捕捉（partial graph capture）、受保护的图捕捉（guarded graph capture）以及即时重新捕捉（just-in-time recapture）改变了这一点。
部分图捕捉让模型可以纳入不受支持的/非 python 构造。当静态图没法为模型的那一部分生成时，就会插入一个 graph break，并且将在部分图之间以动态图模式执行不受支持的构造。
受保护的图捕捉会检查捕捉到的图对执行是否有效。保护就是需要重新编译的变更。这一点很重要，因为同样的代码多次执行并不会多次重新编译。
如果捕捉的图对执行无效，则即时重新捕捉可以重新捕捉静态图。
什么时候该用哪种图捕捉模式
PyTorch 的目标是建立一个有着流畅 UX 的统一前端。这个前端会利用 Dynamo 来生成图。这个解决方案的用户体验不会发生变化，但性能可以得到显著提升。捕捉图意味着执行可以基于大量计算资源并行自行。
经过 Dynamo 与 AOT Autograd 之后，就可以将优化的 FX 图传递给 PyTorch 原生编译器级的 TorchInductor。硬件公司也可以将该图传递给自己的后端编译器。
TorchInductor
TorchInductor 是 python 原生的深度学习编译器，可以为多个加速器和后端生成快速代码。Inductor 可接受具有约 250 个算子的 FX 图，然后将算子减少到约 50 个。Inductor 然后会进入调度阶段，对算子进行融合，并规划好内存的使用。
之后，Inductor 会转入“Wrapper Codegen ”，后者会生成在 CPU、GPU 或其他 AI 加速器上运行的代码。该 Wrapper codegen 取代了编译器技术栈解释器的部分，它可以调用内核及分配内存。后端代码生成部分会利用适用于 GPU 的 OpenAI Triton 并输出 PTX 代码。对于 CPU，会有一个英特尔编译器生成 C++（也适用于非英特尔 CPU）。
未来 Inductor 将支持更多硬件，但它的关键作用在于大大降低了编译器团队在为 AI 硬件加速器写编译器时的工作量。此外，代码针对性能进行了更多的优化。对内存带宽和容量的要求显著降低了。
我们不想开发只支持 GPU 的编译器。我们希望它可以扩展，支持各种硬件后端，而且可以扩充，除了 [ OpenAI ] Triton 以外还有一个C++编译器可强制实现这种通用性。
Jason Ansel——Meta AI
OpenAI Triton
对于英伟达的机器学习闭源软件护城河来说，OpenAI 的 Triton 是颠覆性的。Triton 直接接受 Python 或通过 PyTorch Inductor 技术栈获得数据。后者将是最常见的用例。之后，Triton 会将输入转换为 LLVM（低级虚拟机） 中间表示，然后生成代码。针对 Nvidia GPU，它可以跳过 Nvidia 的闭源 CUDA 库（如 cuBLAS），而是用开源库（如 cutlass），直接生成 PTX 代码。
使用 CUDA 的一般是专门从事加速计算的人员，但机器学习研究人员与数据科学家却不怎么了解这个东西。要想高效使用可能会很有挑战性，而且需要深入了解硬件架构，这可能会导致开发过程放缓。因此，机器学习专家可能要依赖于 CUDA 专家来对其代码进行修改、优化以及并行化。
更高级语言与较低级语言的这道性能鸿沟被 Triton 弥合了。Triton 内核本身对于典型的机器学习研究人员来说非常清晰，这对于可用性来说是非常重要的。Triton 在 SM 内自动执行内存合并、共享内存管理与调度等功能。Triton 对与按元素进行的矩阵乘法不是特别有用，但后者已经有高效的解决方案了。Triton 对昂贵的按点操作以及降低更复杂操作的开销非常有用，比方说 Flash Attention，后者牵涉到占据了融合操作很大一部分的矩阵乘法。
OpenAI Triton 目前官方支持的只有 Nvidia GPU，但在不久的将来这种情况会发生变化。未来Triton将支持多家其他硬件供应商，这个开源项目的发展势头十分迅猛。其他硬件加速器可以直接集成到 LLVM IR，而后者是 Triton 的一部分，这种能力大大减少了为新硬件开发 AI 编译器栈的时间。
英伟达的软件组织很庞大，但缺乏远见，未能利用其在机器学习硬软件方面的巨大优势，让自己成为机器学习默认的编译器。他们缺乏对可用性的关注，这导致 OpenAI 与 Meta 这些外来者开发出了可移植到其他硬件的软件栈。为什么他们不能给机器学习研究人员开发出像 Triton 这样的“简化”版的 CUDA？而像 Flash Attention 这样的东西，为什么会出自博士生之手，而不是英伟达自己？
译者：boxi。