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