最近语言大模型（LLM）异常火爆，一个非常特别的开源社区正在探索在消费级硬件上微调、提供服务和进行推理的最佳方式。为满足上述需求，出现了许多出色的开源代码库，以HuggingFace生态系统为中心，这些代码库还包括FastChat、Axolotl和LLama.cpp。本文专注于分布式训练策略的具体细节，特别是DeepSpeed和FSDP，并总结了以多GPU和多节点训练为主的不同高效微调方法。显然，当前的趋势是，我们会使用越来越多的计算资源，因此将需要更多GPU来运行更大的模型。在这种情况下，理解这些主题尤为重要，尤其是当你想要将几个3090家庭服务器升级到具有8个A100 80GB的GCP容器时，另外，这对于试图微调自己的语言模型的初创公司或企业来说也很重要。大型科技公司进行的实际大规模训练涉及大量资料，这些内容大部分来自主导了BLOOM-176B训练的Stas Bekman，GPU匮乏的用户（即GPU-poors）关注这些资料的意义并不大。本文从多个优秀的资源中整合了各种观点，着重讨论了HuggingFace生态系统的相关内容，并考虑了一些来自在线资源以及作者本人在2023年暑期实习中所学到的实际情况。1. 在分布式训练和性能优化方面，我们应该关注什么？DeepSpeed和FSDP在背后是如何运作的？2. 不同的分布式训练策略需要的硬件设置和注意事项？4. 一些可以覆盖所有重要微调优化的实用指南，用以在多GPU和多节点设置中训练大模型。（本文作者为加州大学圣地亚哥分校计算机科学系的硕士研究生Sumanth R Hegde。以下内容由OneFlow编译发布，转载请联系授权。原文：https://sumanthrh.com/post/distributed-and-efficient-finetuning/ ）分布式训练涵盖的范围非常广泛，因此本文无法覆盖所有内容。在训练/微调LLM时，我们通常会面对超10亿参数的庞大模型和大规模数据集（超1万亿词元的预训练数据集，超100万词元的监督微调数据集）。我们最终的目标是尽快完成训练，以最大化吞吐量，即希望能够每秒处理尽可能多的样本。LLM在训练过程中需要大量的GPU显存，不仅仅因为模型参数数量庞大（例如，Falcon 40B拥有40亿参数，在BF16格式下仅模型权重就需要约74GB显存），还因为优化器状态所需的显存——例如，使用普通的AdamW优化器，每个参数需要12个字节来存储模型权重的副本、动量和方差参数。因此，我们需要智能的分布式训练策略，以确保每个GPU worker只需处理部分训练状态和数据。1. 数据并行（Data Parallelism, DP）：每个GPU worker获取整个小批量数据的一部分，并在该部分数据上计算梯度。然后在所有worker上对梯度进行平均，以更新模型权重。在其最基本的形式中，比如PyTorch中的DDP，每个GPU存储模型权重的副本、优化器状态以及所处理数据部分的梯度。2. 模型并行/垂直模型并行（MP）：在模型并行中，模型被分割成多个部分，每个部分被放置在不同的GPU上，这被称为垂直模型并行。举例来说，如果有一个包含12层的单一模型，该模型的不同层会被分别放置在3个不同的GPU上。--------------- --------------- -----------------1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 |--------------- --------------- -----------------
在朴素模型并行（naive model parallelism）中，所有GPU在处理相同的数据批次时都需要等待前一个GPU完成计算，然后才能处理数据。这意味着在任何给定时刻，除了一个GPU外，其他GPU实际上都处于闲置状态（因此称为“朴素”）。为改善这种情况，可以使用流水线并行（PP），这种方式通过让不同微批次数据的计算重叠，给你带来并行的错觉。这类似于计算机架构中的经典流水线。参见以下关于GPipe的博文：为了在多个加速器上实现高效训练，GPipe将模型划分到不同加速器，并自动将一个小批次的训练样本分成更小的微批次。通过在这些微批次上进行流水线式的执行，各个加速器可以进行并行计算。3. 张量并行（TP)：在张量并行中，每个GPU通过在GPU worker之间对模型进行水平切片，仅处理张量的一部分。每个worker处理相同的数据批次，计算他们所拥有权重部分的激活值，并交换彼此需要的部分，每个worker计算他们所拥有权重部分的梯度。我们可以将上述各种并行策略结合起来，以实现更好的吞吐量增益。接下来，我们将更详细地了解两种用于数据并行训练的改进方法：零冗余优化器（Zero Redundancy Optimizer）和密切相关的全切片数据并行策略（Fully Sharded Data-Parallel strategies）。注释：我将使用术语“GPU worker”来指代在每个GPU上运行的各个进程。虽然这种表述并不十分准确，但在数据并行设置中，这样说更方便，更易于理解。这是当前最高效、最热门的分布式训练策略之一。DeepSpeed的ZeRO是一种数据并行处理形式，它极大提高了内存效率，其主要思想是利用数据并行训练中的内存冗余和快速GPU间通信的最新改进来提高吞吐量，根据不同阶段，会增加一些通信量。实际上，ZeRO有两个组成部分：ZeRO-DP（数据并行）和ZeRO-R（残留内存）。DeepSpeed团队还提出了一些后续优化措施，这进一步提升了ZeRO的吸引力，例如ZeRO-Offload/Infinity（将计算卸载到CPU/NVMe磁盘）和ZeRO++（实现了灵活的多节点训练和量化权重）。ZeRO-DP的可视化图表如下（来自DeepSpeed的博客文章）：在64个GPU上训练7.5B参数模型时，模型表现如下：2. ZeRO 第一阶段/：内存减少4倍（特定示例），与基准的通信量相同（没有额外的GPU间通信）。3. ZeRO 第二阶段/：内存减少8倍（特定示例），与基准的通信量相同。4. ZeRO 第三阶段/：内存减少64倍（特定示例），通信量为基准的1.5倍（这里的1.5倍针对的是不同硬件设置和模型大小）。PyTorch DDP实现了简单的数据并行性。每个GPU worker都有模型权重、优化器状态和梯度的副本。后向传播后，各个worker间的梯度会被平均化（即全局归约（all-reduce）），并且模型权重会被更新。关于通信量的注释：为理解ZeRO的好处，我认为，重要的是要明确通讯量的确切含义。典型DP中存在一个全局归约步骤，其中每个worker都会发送出其所拥有的梯度数组，然后接收其他worker的梯度数组，以获得平均值。下文摘自ZeRO论文：当前最先进的全局归约实现采用了两步操作：第一步reduce-scatter，它在不同进程上的数据的不同部分会被归约。第二步是all-gather，其中每个进程都在这一步中收集所有进程上的归约数据。这两步的结果就是全局归约。同时，reduce-scatter和all-gather都采用了流水线方式实现，这导致每个步骤需要Ψ个元素（对于具有Ψ个元素的数据）的总数据搬运量。因此，在标准的数据并行性中，每个训练步骤会产生2Ψ的数据搬运量。在这种情况下，“数据”指的是我们的梯度，而进程指的是每个GPU上运行的单个worker。总而言之，我想表达的是：如果你有个参数，那么普通数据并行性将产生通信成本。在这种情况下，只有优化器状态在GPU worker之间进行了分区/分片，而模型权重和梯度在所有worker之间进行了复制。反向传播后，需要进行一次常规全局归约，以便在所有worker间获取平均梯度值。然后，每个worker更新其分区中的优化器状态。Adam方程如下。(w、g、v 和m 分别对应权重、梯度、速度和动量)。需要注意的是，这些都是逐元素操作，在计算梯度后，各个权重分片之间并没有依赖关系。通信量：首先我们会执行一次全局归约操作，将更新后的梯度通信给所有GPU，然后，在更新各自分片的优化器状态之后，每个GPU仍需要从其他GPU获取更新后的权重。ZeRO论文并没有清楚表明这样做是否会增加通信量。在我看来，对ZeRO的第1阶段和第2阶段来说，这种实现实际上是一样的：ZeRO Stage 2 /（优化器状态+梯度分区）在这种情况下，优化器状态和梯度都被分区/分片到不同的worker上。这意味着，两个GPU worker在训练期间不仅要关注不同的微批次数据，还要维护模型参数子集的梯度。关键在于，每个worker都在更新其优化器状态的分区，因此对于一个worker而言，它所需的梯度（或者说，经过归约/平均的梯度）只是对应于该状态分区的梯度。至于实现方式，正如上面提到的，DeepSpeed有效地执行了reduce-scatter操作，其中每个worker对应的梯度在该worker处被平均（而不是对所有参数进行典型的all-reduce操作）。这意味着，在相同通信量下节省了更多内存，也就是说，与DDP相比，这里没有额外的数据搬运成本。
注意：在使用ZeRO Stage 1和2时，仍然需要整个模型适配单个GPU。此外，在使用RAM时，需要注意如下事项：随着进程/GPU数量的增加以及模型大小的扩展（超过40亿参数），模型初始化会占用大量RAM。ZeRO 3对此有所改进。
ZeRO Stage 3 /（优化器状态+梯度+参数分区）对我来说，这是ZeRO最有趣的阶段。除优化器状态和梯度外，第3阶段还跨worker划分了模型参数。来自DeepSpeed的可视化图表如下：使用DeepSpeed ZeRO-3对训练状态分片借用Stas Bekman指南中的一个例子，假设有以下3层模型和4个GPU：La | Lb | Lc---|----|---a0 | b0 | c0a1 | b1 | c1a2 | b2 | c2a3 | b3 | c3
使用DeepSpeed ZeRO 3，GPU的配置方式如下：GPU 0:La | Lb | Lc---|----|---a0 | b0 | c0
GPU 1:La | Lb | Lc---|----|---a1 | b1 | c1
GPU 2:La | Lb | Lc---|----|---a2 | b2 | c2
GPU 3:La | Lb | Lc---|----|---a3 | b3 | c3
在ZeRO-3中，模型的每一层都被水平切片，每个worker存储权重张量的一部分。在前向和后向传播过程中（每个GPU worker 仍然看到不同的微批次数据），不同的GPU worker交换它们所拥有的每一层的部分（按需进行参数通信），并计算激活/梯度。其余部分类似ZeRO Stage 2。很容易看出，ZeRO-3的通信量是基准DDP的1.5倍：在每个训练步骤中，我们需要在前向传播中额外进行一次模型参数的all-gather操作。在这个操作中移动的数据量为（每个GPU），因此总通信量为（参数all-gather）+ （梯度reduce-scatter）+ （all-gather，用于更新的参数）= 3 = 1.5倍DDP。考虑到内存消耗被GPU workerN 削减，这给人留下了深刻印象。ZeRO论文的另一个关键洞察如下：只要有足够的设备来共享模型状态，ZeRO就可以使DP适应任意大小的模型。也就是说，只要有足够多的GPU，在进行数据并行（DP）训练时，我们就不会再受到每个GPU显存的限制（说起来容易做起来难）。我并不想深入探讨这一话题，但ZeRO-R在ZeRO-DP的基础上，通过关注激活的内存消耗和管理内存碎片化做了改进提升。ZeRO-R通过对激活进行分区，减少了激活的内存占用。此外，它还在管理临时缓冲区方面进行了一些改进，你可以将之视为在worker间进行梯度累积和归约期间分配用于存储中间结果的内存。ZeRO-Offload是一种优化技术，可以将优化器和计算从GPU卸载到主机CPU上。在2021年1月发布时，ZeRO-Offload可在1个NVIDIA V100 GPU上实现40 TFLOPS（V100 32 GB vRAM，最大吞吐量为130 TFLOPS），适用于10亿参数模型。而采用PyTorch DDP时，最大值为30 TFLOPS，适用于14亿参数模型，在不耗尽内存的情况下，这是可以运行的最大模型。将计算卸载到CPU的主要问题是，以吞吐量衡量，CPU要比GPU慢上多个数量级。ZeRO-Offload采用的这种策略是，只将较少的密集计算（<，其中是模型大小，是批次大小）卸载到CPU，以使总计算复杂度保持不变（）。在实践中，这意味着诸如范数计算（norm calculation）、权重更新等操作可以在CPU上完成，而前向和后向传播的矩阵乘法需要在GPU上完成。ZeRO-Offload适用于ZeRO的所有阶段（1、2和3）。这里（https://docs.it4i.cz/dgx2/introduction/）给出了用于实验的DGX-2节点规格，该节点配备了16个V100 GPU。需要注意的是，如果处于ZeRO-2设置中，ZeRO-Offload仍然会受到每个GPU可用内存的限制，即在每个GPU上容纳整个模型可能成为瓶颈。ZeRO-Infinity是ZeRO-Offload的改进版本，于2021年4月推出。它允许将计算卸载到磁盘（NVMe内存），并对CPU卸载进行了改进。研究表明，ZeRO-Infinity在一个DGX-2节点上训练时，可以适应10-100万亿参数的模型（万亿级别！）。ZeRO-Infinity通过同时利用CPU和NVMe内存来实现这一点。以下是论文中的可视化图表：以上是ZeRO Infinity在4个数据并行排布（GPU）上的截图，描述了反向传播期间的状态。分区/分片的参数从慢速内存（CPU+ NVMe）移动到GPU，然后被收集起来形成完整的层。在梯度计算之后，它们被聚合、重新分区，然后卸载到慢速内存。分区/分片的参数从慢速内存（CPU+ NVMe）移动到GPU，然后被收集起来形成完整的层。在梯度计算之后，它们被聚合、重新分区，然后卸载到慢速内存。与ZeRO-Offload不同，ZeRO-Infinity是基于ZeRO-3专门构建的。作者在32个DGX-2节点、512个GPU上对模型速度进行了评估。结果表明，ZeRO-Infinity可训练高达20万亿参数的模型，每个GPU的吞吐量可达49 TFlops，而3D并行等替代并行策略可训练的模型规模要小40倍。ZeRO-Infinity要想成为有竞争力的选择，需要满足一些带宽要求，即NVMe-CPU和CPU-GPU通信。关于ZeRO-Offload和ZeRO-Infinity之间的差异，以下是DeepSpeed团队的评论：DeepSpeed首先包含了ZeRO-Offload的卸载功能，ZeRO-Offload是一个用于将优化器和梯度状态卸载到ZeRO-2内的CPU内存的系统。ZeRO-Infinity是ZeRO-3可以使用的下一代卸载功能。相比ZeRO-Offload，ZeRO-Infinity能够处理更多数据，能更有效地利用带宽，并且能更好地实现计算和通信重叠。默认情况下，使用ZeRO-3进行卸载时，ZeRO-Infinity的优化机制会自动生效；使用Stage 1/2进行卸载时，ZeRO-Offload会生效。1. ZeRO-Offload/Infinity 教程: https://www.deepspeed.ai/tutorials/zero-offload/2. ZeRO-Offload - 十亿级模型训练大众化 : https://arxiv.org/abs/2101.068403. ZeRO-Infinity - 打破GPU内存墙，实现超大规模深度学习: https://arxiv.org/abs/2104.07857ZeRO++是DeepSpeed团队对ZeRO-3的最新改进。主要改进如下：1. 量化权重（qwZ）：通过将模型权重量化为int8，将all-gather参数通信量减少一半。2. 层次划分 (hpZ)：层次划分是一种混合分区方案，适用于DeepSpeed ZeRO 3的多节点设置。在这种情况下，模型参数会在一个节点内分片，然后在不同节点之间复制。这意味着，与典型的ZeRO-3全面运行相比，我们无法节省相同数量的内存，但可以避免昂贵的节点间参数通信开销，从而提高整体吞吐量。我更倾向于FSDP中使用的 "混合分片（hybrid sharding）" ，而非 "层次划分"，下文讨论FSDP时，会再次深入讨论这个问题。3. 量化梯度（qgZ）：通过在梯度reduce-scatter操作中使用int4量化数据替换fp16，可以节省更大通信量。（回顾：这是发生在ZeRO 2/3阶段，针对分片梯度进行的梯度聚集和平均）总体而言，与ZeRO-3 相比，ZeRO++通过这三项改进将通信量减少了4倍。