英伟达在机器学习领域的CUDA垄断是如何打破的-OpenAI Triton和PyTorch 2.0(1)

发布于: 雪球转发:0回复:0喜欢:0

主要观点:$英伟达(NVDA)$ 庞大的软件组织缺乏远见,无法利用他们在机器学习硬件和软件方面的巨大优势,成为机器学习的默认编译器。他们缺乏对可用性的关注,这使得OpenAI和Meta的局外人能够创建一个可移植到其他硬件的软件堆栈。为什么他们没有为机器学习研究人员构建一个像Triton这样的“简化”CUDA?像闪存注意力这样的东西,为什么它来自博士生而不是英伟达

在过去的十年里,机器学习软件开发的格局发生了重大变化。许多框架来来去去,但大多数框架都严重依赖于利用英伟达的CUDA,并在英伟达图形处理器上表现最佳。

然而,随着PyTorch 2.0和OpenAI的Triton的到来,英伟达在这一领域的主导地位,主要是由于其软件护城河,正在被打破

这份报告将涉及以下主题,如为什么谷歌的TensorFlow输给了PyTorch,为什么谷歌还没有能够公开利用其在人工智能领域的早期领导地位,机器学习模型训练时间的主要组成部分,内存容量/带宽/成本墙,模型优化,为什么其他人工智能硬件公司到目前为止还没有能够削弱英伟达的主导地位,为什么硬件将开始变得更加重要,英伟达在CUDA的竞争优势是如何被抹去的,以及英伟达的一个竞争对手在训练硅的大型云上取得的重大胜利。

一个直接的总结是,机器学习模型的默认软件堆栈将不再是英伟达的闭源CUDA。球在英伟达的球场上,他们让OpenAI和Meta控制软件堆栈。那个生态系统因为英伟达的专有工具失败而建立了自己的工具,现在英伟达的护城河将被永久削弱。

TensorFlow与PyTorch

几年前,框架生态系统相当分散,但TensorFlow是领跑者。谷歌看起来准备控制机器学习行业。他们拥有先行者优势,拥有最常用的框架TensorFlow,并设计/部署了唯一成功的人工智能application-specific加速器TPU。

然而,PyTorch赢了。谷歌未能将其先发优势转化为对新兴机器学习行业的主导地位。如今,谷歌在机器学习社区中有些孤立,因为它缺乏对PyTorch和图形处理器的使用,而倾向于自己的软件堆栈和硬件。在典型的谷歌风格中,他们甚至有一个名为Jax的第二框架,直接与TensorFlow竞争。

人们甚至无休止地谈论谷歌在搜索和自然语言处理领域的主导地位因大型语言模型而减弱,尤其是那些来自OpenAI和各种利用OpenAI API或正在构建类似基础模型的初创公司的语言模型。虽然我们认为这种厄运和悲观被夸大了,但这种故事是另一天的事了。尽管存在这些挑战,谷歌仍然处于最先进机器学习模型的最前沿。他们发明了transformers,并在许多领域保持着最先进的水平(PaLM、LaMBDA、Chinchilla、MUM、TPU)。

回到PyTorch获胜的原因。虽然有一些从Google手中夺走控制权的因素,但主要是由于PyTorch与TensorFlow相比提高了灵活性和可用性。如果我们将其归结为第一个主要级别,PyTorch与TensorFlow的不同之处在于使用“Eager mode”而不是“Graph Mode”。

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

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

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

虽然TensorFlow现在默认有Eager模式,但研究界和大多数大型科技公司已经围绕PyTorch定居下来。这一点的例证是,几乎所有的生成人工智能模型都基于PyTorch成为新闻。谷歌的生成人工智能模型基于Jax,而不是TensorFlow。

当然,还有很多使用TensorFlow和Keras等其他框架的图像网络,但是新模型开发的计算预算都流向了PyTorch模型。有关PyTorch获胜的更深层次的解释,请参阅此处。一般来说,如果你在NeurIPS(主要的人工智能会议)的大厅里走来走去,所有的生成人工智能,非谷歌的工作都是用PyTorch完成的。

机器学习训练组件

如果我们将机器学习模型训练归结为最简单的形式,机器学习模型的训练时间有两个主要的时间分量。

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

2、内存(带宽):等待数据或层权重到达计算资源。bandwidth-constrained操作的常见示例是各种归一化、逐点操作、SoftMax和ReLU。

过去,机器学习训练时间的主导因素是计算时间,等待矩阵乘法。随着英伟达图形处理器的不断发展,这很快就不再是主要关注点

Nvidia的FLOPS通过利用摩尔定律增加了多个数量级,但主要是架构变化,如张量核心和低精度浮点格式。相比之下,内存没有遵循相同的路径

如果我们回到2018年,当时BERT模型是最先进的,英伟达V100是最先进的图形处理器,我们可以看到矩阵乘法不再是提高模型性能的主要因素。从那以后,最先进的模型在参数数量上增长了3到4个数量级,最快的图形处理器在FLOPS上增长了一个数量级。

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

存储墙

随着模型规模的不断飙升,大型语言模型仅就模型权重而言就需要100s千兆字节,如果不是TB的话。百度和Meta部署的生产推荐网络需要数十兆字节的内存来处理其庞大的嵌入表。大型模型训练/推理的大部分时间不是花在计算矩阵乘法上,而是等待数据到达计算资源。显而易见的问题是架构师为什么不把更多的内存放在离计算更近的地方。答案是$$$。

内存遵循从紧密和快速到缓慢和廉价的层次结构。最近的共享内存池在同一芯片上,通常由SRAM制成。一些机器学习专用集成电路试图利用巨大的SRAM池来保持模型权重,但这种方法存在问题。即使是脑电图(Cerebras)价值约250万美元的晶圆级芯片,芯片上也只有40GB的SRAM。没有足够的内存容量来保持100B+参数模型的权重。

英伟达的架构一直在芯片上使用更少的内存。当前一代A100有40MB,下一代H100有50MB。台积电5nm工艺节点上的1GB SRAM将需要约200mm^2的硅。一旦实施相关的控制逻辑/结构,这将需要超过400mm^2的硅,或约占英伟达数据中心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纳秒),但它也便宜得多(1美元GB对100美元GB)。

几十年来,动态随机存取存储器一直遵循摩尔定律。当戈登·摩尔创造这个术语时,英特尔的主要业务是动态随机存取存储器。他对动态随机存取存储器晶体管密度和成本的经济预测一直持续到2009年。然而,自2012年以来,动态随机存取存储器的成本几乎没有提高。

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

虽然容量是一个重要的瓶颈,但它与另一个主要瓶颈——带宽——密切相关。增加的内存带宽通常是通过并行性获得的。虽然今天标准DRAM每GB只有几美元,但为了获得机器学习所需的大量带宽,英伟达使用HBM内存,这是一种由DRAM的3D堆叠层组成的设备,需要更昂贵的封装。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多播(L2 multicast)。这个想法是不同的SM(think cores)可以直接写入另一个SM的SRAM(共享内存/L1缓存)。这有效地增加了缓存的大小,并减少了DRAM读/写所需的带宽。未来的架构将依赖于向内存发送更少的操作,以最大限度地减少内存墙的影响。应该注意的是,更大的模型往往会实现更高的利用率,因为FLOPS要求以指数方式扩展,而内存带宽和容量需求往往会以更线性的方式扩展。

Operator融合-解决方法

就像训练机器学习模型一样,知道你所处的状态可以让你缩小对重要优化的范围。例如,如果你把所有的时间都花在内存传输上(即你处于内存带宽受限的状态),那么增加GPU的FLOPS不会有帮助。另一方面,如果你把所有的时间都花在执行大型笨重的matmuls(即计算受限的状态)上,那么将模型逻辑重写为C++以减少开销也不会有帮助。

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

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

这种优化通常涉及编写自定义CUDA内核,但这比使用简单的python脚本要困难得多。作为一种内置的折衷方案,随着时间的推移,PyTorch在PyTorch中本地稳步实现了越来越多的运算符。其中许多运算符只是将多个常用操作融合到一个更复杂的函数中。

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

我们会说软件开发人员很懒,但老实说,几乎所有人都很懒。如果他们习惯了PyTorch中的一个新运算符,他们会继续使用它。开发人员可能甚至没有意识到性能的提高,而是使用那个运算符,因为这意味着编写更少的代码。

此外,并非所有操作都可以融合。决定融合哪些操作以及将哪些操作分配给芯片和集群级别的特定计算资源通常要花费大量时间。在哪里融合哪些操作的策略虽然通常相似,但确实因架构而异。