新智元报道
编辑：Aeneas 英智
【新智元导读】Hugging Face发布了「超大规模实战手册」，在512个GPU上进行超过4000个scaling实验。联创兼CEO Clement对此感到十分自豪。
最近，Hugging Face发布了一个「超大规模训练手册」，教我们如何在GPU集群上训练LLM。
这项震撼的研究，在512个GPU上做了超过4000个Scaling实验，并测量了吞吐量（标记的大小）和GPU利用率（标记的颜色）。
HuggingFace联创兼CEO Clement表示，自己的理想，就是能生活在这样一个世界：所有的公司和组织无论大小，无论是否富有，都能训练自己的AI。
未来的世界，就应该是这个样子。
因此，当发布这份大规模训练手册时，他感到十分自豪。
网友们纷纷表示，这就是自己爱Hugging Face的原因。
未来我们应该拥有的就是民主化的AI和去中心化数据，让全球大脑共同工作。从此，AI不再是为1%的人服务，而是让99%的人参与其中！
本文将从基础入手，介绍如何将LLM训练规模从一块GPU扩展到数十块、数百块甚至数千块GPU。
随着训练集群规模不断扩大，数据并行、张量并行、流水线并行、上下文并行，以及ZeRO和内核融合等技术相继被提出，来确保GPU的高效利用。
在深入代码和实验之前，先要理解每种方法的工作原理、优缺点和适用场景。例如，LLM在训练过程中的哪些部分占用了最多的显存，训练过程中的哪个阶段会成为瓶颈。
同时，还会介绍如何通过并行来解决显存限制，提高吞吐量。
这样一来，就能理解下面这个用于计算Transformer模型显存占用的小工具是如何工作的。
除了理论分析，还提供了一个工具，用于预测训练过程中显存的实际使用情况：
本文运行了4100多次分布式实验，用了512块GPU，以探索可能的分布式训练架构和模型大小的影响。
概览
本文内容广泛，作者将内容总结如下：
训练过程有三个关键挑战：
显存占用：如果某个训练步骤超出了显存容量，训练就无法继续。
计算效率：希望GPU大部分时间都在执行计算，而不是浪费在数据传输或等待其他GPU执行任务。
通信开销：减少通信开销，以免GPU处于空闲状态。需充分利用节点内部和节点之间的带宽，尽量让通信和计算过程重叠进行，以提高训练效率。
在很多情况下，可以在计算、通信和显存中进行取舍，如通过重计算或张量并行，找到合适的平衡点。
在单个GPU上训练模型时，通常包含三个步骤：前向传播、反向传播和优化步骤。
在预训练中，批大小通常以token为单位。这使得训练的计算量通常与输入序列长度无关。
近期LLM的批大小通常在4M到60M个token。Llama 1的训练批大小约为4M个token，DeepSeek的训练批大小约为60M个token。
第一个挑战已经出现：「显存不足」。
训练时，显存需存储以下内容：模型权重、模型梯度、优化器状态和计算梯度所需的激活值。
如何根据这些变量，快速确定显存使用情况呢？一个简单的方法是通过实验测量。
分析显存使用情况
用PyTorch分析器，可以了解训练过程中显存的分配方式。显存利用率在训练过程中，会有很大的变化。
接下来，探讨如何在扩展训练规模的过程中，最大化计算效率，同时确保激活值、参数、梯度和优化器状态的显存需求在限制范围内。
对于一个简单的Transformer LLM，参数量可以用以下公式计算：
显存需求可通过总参数乘以每个参数占用的字节数来计算。出于稳定性及效率的考虑，通常采用混合精度。
接下来，估算模型的显存需求：
一旦模型参数达到7B，权重和优化器状态的显存需求就开始大幅增加，并超过典型GPU显存容量，例如H100的80G。
相比于权重、梯度和优化器状态，激活值的计算更加复杂，部分原因是它取决于模型的输入。
现在，介绍第一种技术——激活值重计算。
激活值重计算
激活值重计算的基本思想是在前向传播中丢弃一些激活值，以节省显存。并通过增加一些计算量，在反向传播中动态计算这些激活值。
使用重计算时，通常只在模型架构的几个关键点存储激活值，丢弃其余的激活值，并在反向传播中从最近保存的激活值开始重新计算它们。
选择要存储的关键激活值有全量和选择性等策略。接下来看到，重计算如何减少显存占用，以及如何在节省显存和增加计算成本之间取得良好的平衡。
对于规模较小的模型，长序列的激活值产生的影响更大，因此重计算的效果更显著。
如今，大多数训练框架都使用FlashAttention，原生集成了激活值重计算的优化策略。
激活值重计算会稍微增加浮点运算次数，但它减少了内存访问开销。这种权衡在GPU上特别有利，计算速度通常更快，同时降低了显存占用。
然而，激活值仍然与批大小呈线性相关，当批大小增加时，激活值的显存可能又会成为问题。
还有第二个方法，梯度累积来救场！
梯度累积
梯度累积是一种避免显存爆炸的方法，原理是将批量数据拆分为多个微批次，依次进行前向传播和反向传播。
通过梯度累积，全局批大小可通过以下公式计算：
梯度累积还可以与激活重计算相结合，进一步减少显存占用。
梯度累积能通过仅计算部分微批次，来减少激活值占用的显存。每个微批次的前向和反向传播可以并行运行，看来是时候考虑多个GPU了！
在扩展到多个GPU前，介绍分布式训练工具箱中最有用的工具之一：分析器。
PyTorch分析器
分析器能精确追踪和可视化训练过程中的情况，展示了：
CPU线程异步启动内核到GPU。
多个CUDA流并行处理计算和通信任务。
内核执行时间和内存分配。
首先介绍数据并行技术，它是梯度累积的并行版本。
数据并行
数据并行的核心思想是在多个GPU上运行，并在每个GPU上并行处理不同微批次的数据。
每个GPU上的梯度是不同的，为了让不同GPU上的模型保持同步，用all-reduce操作对模型的梯度进行平均。
由于不希望GPU处于空闲状态，应尽可能地让通信和计算同时进行。这里有三种优化方法：将梯度同步与后向传播重叠进行、梯度分桶和与梯度累积相结合。
重新审视全局批大小
结合新引入的数据并行和梯度累积参数来更新批大小：
给定一个目标全局批大小，可以通过调整梯度累积步数和并行进程数来加快训练速度。
当GPU数量超过限制时，吞吐量开始显著下降。
当数据并行达到一定规模后，会出现明显的通信开销瓶颈。对于更大的模型或者批大小，还有其他选择吗？
幸运的是，确实有解决方案。这些方案将一些张量移到CPU上，或将权重、梯度、优化器等张量拆分到多个GPU上。
拆分主要有两种方法：并行化（张量并行、上下文并向或流水线并行）和共享（如DeepSpeed Zero或PyTorch FSDP）。两种方法相互独立，也可以结合使用！
共享模式与数据并行密切相关，首先来研究ZeRO方法。
ZeRO（零冗余优化器）
DeepSpeed ZeRO是一种旨在减少LLM训练中内存冗余的优化技术。
数据并行是一种高效的方法，但在每个实例上简单复制优化器状态、梯度和参数会引入大量的内存冗余。
ZeRO通过在数据并行维度上对优化器状态、梯度和参数进行分区，消除了内存冗余。
ZeRO有三个优化阶段：
ZeRO-1：优化器状态分区
ZeRO-2：优化器状态+梯度分区
ZeRO-3（也称为FSDP）：优化器状态+梯度+参数分区
ZeRO的理念是在数据并行进程之间，对这些对象进行分区，每个节点仅存储一部分，需要时进行重建。
ZeRO-1：优化器状态分区
在ZeRO-1中，优化器状态被分成N_d等份，N_d是数据并行度。在优化步骤中，只有1/N_d的权重会被更新。
对梯度执行reduce-scatter操作。
与传统的数据并行相比，ZeRO-1将all-reduce梯度通信替换为reduce-scatter操作，并在优化器后添加了一个全参数的all-gather操作。
ZeRO-2：增加梯度分区
在反向传播中，不再对梯度执行all-reduce操作，而是仅执行reduce-scatter操作。只传播1/N_d的梯度，与ZeRO-1相比，节省了更多内存。
梯度分区后，内存有所减少。随着N_d的增加，最多可节省8倍内存。在通信方面，与ZeRO-1相比，唯一的区别在于可以动态释放内存。
ZeRO-3：增加参数分区
在第三阶段，将之前对优化器状态和梯度进行分区的做法，扩展到模型参数。
如果模型的所有部分都被分布式存储，在前向传播中，具体操作如下：
在进行前向传播时，遍历各层，获取所需的参数，并在不再需要时将其从内存中清除。反向传播与之类似，只是流程相反：
由于需要在前向传播和反向传播中持续进行all-gather操作，与ZeRO-2相比，每个训练步骤会增加（2×层数-1）次额外的操作。
借助ZeRO技术，可以在数据并行中，将参数、梯度和优化器状态进行分区。
然而，这里有一个限制，ZeRO无法对激活值内存进行处理。这部分内存会随着序列长度和批大小的增加而增加。
为克服这些问题，是时候探索一种新的并行方式了——张量并行。与严重依赖参数通信的ZeRO方法不同，张量并行提出将参数、梯度、优化器状态和激活值分布到多个设备上，而无需在各GPU之间进行模型参数的通信。
张量并行
当激活值内存占用超过预算时，ZeRO就会遇到瓶颈。张量并行是在ZeRO基础上，针对激活内存超预算问题的优化技术。
利用矩阵乘法特性，通过按列或按行分区，将张量分布到多个GPU上计算。列线性需广播输入矩阵、分割权重矩阵列，用all-reduce组合结果；行线性要分割权重矩阵行和输入，分别使用scatter和all-reduce操作。
张量并行能减少矩阵乘法激活内存，在多GPU间分布模型参数、梯度、优化器状态，使7B参数模型可在单节点8个GPU上运行。
缺点是跨节点通信慢，当张量并行度超过8个GPU时，通信开销明显，从TP=8到TP=16、TP=16到TP=32性能显著下降。层归一化和随机失活等操作仍需收集完整激活值。
序列并行
为解决层归一化和随机失活需完整激活值的问题，引入序列并行技术。
序列并行的优势是减少最大激活值存储大小，仅使用张量并行时需存储形状为 (b,s,h) 的激活值，使用序列并行后可减少到 。
这有助于节省激活值内存，能增大批大小和序列长度，如在70B参数模型中，使用TP/SP=16时可处理16k token的序列长度，优于单纯使用张量并行的情况。
随着张量并行度增加，计算效率和显存容量需要权衡。从TP=8提升到TP=16时性能下降明显，因为涉及节点内到节点间的通信转变。
更高的并行度虽能减少激活内存以处理更大批次，但会降低单GPU吞吐量，超过单节点GPU数量对应的阈值时更为明显。
序列长度增加时，TP区域的激活值内存仍会激增；模型过大时，TP=8也无法适配，会因节点间连接导致速度大幅下降。可分别用上下文并行和流水线并行解决这些问题。
上下文并行
借鉴序列并行按序列长度拆分的思路，对已应用张量并行的模块沿序列长度和另一个维度进行拆分，在整个模型上应用序列拆分，而非仅在模型的序列并行区域。
这种方式对多数模块无影响，因为这些模块中每个token独立处理，且无需像张量并行那样进行高成本的通信，仅分割输入，不拆分权重矩阵。计算梯度后，通过all-reduce操作同步上下文并行组内的梯度。
注意力模块中每个token需访问其他所有token的键/值对。由于上下文并行按序列维度拆分输入，注意力模块需在GPU间进行全面通信以交换键/值数据。
为高效处理这种通信，引入了环形注意力（Ring Attention）技术。
以4个GPU和4个token的输入为例，每个GPU先异步将自身的键/值对发送给其他GPU，在等待时计算已有数据的注意力分数。理想状态下，计算完成前能收到下一个键/值对，可立即开始下一轮计算。
每个GPU在每个时间步执行三个操作：非阻塞式发送当前键和值（最后一步除外）、本地计算注意力分数、等待接收前一个GPU的键/值后循环执行。
环形注意力的简单实现会因因果注意力矩阵的形状导致GPU计算负载不均衡。
SoftMax按行计算，GPU1因起始就有完整行的token数据可立即计算，且无需接收其他GPU信息；GPU2则需等待第二轮接收更多数据才能计算，GPU1计算量明显少于其他GPU。
Zig-Zag环形注意力机制
当前需更好的方式分配输入序列，Zig-Zag注意力摒弃顺序分配token至GPU，而是采用混合排序，使每个GPU上都有早期和晚期token，实现计算在各GPU上的平衡分布。
由于每个GPU完成计算需要其他GPU的信息，存在两种常见方式来重叠计算和通信：AllGather和All-to-All（环形）实现。
张量并行可在单个节点上拆分模型处理大型模型，上下文并行则用于解决长序列导致的激活量激增问题。但张量并行在跨节点扩展时效果不佳。
那么，如果模型权重无法轻松地在一个节点上存储，该怎么办呢？
这时，流水线并行（Pipeline Parallelism）就派上用场了！
流水线并行
张量并行扩展到超过单个节点的GPU数量（一般4或8个）时，会受到节点间连接低带宽网络影响，性能下降明显。
对于70B参数以上的模型，单节点4-8个GPU难以承载其权重规模，因此需要流水线并行技术。
将模型的各层分布到多个GPU上，如8个GPU时，可把第1-4层放于GPU1，第5-8层放于GPU2等。这样每个GPU仅需存储和处理部分模型层，减少了单个GPU的内存需求。
但由于每个GPU仍需处理完整批次数据，激活内存不会因层的划分而减少，且激活张量需在GPU间按流水线顺序传递。流水线并行中的数据处理具有顺序性，GPU利用率不高。
全前向全反向（AFAB）调度
由于计算是顺序进行的，存在效率不高的问题。GPU存在空闲时间，会降低利用率。通过公式推导可知，空闲时间与流水线并行度相关，并行度越高，空闲时间越长，利用率越低。
将批次数据拆分成更小的微批次进行并行处理。AFAB调度先进行所有前向传播，再进行所有反向传播，保留了模型训练代码的总体结构，易于实现。计算表明，增加微批次数量可减小空闲时间占比，提高效率。
由于反向传播前需保存所有激活值，这种方法会导致内存迅速耗尽。
因此需要探索新方法，如在进行前向计算时就开始反向传播，以尽早丢弃反向传播所需的激活值，避免显存爆炸。
One-forward-one-backward调度
交替执行一次前向和一次反向传播，尽早开始反向传播。该方式虽未显著提升训练效率，但能减少激活内存占用，且可增加微批次以减小空闲时间。
微批次数量等于或小于流水线并行度-1时，性能低且随并行度增加而下降。使用更多微批次有助于提升低并行度性能，但高并行度时仍受限。
受全局批大小限制，不能随意增加微批次，并行度增加时空闲时间会相应增加。
微批次数量较少时，跨节点扩展性能下降仅14%，优于张量并行，在跨节点分布式训练中有吸引力。
交错阶段技术
不同于简单按模型深度划分，交错阶段如将奇数层和偶数层分别置于不同GPU，形成「循环流水线」。微批次前向传播时在GPU间循环。
虽增加了通信量，但每次前向和反向传播时间因v（每个GPU的阶段数或模型块数）而减少，可通过增加微批次和交错阶段减小空闲时间。
交错阶段使调度更复杂，需在特定时刻决定优先策略，即「深度优先」（尽快通过模型）或「广度优先」（尽可能填满流水线）。
Llama 3.1的流水线并行方法采用了带交错阶段的单前向单反向设置，深度优先和广度优先的优先级设置可调节。
零气泡和双管道技术
为减少空闲时间提出了更复杂方法，关键是细粒度拆分操作并交错执行。如DeepSeek V3/R1的DualPipe。
ZeroBubble发现矩阵乘法反向传递中，输入反向操作（B）和权重反向操作（W）可分离，W可在对应B之后灵活安排，用于填补流水线空闲时间。
DeepSeek的DualPipe在V3技术报告中扩展了分解方法，针对从PP维度两端传播的两个数据流交错，以进一步减少 GPU空闲时间。
专家并行
MoE模型近年来因GPT-4、Mixtral、DeepSeek-V3/R1等模型受到关注。其基本思想是每一层不采用单个前馈模块，而是设置多个并行模块，对token进行不同处理。
MoE层的设计使专家并行易于实现，因前馈层完全独立。与张量并行（TP）相比，专家并行更轻量，无需拆分矩阵乘法，只需将token隐藏状态路由到合适专家。
实际中，专家并行（EP）常与其他并行方式结合使用。因EP仅影响MoE层，不分片输入token，若仅用EP，GPU处理非MoE模块时会有冗余计算。EP高效运行的技巧与模型设计紧密相关。
更多详细内容请查看原文！
参考资料：
https://huggingface.co/spaces/nanotron/ultrascale-playbook
