Horace He：为万亿乘万亿次运算构建机器学习系统

我在 Meta 的 PyTorch 编译器团队工作。今天我想解释，我们如何为涉及万亿乘万亿次浮点运算的工作负载构建机器学习系统，范围从单个 GPU 到极大规模的分布式运行。

基础设施建设的规模非同寻常：新闻报道着庞大的 GPU 集群和新的电力安排，领先的 AI 公司则筹集数十亿美元。大量投资最终都用于支持惊人数量的浮点运算，尤其是重复数百万次迭代的矩阵乘法。

Epoch AI 的一张图展示了训练算力的增长。我估计领先模型正接近约 10^26 次浮点运算。这些数字背后有许多层软件，从模型 API 到 PyTorch、Triton、CUDA 和编译器工具链。

基础设施工作很有成就感，因为在广泛使用的一层做出一点改进，就能让许多人受益。但就像简单接口背后的语言模型一样，其下的基础设施也可能是我们并不完全理解的庞大系统，由我们或许从未见过的人构建。

机器学习系统与许多其他系统不同：模型逻辑往往很简单，人们对性能的期望却极高。《The Bitter Lesson》认为，能够随算力扩展的简单方法，一再胜过建立在更专门的巧思之上的方法。

Andrej Karpathy 的 llama2.c 展示了这种简单性：一个 Llama 2 实现只需不到一千行 C 代码、没有依赖，其中包括循环和矩阵乘法。这个版本运行得并不快，但它的简洁说明了庞大计算的基本内容。

一个常用指标是模型 FLOP 利用率，即模型达到硬件理论浮点运算峰值的比例。CPU 要达到 100%，意味着所有核心始终发出最大宽度的 SIMD 指令。普通 CPU 程序远达不到这个水平，而大型机器学习训练通常能达到约 50%。

这个领域也在不断集中。Transformer 已在语言、视觉和音频领域占据主导，训练前沿模型的机构相对较少。这改变了基础设施提供者与运行最大规模任务的用户之间的关系。

改进系统有两种方式：为所有人优化底层机制，或者为用户提供编程模型和工具，让他们表达并调优自己的工作。机器学习框架的历史表明，这种选择可能非常重要。

Caffe 等早期框架使用包括 Protobuf 文件在内的声明式描述来指定网络。手写这些描述有很多限制，于是 TensorFlow 1 等图构建 API 让程序员通过领域专用语言构建计算图。

图构建器仍然让用户编写的内容与 GPU 上实际执行的内容之间存在距离。PyTorch 在 2016 至 2017 年左右凭借命令式的即时执行模型取得成功：调用一个操作，从概念上说它就会运行。这比函数经过中间图和一系列变换更容易理解。

我有时开玩笑说，如果只给一天时间训练模型，用 Python 会比用 C++ 更快得到结果。原因不是 Python 语言本身运行更快，而是它简单、可扩展的执行模型让人们围绕它建立了庞大的机器学习基础设施生态。

PyTorch 的设计处于类似位置：简单的执行模型不会自动优化一切，却让其他人很容易在其上构建框架。早期，与基于计算图的竞争者相比，这种简单性并不一定损失多少性能。

原因之一是矩阵乘法占了超过 90% 的运行时间，而 cuBLAS 和 cuDNN 等厂商库已经对其进行了高度优化。框架编译器能优化的其他工作相对很少。

另一个原因是异步执行。CPU 将 GPU 工作安排进队列，GPU 同时处理队列。我会想象 Gromit 在行驶中的火车前铺轨：只要 Python 调度工作的速度快于 GPU 消耗工作的速度，调度开销实际上就能被隐藏。

2017 年 NVIDIA 推出 Tensor Cores 后，这种平衡改变了。这些 GPU 单元直接执行矩阵乘法，使矩阵乘法吞吐量与其他 GPU 操作之间的差距大幅扩大。矩阵乘法加速后，程序其余部分占据的运行时间比例就更大了。

机器学习编译器随之出现，以优化剩余工作。用户仍按命令式编程模型编写代码，但实现会捕获计算图并进行优化。例如，torch.compile 在 Python 字节码解释器层拦截执行，有点像专为 PyTorch 程序设计的即时编译器。

从历史上看，图构建系统既向程序员暴露计算图，也在执行时使用计算图；早期 PyTorch 暴露并执行即时操作。现代框架通常保留即时执行的用户接口，同时在内部捕获计算图，尽管 PyTorch、JAX、TinyGrad 和 MLX 的捕获方式不同。

我把 GPU 运行时间大致分成三类：计算、GPU 内部的内存搬运，以及 GPU 空闲时间等开销。这种划分有助于判断哪类优化能真正改善特定工作负载。

在现代硬件上，矩阵乘法可能比其他算术运算快得多。在我的 GPU 例子中，FP32 算力约为 67 TFLOPS，而面向矩阵乘法的 TF32 算力约为 1,000 TFLOPS。因此，即使非矩阵乘法运算完全占满硬件，与矩阵乘法峰值相比也可能显得很差。

一个意外的细节是，即使内存访问模式和控制流完全相同，矩阵里的值也可能影响性能。使用全零、高斯采样和均匀分布数据的基准测试，可能得到不同结果。

原因在于功耗。晶体管切换会消耗动态功率，不同位模式导致的切换次数也不同。硬件达到功率或温度上限时可能降频，因此全零矩阵可能比引发更多切换的数据运行得更快。

这可能让基准测试产生误导。有人说模型产生 NaN 后突然看起来更快了：大量张量值变成 NaN 后，硬件切换方式改变，降频也减少了。这种情况下吞吐量提高是计算出错的症状，并非真正的优化。

内存搬运是第二项主要成本。在一个类似 BERT 的工作负载中，张量收缩占浮点运算的 99.8%，却只占运行时间的 61%。算术运算量相对很少的操作仍然可能占用大量时间。

即使数据已经位于 GPU 显存中，也必须搬到计算单元再搬回来。可以想象一个空间有限的工厂和远处更大的仓库：在两者之间搬运物资会消耗带宽和时间，尽管搬运本身不做有用的算术运算。

假设我们分别执行加法、激活函数和另一个逐元素操作。在即时执行模式下，每项操作都可能让中间数据从内存到计算单元再返回。把结果送回仓库，只为了下一步立刻取回来，是一种浪费。

算子融合把多项操作组合到一个 GPU 内核中：只把数据送到计算单元一次，执行多个操作，再写回结果。我认为这是机器学习编译器中最重要的优化。它需要跨操作观察程序，而最简单的即时执行模型不会这样做。

融合也涉及保存值和重新计算值之间的取舍。利用附近已有的数据重新计算，可以避免内存访问。尽管增加了算术运算，它仍可能降低内存峰值并改善运行时间。

反向传播使这个问题尤其重要。训练先按层执行前向传播，再按相反顺序执行反向传播。两个过程之间可能需要保留中间结果，也就是激活值；它们的生命周期较长，可能导致内存不足。

第三类是开销。如果 CPU 无法足够快地将工作加入队列，GPU 就会闲等，就像火车追上了铺轨的人。CUDA Graphs 和开销更低的生成包装代码，是减少这类调度开销的办法。

既然只有少数占主导地位的架构，一次训练又要花费数十亿美元，人们可能会问：为什么还需要编译器？为什么不安排专家手工优化模型？事实上，有些团队确实避开编译器，他们的理由揭示了一个真正的易用性问题。

编译器能放大工作成果：实现一次优化，就有许多用户受益。但试想一个库：它并不总能工作，没有文档说明何时能工作，更新后行为还可能改变。如果这里的“工作”指达到所需性能，编译器优化有时就会给人这种感觉。

一篇关于 ISPC 的有影响力的讨论指出：自动向量化不是一种编程模型。如果向量化可能失败，依赖 SIMD 性能的程序员就必须弄清失败原因，调整代码，或想办法让编译器应用优化。

明确的编程模型能给用户更可靠的依据。SIMD 内建函数指定了映射到 SIMD 指令的操作；用户可以根据这个约定推理，而不必逆向理解优化器。一个始终有效的优化可以成为编程模型的一部分。

编译器通常把复杂实现隐藏在简单接口后面。当它没有应用某项优化，用户就不得不面对这些实现细节，努力弄清代码为什么没有内联或向量化。在机器学习中，数值精度让情况更加微妙。

硬件不断向更低精度的算术运算发展，从 16 位到 8 位，再到较新硬件中的 4 位格式。微小的数值差异可能产生巨大影响；单独看起来更好的优化，也可能破坏算法所依赖的假设。

我曾在一个 FlashAttention 实现中遇到与融合乘加（FMA）有关的 NaN。FMA 先完成乘法和加法，最后只舍入一次，结果通常比分别舍入更精确。但这一改变可能破坏依赖两次计算得到完全相同舍入值的代码。

该计算在求指数之前减去最大值，以避免过大的指数。如果最大值来自已经舍入的乘积，而后面的表达式只在一条分支使用 FMA，两个表达式就可能以算法未预期的方式出现差异。在这个例子中，开启 FMA 会产生 NaN，关闭它则不会。

编译器也难以处理代数重写。FlashAttention 依赖一种在线 softmax 改写，去掉一次全局同步。这对人来说是可以处理的数学问题，但通常不能指望编译器从一串普通操作中自行发明这种改写。

注意力计算可以写成一次矩阵乘法、一次 softmax，再接一次矩阵乘法。编译器可以尝试匹配这一模式，并替换成融合内核。但如果用户编写自己的 softmax，就可能无意中破坏匹配，导致内存占用突然大增、执行变慢。

一种应对方法是提供单一的融合注意力算子。用户因此能获得可预测的快速操作，但单体 API 很难表达每个新想法。研究人员不断提出滑动窗口、ALiBi、分页和邻域注意力等变体。

随着变体增加，融合算子会不断积累标志位和特殊情况；已有的特殊行为也不能简单删除，否则会破坏用户代码。即使选项很多，仍会有人需要算子无法表达的注意力变体。

我们不必只在不可预测的优化器和缺乏灵活性的单体方案之间选择。更实用的编程模型可以把复杂的手写 FlashAttention 内核，与编译器能够自动生成的简单用户修改结合起来。

FlexAttention 就是一个例子。它使用 torch.compile，同时为用户提供可靠的融合注意力内核及其内存特性。用户可以在这个模型内尝试不同方案，而不必了解底层编译器的每个实现细节。

对于滑动窗口因果掩码，用户可以检查查询位置是否在键值位置之后或与之相同，以及两者距离是否处于窗口内，再组合这些条件。内核机制负责实现，用户提供变化部分。

在这种方法出现之前，对注意力掩码、偏置、训练和推理的支持由手工实现的内核拼凑而成。灵活的抽象能覆盖更多组合。有位用户甚至表达了不同大小分子图的掩码，这是我设计 API 时没有想到的用途。

我常用一个数学类比来思考编程模型：与其更用力地敲核桃，不如改变解决问题的方式，让它更容易打开。换一种抽象，可以把脆弱的优化变成用户能够依赖的属性。

看看 CUDA 的并行性。问编译器什么时候会自动让 CUDA 程序并行，其实没有抓住重点：并行执行已经内置于其编程模型中。程序仍可能出错或死锁，但并行性不是一次可能成功也可能失败的自动向量化决策。

扩展到单个 GPU 之外会带来另一组问题。传统分布式服务通常扩展许多小型、可互换的查询，也可能为可用性接受额外硬件成本。大型训练更像一个巨大的任务，需要频繁全局同步，对性能的约束也非常严格。

组织并行的一种思路是想象一个工作立方体：批次维度上的不同任务、任务内部的操作，以及时间。数据并行、张量并行和流水线并行分别沿着这幅图的不同维度划分工作。

数据并行把不同数据或任务分配给不同 GPU，但每一步之后都必须同步梯度。我们不能无限增大批次来避免通信：批次过大可能使训练无法按预期收敛。完全分片数据并行也避免了简单地复制所有参数。

张量并行将同一个任务拆分到多个 GPU。由于设备协作处理同一任务，通信与计算往往更难重叠；但硬件空闲造成的性能损失又使重叠很有价值。

异步张量并行可以在一次操作内部创建小型流水线。它把通信和矩阵乘法划分成多个部分，让一部分数据传输时，另一部分已经开始计算。这样即使在张量并行操作内部，也能实现通信与计算的重叠。

流水线并行则在不同时间把任务的连续部分分配给不同 GPU。全局同步可能在流水线中留下空档，而反向传播进一步增加了调度选择，因为必须同时安排前向和反向微批次。

训练流水线在某些时刻可以选择运行前向或反向微批次。不同调度方案利用这种选择减少空闲时间。Llama 3 的训练图结合了数据、张量和流水线并行，以及上下文并行。

编译器过去一直难以跟上分布式优化的进展。在一组固定选项内，自动模型或许能选出好的配置。但许多进展来自发明新的并行维度，扩大搜索空间，而不只是搜索原有空间。

容错在大规模运行时成为另一项约束。数千个 GPU 全局同步工作时，硬件、网络或用户相关故障都可能中断整个任务。对于较小集群看似可以接受的故障率，在集群扩大后会成为频繁的障碍。

Llama 3 的容错数字说明了这个问题：他的例子中，约 16,000 个 GPU 每隔 1.8 小时出现一次错误，而 131,000 个 GPU 大约每 15 分钟就会出一次错误。当规模足够大时，一次训练步骤还没完成就可能出现故障。

机器学习已经成为极其重要的计算工作负载，但它的社会和技术特征意味着传统系统和编译器方法并不总能直接套用。工作负载不断演变；编译器团队不能闭门优化数年，忽视这些变化。

因此，最有意思的问题不仅是实现哪些优化，还包括什么样的编程模型能给人们可靠的工具，让他们表达和开发自己的优化。这些模型可以帮助用户构建下一代超大规模训练系统。