十年来，机器学习软件开发的格局发生了重大变化。许多框架如雨后春笋般涌现，但大多数都严重依赖于英伟达的 CUDA，并在英伟达的 GPU 上才能获得最佳的性能。然而，随着 PyTorch 2.0 和 OpenAI Triton 的到来，英伟达在这一领域的主导地位正在被打破。
谷歌早期在机器学习模型架构、训练、模型优化方面都具有很大优势，但现在却难以充分发挥这些优势。而在硬件方面，其他 AI 硬件公司很难削弱英伟达的统治地位。直到 PyTorch 2.0 和 OpenAI Triton 出现，机器学习模型的默认软件堆栈将不再是英伟达的闭源 CUDA。
TensorFlow vs. PyTorch
类似的竞争出现在机器学习框架中。几年前，框架生态系统相当分散，但 TensorFlow 是领跑者。表面看来，谷歌稳操机器学习框架行业，他们凭借 TensorFlow 设计了 AI 应用特定加速器 TPU，从而获得了先发优势。
然而现在看来，PyTorch 赢了，谷歌未能将其先发优势转化为新兴 ML 行业的主导地位。如今，谷歌在机器学习社区中似乎有些被孤立了，因为它没有使用 PyTorch 和 GPU，而是使用自己的软件堆栈和硬件。甚至，谷歌研发了第二个机器学习框架 ——JAX，直接与 TensorFlow 竞争，这是典型的「谷歌行为」。
一些人认为，由于大型语言模型的兴起，特别是 OpenAI 的大型语言模型和各种利用 OpenAI API 构建的语言模型，谷歌在搜索和自然语言处理方面的主导地位正在减弱。也许这种观点过于悲观，毕竟当前大多数模型的基础架构仍是谷歌开发的 transformer。
那么，为什么 PyTorch 能大获全胜呢？主要原因是相对于 TensorFlow，PyTorch 具有更高的灵活性和可用性。PyTorch 与 TensorFlow 主要的不同之处在于使用 Eager 模式而非 Graph 模式。
Eager 模式可以说是标准的脚本执行方法，与普通的 Python 代码没什么区别。这使得调试和理解代码更加容易，因为用户可以看到中间操作的结果，以及模型是如何运行的。
相反，Graph 模式分为两个阶段。第一阶段代表要执行操作的计算图，其中的节点代表操作或变量，而节点之间的边表示它们之间的数据流。第二阶段是延迟执行计算图的优化版本。
这种两阶段的方法使得理解和调试代码更具挑战性，因为在图执行结束之前用户无法看到发生了什么。这类似于「解释型」与「编译型」语言，如 Python 与 C++，调试 Python 更容易，因为它是解释型语言。
虽然 TensorFlow 现在也默认使用 Eager 模式，但研究社区和大多数大型科技公司都选择使用 PyTorch。
机器学习训练组件
如果将机器学习模型训练简化为最简单的形式，影响机器学习模型训练的因素主要有两点：
- 计算 ( FLOPS ) ：在每一层内运行密集矩阵乘法；
- 内存带宽。
以前，影响机器学习训练时间的主要因素是计算时间，等待系统执行矩阵乘法。随着英伟达 GPU 的不断发展，这很快就不再是主要问题了。
英伟达利用摩尔定律将 FLOPS 提高了多个数量级，但主要是架构变化 —— 张量核（tensor core）和更低精度的浮点数格式。相比之下，存储方面没有太大的变化。
2018 年，最先进的模型是 BERT，英伟达 V100 是最先进的 GPU，那时矩阵乘法已经不再是提高模型性能的主要因素。之后，模型在参数数量上增长了 3 到 4 个数量级，而最快的 GPU 在 FLOPS 上增长了 1 个数量级。
即使在 2018 年，纯计算绑定的工作负载也占 FLOPS 的 99.8%，但仅占运行时的 61%。与矩阵乘法相比，归一化和逐点运算（pointwise ops）使用的 FLOPS 仅为矩阵乘法的 1/250 和 1/700，但它们消耗了近 40% 的模型运行时间。
内存墙
随着模型规模的不断飙升，大型语言模型（LLM）仅用于模型权重的内存就需要 100 GB 以上。百度和 Meta 部署的产品推荐网络需要数十 TB 的内存来存储其海量嵌入表。大型模型训练 / 推理中的大部分时间都没有花在计算矩阵乘法上，而是在等待数据传输。显然，问题在于为什么架构师不将更多内存放在更靠近计算的位置，问题的答案也是可想而知的 —— 成本。
最近的共享内存池一般是同一芯片上的 SRAM。一些机器学习 ASIC 试图利用巨大的 SRAM 池来保存模型权重。但即使是 Cerebras 的价值约 5000000 美元的晶圆级芯片也只有 40GB 的 SRAM。内存容量不足以容纳 100B+ 参数模型的权重。
英伟达设计的芯片片上内存容量要少得多 ——A100 有 40MB，H100 有 50MB。台积电 5 纳米芯片上的 1GB SRAM 需要约 200 平方毫米的硅，要实现相关的控制逻辑 / 结构，将需要超过 400 平方毫米的硅。鉴于 A100 GPU 的成本在 1 万美元以上，而 H100 更是接近 2 万美元，从经济角度看，这种方法是不可行的。即使忽略英伟达在数据中心 GPU 上约 75% 的利润率，对于完全量产的产品，SRAM 内存的成本仍在 100 美元 / GB 左右。
此外，片上 SRAM 存储器的成本不会随着传统摩尔定律工艺技术的缩小而降低太多。同样的 1GB 内存，采用台积电下一代 3nm 制程工艺，成本反而更高。虽然 3D SRAM 将在一定程度上帮助降低 SRAM 成本，但这只是暂时的。
内存层次结构的下一步是紧密耦合的片外内存 DRAM。DRAM 的延迟比 SRAM 高一个数量级（约 100nsVS10ns），但它也便宜得多。几十年来，DRAM 一直遵循着摩尔定律。戈登摩尔创造这个词时，英特尔的主要业务就是 DRAM。他对晶体管密度和成本的预测在 2009 年之前对 DRAM 普遍适用。不过自 2012 年以来，DRAM 的成本几乎没有改善。
然而，人们对内存的需求只增不减。DRAM 现在占服务器总成本的 50%，逐渐形成所谓的「内存墙」。将英伟达 2016 年的 P100 GPU 与最新的 H100 GPU 进行比较，我们会发现内存容量增加到原来的 5 倍（16GB → 80GB），FP16 性能增加到 46 倍（21.2 TFLOPS → 989.5 TFLOPS）。
虽然内存容量是一个重要瓶颈，但另一个瓶颈 —— 内存带宽也非常关键。内存带宽的增加通常是通过并行性获得的。虽然如今标准 DRAM 的价格仅为几美元 / GB，但为了获得机器学习所需的海量带宽，英伟达使用 HBM 内存 —— 一种由 3D 堆叠 DRAM 层组成的设备，需要更昂贵的封装。HBM 的成本大概是 10-20 美元 / GB，包括封装和产量成本。
内存带宽和容量的成本限制问题在英伟达的 A100 GPU 中尤为明显。如果不进行大量优化，A100 只能具有非常低的 FLOPS 利用率。
即使研究人员进行了大量优化，大型语言模型的 FLOPS 利用率也只能达到 60% 左右 。很大一部分时间花费在等待来自另一个计算 / 内存的数据，或者及时重新计算结果以减少内存瓶颈。
从 A100 到 H100，FLOPS 增长到原来的 6 倍以上，但内存带宽仅增长到 1.65 倍。这导致许多人担心 H100 的利用率会很低。A100 需要很多技巧才能绕过内存墙，而 H100 还需要实现更多技巧。
H100 为 Hopper 架构带来了分布式共享内存和 L2 多播 ( multicast ) 。这个想法让一个 SM 中的数据可以直接写入另一个 SM 的 SRAM（共享内存 /L1 Cache）中。这有效地增加了缓存的大小并减少了 DRAM 读 / 写所需的带宽。未来的架构将减少向内存发送的操作数量，以最大限度地减少内存墙的影响。值得注意的是，较大的模型往往会实现更高的利用率，因为 FLOPS 需要按参数量的三次方成比例扩展，而内存带宽和容量需求往往按二次方扩展。
算子融合
如果将所有时间都花在内存传输上（即处于内存带宽限制状态），那么增加 GPU 的 FLOPS 将无济于事。另一方面，如果将所有时间都花在执行大型 matmuls 上，那么即使将模型逻辑重写为 C++ 来减少开销也将无济于事。
PyTorch 之所以能胜过 TensorFlow，就是因为 Eager 模式提高了灵活性和可用性，但转向 Eager 模式并不是只有好处。在 Eager 模式下运行时，每次运算都要从内存中读取、计算，然后在处理下一次运算之前发送到内存。如果不进行大量优化，这会显著增加内存带宽需求。
因此对于在 Eager 模式下执行的模型，有一种主要的优化方法是算子融合。融合运算在一次传递中计算多个函数，以最小化内存读 / 写，而不是将每个中间结果写入内存。算子融合改善了运算符调度、内存带宽和内存大小成本。
这种优化通常涉及编写自定义 CUDA 内核，但这比使用简单的 Python 脚本要难得多。随着时间的推移，PyTorch 中稳定地实现了越来越多的算子，其中许多算子只是简单地将多次常用运算融合到一个更复杂的函数中。
算子的增加让在 PyTorch 中创建模型变得更容易，并且由于内存读 / 写更少，Eager 模式的性能更快。缺点是 PyTorch 在几年内激增到了 2000 多个算子。
我们可以说软件开发人员太懒了，但说实话，又有谁没懒惰过呢。一旦习惯了 PyTorch 中的一个新算子，他们就会继续用它。开发人员甚至可能没有意识到性能在提高，而是继续使用该算子，因为这样就不用编写更多的代码。
此外，并非所有算子都可以融合。决定要融合哪些运算，将哪些运算分配给芯片和集群级别的特定计算资源都需要花费大量的时间。算子在何处融合的策略虽大体相似，但因为架构的不同也会有很大差异。
英伟达曾是王者
算子的增长和默认的地位对英伟达来说是优势，因为每个算子都针对其架构进行了快速优化，但并未针对任何其他硬件进行优化。如果一家 AI 硬件初创公司想要全面实施 PyTorch，那就意味着以高性能支持不断增长的 2000 个算子列表。
因为提取到最大性能需要很多技巧，在 GPU 上训练具有高 FLOPS 利用率的大型模型所需的人才水平越来越高。Eager 模式执行加算子融合意味着开发的软件、技术和模型都在不断地被推动，以适应当前一代 GPU 具有的计算和内存比率。
每个开发机器学习芯片的人都受制于同一个内存墙。ASIC 受制于支持最常用的框架，受制于默认的开发方法、GPU 优化的 PyTorch 代码以及英伟达和外部库的混合。在这种情况下，避开 GPU 的各种非计算包袱而支持更多 FLOPS 和更严格的编程模型的架构意义不大。
然而，易用性第一。打破恶性循环的唯一方法是让在英伟达的 GPU 上运行模型的软件尽可能轻松无缝转移到其他硬件。随着模型架构的稳定和来自 PyTorch 2.0、OpenAI Triton 和 MLOps 公司（如 MosaicML）的抽象成为默认，芯片解决方案的架构和经济性开始成为购买的最大驱动力，而不是英伟达高级软件提供的易用性。
PyTorch 2.0
几个月前，PyTorch 基金会成立，并脱离了 Meta 。除了对开放式开发和治理模型的更改外，2.0 还发布了早期测试版本，并于 3 月全面上市。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 做出如此巨大的贡献，是因为他们希望在自己价值数十亿美元的 GPU 训练集群上以更少的努力实现更高的 FLOPS 利用率。这样他们也有动力使软件堆栈更易于移植到其他硬件，将竞争引入机器学习领域。
在更好的 API 的帮扶下，PyTorch 2.0 还可以支持数据并行、分片、pipeline 并行和张量并行，为分布式训练带来了进步。此外，它在整个堆栈中原生支持动态形状，在许多其他示例中，这更容易支持 LLM 的不同序列长度。下图是主要编译器首次支持从训练到推理的 Dynamic Shapes：
PrimTorch
对于除英伟达 GPU 之外的每个机器学习 ASIC 来说，为 PyTorch 编写一个完全支持所有 2000 多个算子的高性能后端并非易事。PrimTorch 将算子的数量减少到约 250 个原始算子，同时还保持 PyTorch 最终用户的可用性不变。PrimTorch 使 PyTorch 的不同非英伟达后端的实现变得更加简单和易于访问。定制硬件和系统供应商可以更轻松地推出他们的软件堆栈。
TorchDynamo
转向图模式需要可靠的图定义。为了实现这一转向，Meta 和 PyTorch 已经尝试了大约 5 年的时间，但是他们提出的每个解决方案都存在明显的缺点。最后，他们用 TorchDynamo 破解了这个难题。TorchDynamo 将摄取任何 PyTorch 用户脚本，包括调用外部第三方库的脚本，并生成 FX 图。
Dynamo 将所有复杂算子减少到 PrimTorch 中的约 250 个原始算子。一旦图形成，未使用的算子将被丢弃，图会决定哪些中间算子需要存储或写入内存、哪些可能被融合。这极大地减少了模型内的开销，同时对用户来说也是「无缝」的。
在测试的 7000 个 PyTorch 模型中，TorchDynamo 已经适用于 99% 以上的模型，包括来自 OpenAI、HuggingFace、Meta、英伟达、Stability.AI 等的模型，而无需对原始代码进行任何更改。测试的 7000 个模型是从 GitHub 上使用 PyTorch 的最受欢迎项目中随机挑选出来的。
谷歌的 TensorFlow/Jax 和其他图模式执行 pipeline 通常要求用户确保他们的模型适合编译器架构，以便可以捕获图。Dynamo 通过启用部分图捕获、受保护的图捕获和即时重新捕获来改变这一点。
部分图捕获允许模型包含不受支持的 / 非 python 构造。当无法为模型部分生成图时，将插入图中断，并且将在部分图之间以 eager 模式执行不支持的构造。
受保护的图捕获会检查捕获的图是否对执行有效。「保护」的意思是一种需要重新编译的更改。这很重要，因为多次运行相同的代码不会多次重新编译。如果捕获的图对于执行无效，则即时重新捕获允许重新捕获图。
PyTorch 的目标是创建一个具有流畅 UX 的统一前端，该前端利用 Dynamo 生成 graph。该解决方案的用户体验不会发生变化，但性能可以得到显着提升。捕获图可以在大量计算资源上更有效地并行执行。
随后，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 的编译器，而是想要支持各种硬件后端。
OpenAI Triton
对英伟达的机器学习闭源软件来说，OpenAI Triton 是一个颠覆性的存在。Triton 直接采用 Python 或通过 PyTorch Inductor 堆栈提供数据，后者是最常见的用法。Triton 负责将输入转换为 LLVM 中间表征，并生成代码。英伟达 GPU 将直接生成 PTX 代码，跳过英伟达的闭源 CUDA 库（如 cuBLAS），转而使用开源库（如 cutlass）。
CUDA 在加速计算领域很受欢迎，但在机器学习研究人员和数据科学家中却鲜为人知。使用 CUDA 可能会带来重重挑战，并且需要深入了解硬件架构，这可能导致开发过程变慢。因此，机器学习专家可能就要依赖 CUDA 专家来修改、优化和并行化他们的代码。
Triton 弥补了这一缺陷，使高级语言实现了与低级语言相当的性能。Triton 内核本身对典型的 ML 研究者来说非常清晰，这对可用性来说非常重要。Triton 在 SM 中自动执行内存合并、共享内存管理和调度。Triton 对逐元素矩阵乘法不是特别有用，但矩阵乘法已经可以非常高效地完成。Triton 对于成本高昂的逐点运算和减少复杂操作的开销非常有用。
OpenAI Triton 目前仅正式支持英伟达的 GPU，但在不久的将来会发生变化，将支持多个其他硬件供应商。其他硬件加速器可以直接集成到 Triton 的 LLVM IR 中，这大大减少了为新硬件构建 AI 编译器堆栈的时间。
英伟达庞大的软件体系缺乏远见，无法利用其在 ML 硬件和软件方面的巨大优势，也就没能成为机器学习的默认编译器。他们缺乏对可用性的关注，而 OpenAI 和 Meta 也正是得益于此才能够创建出可移植到其他硬件的软件堆栈。
原文链接：https://www.semianalysis.com/p/nvidiaopenaitritonpytorch
