深度学习：大规模模型分布式训练框架DeepSpeed

DeepSpeed简介

随着机器学习模型变得越来越复杂和庞大，训练这些模型所需的计算资源也在不断增加。特别是在自然语言处理（NLP）等领域，模型大都达到了数十亿甚至数千亿参数的规模，这就需要多卡或者多节点分布式训练。为了有效地训练这些巨型模型，微软Microsoft推出了一个名为DeepSpeed的开源大规模模型分布式训练框架，提供了一些优化策略和工具，旨在支持大规模模型训练和更高效的训练过程。

DeepSpeed的设计初衷就是简化并优化大规模模型的训练。DeepSpeed 实现了3D并行，即三种并行方法的灵活组合：ZeRO 支持的数据并行，流水线并行和张量切片模型并行。这些技术可以极大地减少单个GPU的内存消耗，允许在有限资源下训练更大的模型。此外，DeepSpeed通过ZeRO-Offload可以同时利用CPU和GPU内存来训练大模型，极大地降低显存消耗。

DeepSpeed提供了高度优化的数据加载和网络通信工具，这些工具可以减少通信量、显著提高多GPU和多节点环境下的训练效率。该框架还支持混合精度训练，进一步提升了计算速度和资源利用率。此外，DeepSpeed提出了一种稀疏注意力核，与经典的稠密 Transformer 相比，它支持的输入序列长一个数量级，并在保持相当的精度下获得最高 6 倍的执行速度提升。

此外，DeepSpeed致力于优化大规模训练，但为了注重用户体验。DeepSpeed提供了易于集成的API，使得将现有的PyTorch模型迁移到DeepSpeed框架变得轻而易举，而无需进行大量的代码重写。

DeepSpeed是一个活跃的开源项目，其在GitHub上持续更新和维护，其Github链接为https://github.com/microsoft/DeepSpeed。

DeepSpeed核心特点

DeepSpeed的核心特点如下：

DeepSpeed在模型训练上提供了ZeRO、3D-Parallelism、DeepSpeed-MoE、ZeRO-Infinity 等创新技术，使大规模深度学习训练变得有效、高效，大大提高了易用性，并在可能的规模方面重新定义了深度学习训练格局。
DeepSpeed在模型推理上汇集了Tensor、Pipeline、Expert和ZeRO-parallelism等并行技术的创新，并将它们与高性能定制推理内核、通信优化和异构内存技术相结合，以前所未有的规模实现推理，同时实现无与伦比的延迟、吞吐量和性能。
DeepSpeed在模型压缩上，提供易于使用且组合灵活的压缩技术来压缩他们的模型，同时提供更快的速度、更小的模型大小并显着降低的压缩成本。
DeepSpeed团队推出了一项名为 DeepSpeed4Science 的新计划，旨在通过人工智能系统技术创新构建独特的能力，帮助领域专家解开当今最大的科学谜团。

DeepSpeed如何工作？

下面看一看DeepSpeed的一些核心组件和它们是如何工作的。

1. ZeRO优化器
   ZeRO（Zero Redundancy Optimizer）是DeepSpeed的一个关键创新，旨在解决大模型训练中的内存瓶颈问题。它通过优化数据并行策略中的冗余数据来实现内存节省。简单来说，ZeRO将模型参数、梯度和优化器状态等分散存储在多个GPU上，而不是在每个GPU上存储一份完整的副本。它在训练期间使用动态通信计划，以在分布式设备之间共享必要的数据，以保持计算粒度和数据并行性的通信量。这样做可以显著减少每个GPU的内存负担，使得训练更大的模型成为可能。ZeRO的级别分类如下：

• ZeRO-0：禁用所有类型的分片，仅使用 DeepSpeed 作为 DDP (Distributed Data Parallel)；
• ZeRO-1：分割优化器状态，减少了4倍的内存，通信容量与数据并行性相同；
• ZeRO-2：分割优化器状态与梯度，8x内存减少，通信容量与数据并行性相同；
• ZeRO-3：分割优化器状态、梯度与参数，内存减少与数据并行度和复杂度成线性关系。
• ZeRO-Infinity 是ZeRO-3的拓展。允许通过使用 NVMe 固态硬盘扩展 GPU 和 CPU 内存来训练大型模型。ZeRO-Infinity 需要启用 ZeRO-3。

2. ZeRO-Offload技术
   ZeRO-Offload进一步扩展了ZeRO的能力，通过将部分计算任务卸载到CPU上，减少了GPU上的计算需求。这允许在GPU资源受限的情况下训练大型模型，同时还能够有效利用CPU资源。
3. 参数分片
   在DeepSpeed中，参数分片是另一种减少GPU内存需求的手段。通过将模型的参数分割成更小的片段，并在训练过程中只在必要时将它们加载到内存中，DeepSpeed可以进一步减少单个GPU所需的内存量，从而允许更大模型的训练。
4. 混合精度训练
   混合精度训练是指在训练过程中同时使用FP16（半精度浮点数）和FP32（单精度浮点数）两种精度的技术。使用FP16可以大大减少内存占用，从而可以训练更大规模的模型。在使用混合精度训练时，需要使用一些技术来解决可能出现的梯度消失和模型不稳定的问题，例如动态精度缩放和混合精度优化器等。

DeepSpeed如何使用？

DeepSpeed可以与现有的PyTorch代码库无缝集成，开发者仅需进行少量修改即可开始使用。下面是一些基本的步骤在已经搭建好模型的基础上来开始使用DeepSpeed：

1. 安装DeepSpeed：可以通过pip包管理器快速安装pip install deepspeed 。
2. 修改代码以兼容DeepSpeed：这通常包括导入DeepSpeed库import deepspeed、初始化DeepSpeed引擎deepspeed.initialize()，并对数据加载器data_loader和迭代训练（前向传播和反向传播）进行设置。
3. 配置运行环境：根据你的硬件和模型大小，设置配置文件，包括batch size、学习率、优化器选择、内存优化等等。
4. 启动训练：使用DeepSpeed提供的命令行工具来启动训练过程，这个工具可以在多个GPU上分布式地运行你的模型。

此外，DeepSpeed现已集成到一些开源深度学习框架中，例如Transformers、Accelerate、Lightning、MosaicML、Determined和MMEngine。可以结合这些开源框架使用DeepSpeed，例如，transformers框架可以通过Trainer来使用集成的DeepSpeed功能，这种用法需要提供配置文件deepspeed_config.json，详细教程见transformers官网链接：

from transformers import Trainer

deepspeed_config = "./deepspeed_config.json"

model = ...
args = TrainingArguments(
	...
	deepspeed=deepspeed_config,
	...
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    data_collator=data_collator,
    optimizer=optimizer,
)
trainer.train()
trainer.save_model("best")

参考文献

1. [LLM]大模型训练(一)–DeepSpeed介绍
2. [LLM]大模型训练(二)–DeepSpeed使用