OpenAI Triton和PyTorch 2.0快速发展，英伟达在机器学习领域CUDA垄断被打破

TensorFlow 与 PyTorch

几年前，框架生态系统相当分散，但 TensorFlow 是领跑者。谷歌看起来准备控制机器学习行业。他们凭借最常用的框架 TensorFlow 以及通过设计/部署唯一成功的 AI 应用特定加速器 TPU 获得了先发优势。

相反，PyTorch 赢了。谷歌未能将其先发优势转化为对新兴 ML 行业的主导地位。如今，谷歌在机器学习社区中有些孤立，因为它不使用 PyTorch 和 GPU，而是使用自己的软件堆栈和硬件。按照典型的谷歌方式，他们甚至还有一个名为 Jax 的第二个框架，它直接与 TensorFlow 竞争。

由于大型语言模型，尤其是来自 OpenAI 和使用 OpenAI API 或正在构建类似基础模型的各种初创公司的大型语言模型，谷歌在搜索和自然语言处理方面的主导地位甚至没完没了。虽然我们认为这种厄运和阴霾被夸大了，但这是另一个故事。尽管存在这些挑战，谷歌仍处于最先进机器学习模型的前沿。他们发明了transformers ，并在许多领域（PaLM、LaMBDA、Chinchilla、MUM、TPU）保持最先进的水平。

回到为什么 PyTorch 赢了。虽然谷歌在争夺控制权方面存在一些因素，但这主要是由于 PyTorch 与 TensorFlow 相比具有更高的灵活性和可用性。如果我们将其归结为第一个主要级别，PyTorch 与 TensorFlow 的不同之处在于使用“ Eager mode ”而不是“ Graph Mode ”。

Eager 模式可以被认为是一种标准的脚本执行方法。深度学习框架会像任何其他 Python 代码一样，逐行立即执行每个操作。这使得调试和理解您的代码更容易，因为您可以看到中间操作的结果并查看您的模型的行为方式。

相反，Graph 模式有两个阶段。第一阶段是表示要执行的操作的计算图的定义。计算图是一系列相互连接的节点，表示操作或变量，节点之间的边表示它们之间的数据流。第二阶段是延迟执行计算图的优化版本。

这种两阶段方法使理解和调试代码更具挑战性，因为在图形执行结束之前您无法看到发生了什么。这类似于“解释”与“编译”语言，如 Python 与 C++。调试 Python 更容易，主要是因为它是解释型的。

虽然 TensorFlow 现在默认具有 Eager 模式，但研究社区和大多数大型科技公司已经围绕 PyTorch 安定下来。

机器学习训练组件

如果我们将机器学习模型训练简化为最简单的形式，那么机器学习模型的训练时间有两个主要的时间组成部分。

1、计算 (FLOPS)：在每一层内运行密集矩阵乘法

2、内存（带宽）：等待数据或层权重到达计算资源。带宽受限操作的常见示例是各种规范化、逐点操作、SoftMax和ReLU 。

如果我们回到 2018 年，那时 BERT 模型是最先进的，Nvidia V100 是最先进的 GPU，我们可以看到矩阵乘法不再是提高模型性能的主要因素。从那时起，最先进的模型在参数数量上增长了 3 到 4 个数量级，而最快的 GPU 在 FLOPS 上增长了一个数量级。

即使在 2018 年，纯计算绑定的工作负载也占 FLOPS 的 99.8%，但仅占运行时的 61%。与矩阵乘法相比，归一化和逐点运算分别实现了 250 倍和 700 倍的 FLOPS，但它们消耗了模型运行时间的近 40%。

随着模型规模的不断飙升，大型语言模型仅用于模型权重就需要 100 GB（如果不是 TB）。百度和 Meta 部署的生产推荐网络需要数十 TB 的内存来存储其海量嵌入表。大型模型训练/推理中的大部分时间都没有花在计算矩阵乘法上，而是在等待数据到达计算资源。显而易见的问题是，为什么架构师不将更多内存放在更靠近计算的位置。答案是显而易见的——成本。

Nvidia 的体系结构在裸片上一直使用的内存量要少得多。当前一代A100有40MB，下一代H100有50MB。台积电 5 纳米工艺节点上的 1GB SRAM 需要约 200mm^2 的硅。一旦实现了相关的控制逻辑/结构，将需要超过 400mm^2 的硅，或 Nvidia 数据中心 GPU 总逻辑面积的大约 50%。鉴于 A100 GPU 的成本为 1 万美元以上，而 H100 更接近 2 万美元以上，从经济角度来看，这是不可行的。即使忽略 Nvidia 在数据中心 GPU 上约 75% 的毛利率（约 4 倍加价），对于完全量产的产品，每 GB SRAM 内存的成本仍将在 100 美元左右。

此外，片上SRAM存储器的成本不会随着传统摩尔定律工艺技术的缩小而降低太多。同样的1GB内存，采用台积电下一代3nm制程工艺，成本反而更高。虽然 3D SRAM 将在一定程度上帮助降低 SRAM 成本，但这只是曲线的暂时弯曲。

内存层次结构的下一步是紧密耦合的片外内存 DRAM。DRAM 的延迟比 SRAM 高一个数量级（~>100 纳秒对~10 纳秒），但它也便宜得多（$1sa GB 对 $100s GB。）

几十年来，DRAM 一直遵循着摩尔定律。当戈登摩尔创造这个词时，英特尔的主要业务是 DRAM。他对晶体管密度和成本的经济预测在 2009 年之前对 DRAM 普遍适用。不过自 2012 年以来，DRAM 的成本几乎没有改善。

对内存的需求只会增加。DRAM 现在占服务器总成本的 50% 。这就是内存墙，它已经出现在产品中。将 Nvidia 2016年的P100 GPU 与2022 刚刚开始出货的H100 GPU 进行比较，内存容量增加了 5 倍（16GB -> 80GB），但 FP16 性能增加了 46 倍（21.2 TFLOPS -> 989.5 TFLOPS）。

虽然容量是一个重要的瓶颈，但它与另一个主要瓶颈带宽密切相关。增加的内存带宽通常是通过并行性获得的。虽然如今标准 DRAM 的价格仅为每 GB 几美元，但为了获得机器学习所需的海量带宽，Nvidia 使用 HBM 内存，这是一种由3D 堆叠 DRAM 层组成的设备，需要更昂贵的封装。HBM 在每 GB 是10 到 20 美元的范围内，这包括封装和产量成本。

内存带宽和容量的成本限制不断出现在 Nvidia 的 A100 GPU 中。如果不进行大量优化，A100 往往具有非常低的 FLOPS 利用率。FLOPS 利用率衡量训练模型所需的总计算 FLOPS 与 GPU 在模型训练时间内可以计算的理论 FLOPS。

即使领先研究人员进行了大量优化，60% 的 FLOPS 利用率也被认为是大型语言模型训练的非常高的利用率。其余时间是开销，空闲时间花在等待来自另一个计算/内存的数据，或者及时重新计算结果以减少内存瓶颈。

从当前一代的 A100 到下一代 H100，FLOPS 增长了 6 倍以上，但内存带宽仅增长了 1.65 倍。这导致许多人担心 H100 的利用率低。A100需要很多技巧才能绕过内存墙，H100 还需要实现更多技巧。

H100 为Hopper 带来了分布式共享内存和 L2 多播(multicast)。这个想法是不同的 SM（think cores）可以直接写入另一个 SM 的 SRAM（共享内存/L1 缓存）中。这有效地增加了缓存的大小并减少了DRAM 读/写所需的带宽。未来的架构将依赖于向内存发送更少的操作，以最大限度地减少内存墙的影响。应该注意的是，较大的模型往往会实现更高的利用率，因为 FLOPS 需要按 2^n 缩放，而内存带宽和容量需求往往会按 2*n 缩放。

算子融合——解决方法

有分析人士认为，就像训练 ML 模型一样，了解您所处的状态可以让您缩小重要的优化范围。例如，如果您将所有时间都花在内存传输上（即您处于内存带宽限制状态），那么增加 GPU 的 FLOPS 将无济于事。另一方面，如果您将所有时间都花在执行大型 matmuls（如计算绑定机制）上，那么将您的模型逻辑重写为 C++ 以减少开销将无济于事。

而回顾 PyTorch 获胜的原因，这是由于 Eager 模式提高了灵活性和可用性，但转向 Eager 模式并不全是阳光和彩虹。在 Eager 模式下执行时，每个操作都从内存中读取、计算，然后在处理下一个操作之前发送到内存。如果不进行大量优化，这会显着增加内存带宽需求。

在 Eager 模式下执行的模型的主要优化方法之一称为算子融合（operator fusion）。操作被融合，而不是将每个中间结果写入内存，因此在一次传递中计算多个函数以最小化内存读/写。算子融合改善了运算符调度、内存带宽和内存大小成本。

算子的增加使得在 PyTorch 中创建模型变得更容易，并且由于内存读/写更少，Eager 模式的性能更快。缺点是 PyTorch 在几年内激增到 2,000 多个算子。

此外，并非所有算子都可以融合。通常花费大量时间来决定要融合哪些操作以及将哪些分配给芯片和集群级别的特定计算资源。哪些算子在哪里融合的策略虽然大体相似，但根据架构的不同会有很大差异。

算子的增长和默认位置对 Nvidia 有所帮助，因为每个算子都针对其架构进行了快速优化，但并未针对任何其他硬件进行优化。如果一家 AI 硬件初创公司想要全面实施 PyTorch，那就意味着以高性能原生支持不断增长的 2,000 个算子列表。

由于提取最大性能所需的所有技巧，在 GPU 上训练具有高 FLOPS 利用率的大型模型所需的人才水平越来越高。Eager mode 执行加上operator fusion 意味着开发的软件、技术和模型被推动以适应当前一代 GPU 具有的计算和内存比率。

每个开发机器学习芯片的人都依赖于同一个内存墙。ASIC 有责任支持最常用的框架。ASIC 受制于默认的开发方法，GPU 优化的 PyTorch 代码混合了 Nvidia 和外部库。避开 GPU 的各种非计算包袱而支持更多 FLOPS 和更严格的编程模型的架构在这种情况下意义不大。

但易用性为王。

打破恶性循环的唯一方法是让在 Nvidia GPU 上运行模型的软件尽可能轻松地无缝转移到其他硬件。随着模型架构的稳定和来自 PyTorch 2.0、OpenAI Triton和 MLOps 公司（如 MosaicML）的抽象成为默认，芯片解决方案的架构和经济性开始成为购买的最大驱动力，而不是提供给它的易用性Nvidia 的高级软件。

PyTorch 2.0

PyTorch基金会成立并于几个月前从 Meta 的羽翼下撤出。除了对开放式开发和治理模型的更改外，2.0 还发布了早期测试，并于 3 月全面上市。PyTorch 2.0 带来了许多变化，但主要区别在于它添加了一个支持图形执行模型的编译解决方案。这种转变将使正确利用各种硬件资源变得更加容易。

PyTorch 2.0在 Nvidia A100 上的训练性能提升了 86% ，在 CPU 上的推理性能提升了 26% ！这大大减少了训练模型所需的计算时间和成本。这些好处可以扩展到来自AMD 、英特尔、Tenstorrent 、Luminous Computing、特斯拉、谷歌、亚马逊、微软、Marvell 、Meta 、Graphcore 、Cerebras 、SambaNova 等的其他 GPU 和加速器。

对于当前未优化的硬件，PyTorch 2.0 的性能改进会更大。Meta 和其他公司对 PyTorch 的巨大贡献源于这样一个事实，即他们希望在他们价值数十亿美元的 GPU 训练集群上以更少的努力，更容易地实现更高的 FLOPS 利用率。他们也有动力使他们的软件堆栈更易于移植到其他硬件，以将竞争引入机器学习领域。

PyTorch 2.0 还通过更好的 API 支持数据并行、分片、管道并行和张量并行，为分布式训练带来了进步。此外，它在整个堆栈中原生支持动态形状，在许多其他示例中，这使得 LLM 的不同序列长度更容易支持。这是主要编译器首次支持从训练到推理的 Dynamic Shapes。

为 PyTorch 编写一个完全支持所有 2,000 多个算子的高性能后端对于除 Nvidia GPU 之外的每个机器学习 ASIC 来说都是困难的。PrimTorch 将算子的数量减少到约 250 个原始算子，同时还保持 PyTorch 最终用户的可用性不变。PrimTorch 使 PyTorch 的不同非 Nvidia 后端的实现变得更加简单和易于访问。定制硬件和系统供应商可以更轻松地推出他们的软件堆栈。

TorchDynamo

转向图形模式需要一个可靠的图形定义。Meta 和 PyTorch 已经尝试了大约 5 年的时间来实现这一点，但是他们提出的每个解决方案都有明显的缺点。他们终于用 TorchDynamo 破解了这个难题。TorchDynamo 将摄取任何 PyTorch 用户脚本，包括调用外部 3rd 方库的脚本，并生成FX 图形。

Dynamo 将所有复杂算子减少到 PrimTorch 中的约 250 个原始算子。一旦图形成，未使用的算子将被丢弃，图确定哪些中间算子需要存储或写入内存，哪些可能被融合。这极大地减少了模型内的开销，同时对用户来说也是无缝的。

TorchDynamo 已经适用于7,000 个经过测试的 PyTorch 模型中的 99% 以上，包括来自 OpenAI、HuggingFace、Meta、Nvidia、Stability.AI 等的模型，而无需对原始代码进行任何更改。测试的 7,000 个模型是从 GitHub 上使用 PyTorch 的最流行项目中不分青红皂白地挑选出来的。

谷歌的 TensorFlow/Jax 和其他图形模式执行管道通常要求用户确保他们的模型适合编译器架构，以便可以捕获图形。Dynamo 通过启用部分图捕获、受保护的图捕获和即时重新捕获来改变这一点。

Dynamo 和AOT Autograd然后将优化的 FX 图传递给 PyTorch 本机编译器级别 TorchInductor。硬件公司也可以将此图输入到他们自己的后端编译器中。

TorchInductor

TorchInductor 是一个 python 原生深度学习编译器，可以为多个加速器和后端生成快速代码。Inductor 将采用具有约 250 个算子的 FX 图，并将它们降低到约 50 个算子。

Inductor 然后进入调度阶段，在该阶段融合运算符，并确定内存规划。

Inductor 然后进入“Wrapper Codegen”，它生成在 CPU、GPU 或其他 AI 加速器上运行的代码。包装器 codegen 取代了编译器堆栈的解释器部分，可以调用内核和分配内存。后端代码生成部分利用适用于 GPU 的 OpenAI Triton 并输出 PTX 代码。对于 CPU，英特尔编译器生成 C++（也适用于非英特尔 CPU）。

未来他们将支持更多硬件，但关键是 Inductor 大大减少了编译器团队在为其 AI 硬件加速器制作编译器时必须做的工作量。此外，代码针对性能进行了更优化。显着降低了内存带宽和容量要求。

分析人士表示，我们不想构建只支持 GPU 的编译器。我们想要一些可以扩展以支持各种硬件后端的东西，并且拥有 C++ 和 [OpenAI] Triton 会强制实现这种通用性。

OpenAI 的Triton

OpenAI 的 Triton 对 Nvidia 的机器学习闭源软件护城河具有颠覆性的角度。Triton 直接采用 Python 或通过PyTorch Inductor 堆栈提供数据。后者将是最常见的用例。Triton 然后将输入转换为 LLVM 中间表示，然后生成代码。对于 Nvidia GPU，它直接生成 PTX 代码，跳过 Nvidia 的闭源 CUDA 库（例如 cuBLAS），转而使用开源库（例如 cutlass）。

CUDA 通常被专门从事加速计算的人员使用，但在机器学习研究人员和数据科学家中却鲜为人知。高效使用可能具有挑战性，并且需要深入了解硬件架构，这可能会减慢开发过程。因此，机器学习专家可能依赖 CUDA 专家来修改、优化和并行化他们的代码。

Triton 弥合了使高级语言能够实现与使用低级语言的语言相当的性能的差距。Triton 内核本身对于典型的 ML 研究人员来说非常清晰，这对于可用性来说非常重要。Triton 在 SM 中自动执行内存合并、共享内存管理和调度。Triton 对逐元素矩阵乘法不是特别有用，这已经非常有效地完成了。Triton 对于昂贵的逐点操作和减少更复杂操作的开销非常有用，例如Flash Attention涉及矩阵乘法作为较大融合操作的一部分。

OpenAI Triton 目前仅正式支持 Nvidia GPU，但在不久的将来这种情况会发生变化。未来将支持多个其他硬件供应商，这个开源项目正在获得令人难以置信的动力。其他硬件加速器直接集成到作为 Triton 一部分的 LLVM IR 的能力大大减少了为新硬件构建 AI 编译器堆栈的时间。

Nvidia 庞大的软件组织缺乏远见，无法利用其在 ML 硬件和软件方面的巨大优势，成为机器学习的默认编译器。他们缺乏对可用性的关注，这使得 OpenAI 和 Meta 的外部人员能够创建可移植到其他硬件的软件堆栈。