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

2024-01-03 17:31:49

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。
  1. ZeRO-Offload技术
    ZeRO-Offload进一步扩展了ZeRO的能力,通过将部分计算任务卸载到CPU上,减少了GPU上的计算需求。这允许在GPU资源受限的情况下训练大型模型,同时还能够有效利用CPU资源。
  2. 参数分片
    在DeepSpeed中,参数分片是另一种减少GPU内存需求的手段。通过将模型的参数分割成更小的片段,并在训练过程中只在必要时将它们加载到内存中,DeepSpeed可以进一步减少单个GPU所需的内存量,从而允许更大模型的训练。
  3. 混合精度训练
    混合精度训练是指在训练过程中同时使用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使用

文章来源:https://blog.csdn.net/weixin_43603658/article/details/134783077
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。