Hugging Face研究员Stas Bekman的客座博客文章
随着最近的机器学习模型增长速度远远超过新发布的显卡GPU内存的增加速度，许多用户无法在他们的硬件上训练甚至只是加载其中一些庞大的模型。虽然目前正在努力将其中一些庞大的模型精简为更易管理的大小，但这一努力并没有及时产生足够小的模型。
2019年秋季，Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase和Yuxiong He发表了一篇论文：《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》，该论文中包含了许多关于如何使硬件做出比之前认为可能的更多事情的巧妙新想法。不久之后，DeepSpeed被发布，它为世界提供了该论文中大部分想法的开源实现（其中一些想法仍在进行中），与此同时，Facebook的一个团队发布了FairScale，该团队也实现了ZeRO论文中的一些核心思想。
如果您使用Hugging Face Trainer，在transformers v4.2.0版本中，您可以实验性地支持DeepSpeed和FairScale的ZeRO功能。新的–sharded_ddp和–deepspeed命令行Trainer参数分别提供了FairScale和DeepSpeed的集成。这是完整的文档。
本博客文章将介绍您如何从ZeRO中受益，无论您是否只拥有单个GPU还是整个GPU堆栈。
让我们使用t5-large模型和transformers GitHub仓库中examples/seq2seq下的finetune_trainer.py脚本进行一个小的微调翻译任务实验。
我们有2x 24GB（Titan RTX）的GPU进行测试。
这只是一个概念验证基准测试，所以肯定还有进一步的改进空间，因此我们将在2000个训练样本和500个评估样本的小样本上进行基准测试以进行比较。默认情况下，评估使用大小为4的波束搜索，所以比使用相同数量的样本进行训练要慢，这就是为什么在这些测试中使用了4倍较少的评估样本的原因。
这是我们基线的关键命令行参数：
export BS=16
python -m torch.distributed.launch --nproc_per_node=2 ./finetune_trainer.py \
--model_name_or_path t5-large --n_train 2000 --n_val 500 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro [...]我们只是使用DistributedDataParallel(DDP)来提高基线的性能。在遇到内存溢出错误之前，我能够适应批量大小（BS）为16。
请注意，为了简单起见并使其更易于理解，我只显示了对于此演示重要的命令行参数。您可以在这篇文章中找到完整的命令行。
接下来，我们将每次增加以下之一重新运行基准测试：
--fp16--sharded_ddp（fairscale）--sharded_ddp --fp16（fairscale）--deepspeed无CPU卸载--deepspeed有CPU卸载
由于每种技术的关键优化在于更高效地利用GPU内存，我们将尝试不断增加批量大小，并期望训练和评估完成得更快（同时保持指标稳定或甚至改进一些，但我们这里不专注于这些）。
请记住，训练和评估阶段非常不同，因为在训练期间，模型权重被修改，梯度被计算，并存储了优化器状态。在评估期间，这些都不会发生，但在这个特定的翻译任务中，模型将尝试搜索最佳假设，因此它实际上必须在满意之前多次运行。这就是为什么它不快，特别是当模型很大时。
让我们来看看这六次测试的结果：
很容易看出，无论是FairScale还是DeepSpeed在总训练和评估时间，以及批次大小方面都提供了很大的改进。目前，DeepSpeed实现了更多的魔法，并且似乎是短期的赢家，但是Fairscale更容易部署。对于DeepSpeed，您需要编写一个简单的配置文件并更改命令行的启动器，而对于Fairscale，您只需要添加--sharded_ddp命令行参数，因此您可能希望先尝试它，因为它是最低成本的。
根据80:20法则，我只花了几个小时进行这些基准测试，并且没有尝试通过调整命令行参数和配置来挤压每个MB和秒，因为从简单的表格中很明显可以看出您想要尝试的下一步。当您面临一个需要运行数小时甚至数天的真实项目时，一定要花更多时间来确保您使用最优的超参数，以更快速、以最小成本完成工作。
如果您想要自己尝试这个基准测试，或者想要了解更多有关用于运行它的硬件和软件的详细信息，请参考此文章。
尽管Fairscale只在多个GPU上为我们提供了提升，但DeepSpeed即使对于只有一个GPU的人来说也是一份大礼。
让我们试试不可能的事情 – 在一张24GB的RTX-3090卡上训练t5-3b。
首先，让我们尝试使用普通的单GPU设置来微调庞大的t5-3b：
export BS=1
CUDA_VISIBLE_DEVICES=0 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 [...]没有奖励，即使使用BS=1我们也得到：
RuntimeError: CUDA内存不足。尝试分配64.00 MiB（GPU 0; 23.70 GiB总容量；
已分配21.37 GiB; 45.69 MiB空闲; PyTorch总共预留了22.05 GiB）。请注意，与之前一样，我只显示了重要部分，完整的命令行参数可以在此处找到。
现在将您的transformers更新到v4.2.0或更高版本，然后安装DeepSpeed：
pip install deepspeed然后让我们再试一次，这次在命令行中添加DeepSpeed：
export BS=20
CUDA_VISIBLE_DEVICES=0 deepspeed --num_gpus=1 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 --deepspeed ds_config_1gpu.json [...]et voila！我们得到了一个批次大小为20的训练结果。我可能还可以进一步提高它。程序在BS=30处由于内存不足而失败。
以下是相关的结果：
2021-01-12 19:06:31 | INFO | __main__ | train_n_objs = 60
2021-01-12 19:06:31 | INFO | __main__ | train_runtime = 8.8511
2021-01-12 19:06:35 | INFO | __main__ | val_n_objs = 10
2021-01-12 19:06:35 | INFO | __main__ | val_runtime = 3.5329我们无法将这些与基准进行比较，因为基准甚至无法启动并立即由于内存不足而失败。
非常惊人！
我只使用了一个小样本，因为我主要是为了能够训练和评估这个通常无法适应24GB GPU的巨大模型。
如果您想要自己尝试这个基准测试，或者想要了解更多有关用于运行它的硬件和软件的详细信息，请参考此文章。
由于transformers只是集成了这些神奇的解决方案，并不是它们的发明者，因此我将分享一些资源，您可以在其中了解所有细节。但是以下是一些快速见解，可能有助于理解ZeRO是如何实现这些惊人的功能的。
ZeRO的关键特性是将分布式数据存储添加到了非常熟悉的数据并行训练的概念中。
每个GPU上的计算与数据并行训练完全相同，但参数、梯度和优化器状态以分布式/分区的方式存储在所有GPU上，并且只在需要时提取。
下面的图表来自这篇博文，说明了它的工作原理：
ZeRO的巧妙方法是将参数、梯度和优化器状态平均分配到所有GPU上，并且每个GPU只有一个分区（也称为切片）。这导致GPU之间的数据存储没有重叠。在运行时，每个GPU通过向参与的GPU请求所需的信息来动态构建每一层的数据。
这个想法可能很难理解，你可以在这里找到我的解释。
截至目前，FairScale和DeepSpeed只对优化器状态和梯度进行分区（切片）处理。DeepSpeed和FairScale中模型参数的分区（切片）处理将很快推出。
另一个强大的功能是ZeRO-Offload（论文）。这个功能将一些处理和内存需求转移到主机的CPU上，从而允许更多的内容适配到GPU上。你在使用24GB GPU运行t5-3b时看到了它的显著影响。
许多人在pytorch论坛上抱怨的另一个问题是GPU内存碎片化。通常会出现OOM错误，错误信息可能如下所示：
RuntimeError: CUDA out of memory. Tried to allocate 1.48 GiB (GPU 0; 23.65 GiB total capacity;
16.22 GiB already allocated; 111.12 MiB free; 22.52 GiB reserved in total by PyTorch)程序想要分配大约1.5GB的内存，而GPU仍有6-7GB的未使用内存，但它报告只有大约100MB的连续空闲内存，然后由于OOM错误而失败。这是因为不同大小的内存块被分配和释放多次，随着时间的推移会产生空洞，导致内存碎片化，即有大量未使用的内存，但没有所需大小的连续块。在上面的示例中，程序可能可以分配100MB的连续内存，但显然无法得到1.5GB的单个块。
DeepSpeed通过自己管理GPU内存来解决这个问题，并确保长期内存分配与短期内存分配不混合，从而减少了碎片化。虽然论文没有详细说明，但源代码是可用的，因此可以看到DeepSpeed是如何实现的。
正如ZeRO代表零冗余优化器一样，很容易看出它名副其实。
除了DeepSpeed中即将推出对模型参数分区的支持外，它已经发布了一些我们尚未探索的新功能。这些功能包括DeepSpeed稀疏注意力和1位Adam，它们被认为可以减少内存使用并显著减少GPU间通信开销，从而实现更快的训练和支持更大的模型。
我相信我们还将看到FairScale团队带来的新礼物。我认为他们也在研究ZeRO第三阶段。
更令人兴奋的是，ZeRO正在集成到pytorch中。
如果你对这篇博文中分享的结果感兴趣，请点击这里了解如何使用transformers Trainer与DeepSpeed和FairScale。
当然，你也可以根据每个项目的说明修改自己的训练器来集成DeepSpeed和FairScale，或者你可以“作弊”，看看我们在transformers Trainer中是如何做到的。如果选择后者，请在源代码中使用grep查找deepspeed和/或sharded_ddp。
好消息是ZeRO不需要修改模型，只需要对训练代码进行必要的修改。
如果在这些项目的集成部分遇到任何问题，请在transformers中提交Issue。
但如果你在DeepSpeed和FairScale的安装、配置和部署方面遇到问题，你需要向相关领域的专家提问，请使用DeepSpeed Issue或FairScale Issue。
虽然您不必真正理解这些项目的工作原理，只需通过transformers Trainer部署它们，但如果您想了解其中的原因和方法，请参考以下资源。
-
FairScale GitHub
-
DeepSpeed GitHub
-
论文：ZeRO：面向训练万亿参数模型的内存优化。这篇论文非常有趣，但非常简洁。
-
这里有一段关于该论文的很好的视频讨论，有可视化效果
-
论文：ZeRO-Offload：实现亿级规模模型训练的民主化。刚刚发布-这篇论文详细介绍了ZeRO Offload功能。
-
DeepSpeed配置和教程
-
除了论文外，我强烈推荐阅读以下详细的博客文章和图表：
- DeepSpeed：为所有人提供极限规模的模型训练
- ZeRO & DeepSpeed：新的系统优化使得训练超过1000亿参数的模型成为可能
- Turing-NLG：微软的一个170亿参数的语言模型
-
DeepSpeed在GitHub上的示例
在将这些项目集成到transformers中的过程中，我们对FairScale和DeepSpeed开发团队提供的支持表示非常惊讶。
特别感谢以下人员：
- Benjamin Lefaudeux @blefaudeux
- Mandeep Baines @msbaines
来自FairScale团队的：
- Jeff Rasley @jeffra
- Olatunji Ruwase @tjruwase
- Samyam Rajbhandari @samyam
来自DeepSpeed团队的慷慨和关怀支持，以及对我们遇到问题的及时解决。
还要感谢HuggingFace提供用于运行基准测试的硬件访问权限。
Sylvain Gugger @sgugger和Stas Bekman @stas00负责集成这些项目。