DeepSpeed: 大模型训练框架 | 京东云技术团队
背景:
目前,大模型的发展已经非常火热,关于大模型的训练、微调也是各个公司重点关注方向。但是大模型训练的痛点是模型参数过大,动辄上百亿,如果单靠单个 GPU 来完成训练基本不可能。所以需要多卡或者分布式训练来完成这项工作。
一、分布式训练
1.1 目前主流的大模型分布式训练主要包括两种:
数据并行训练
模型并行训练
二、DeepSpeed
DeepSpeed 是由 Microsoft 提供的分布式训练工具,旨在支持更大规模的模型和提供更多的优化策略和工具。对于更大模型的训练来说,DeepSpeed 提供了更多策略,例如:Zero、Offload 等。
2.1 基础组件
分布式训练需要掌握分布式环境中的基础配置,包括节点变化、全局进程编号、局部进程编号、全局总进程数、主节点等。这些组件都跟分布式训练紧密相关,同时组件之间也有非常大的联系,例如通信联系等。
2.2 通信策略
既然是分布式训练,那机器之间必须要保持通信,这样才可以传输模型参数,梯度参数等信息。
DeepSpeed 提供了 mpi、gioo、nccl 等通信策略
我们在使用 DeepSpeed 进行分布式训练的时候,可以根据自身的情况选择合适的通信库,通常情况下,如果是 GPU 进行分布式训练,可以选择 nccl。
2.3 Zero(零冗余优化器)
Microsoft 开发的 Zero 可以解决分布式训练过程中数据并行和模型并行的限制。比如: Zero 通过在数据并行过程中划分模型状态(优化器、梯度、参数),来解决数据并行成可能出现内存冗余的情况(正常数据并行训练,模型全部参数是复制在各个机器上的);同时可以在训练期间使用动态通信计划,在分布式设备之间共享重要的状态变量,这样保持计算粒度和数据并行的通信量。
Zero 是用于大规模模型训练优化的技术,它的主要目的是减少模型的内存占用,让模型可以在显卡上训练,内存占用主要分为 Model States 和 Activation 两个部分,Zero 主要解决的是 Model States 的内存占用问题。
Zero 将模型参数分成三个部分:
Zero 的级别如下:
2.4 Zero-Offload:
相比 GPU,CPU 就相对比较廉价,所以 Zero-Offload 思想是将训练阶段的某些模型状态放(offload)到内存以及 CPU 计算。
Zero-Offload 不希望为了最小化显存占用而让系统计算效率下降,但如果使用 CPU 也需要考虑通信和计算的问题(通信:GPU 和 CPU 的通信;计算:CPU 占用过多计算就会导致效率降低)。
Zero-Offload 想做的是把计算节点和数据节点分布在 GPU 和 CPU 上,计算节点落到哪个设备上,哪个设备就执行计算,数据节点落到哪个设备上,哪个设备就负责存储。
Zero-Offload 切分思路:
下图中有四个计算类节点:FWD、BWD、Param update 和 float2half,前两个计算复杂度大致是 O(MB), B 是 batch size,后两个计算复杂度是 O(M)。为了不降低计算效率,将前两个节点放在 GPU,后两个节点不但计算量小还需要和 Adam 状态打交道,所以放在 CPU 上,Adam 状态自然也放在内存中,为了简化数据图,将前两个节点融合成一个节点 FWD-BWD Super Node,将后两个节点融合成一个节点 Update Super Node。如下图右边所示,沿着 gradient 16 和 parameter 16 两条边切分。
Zero-Offload 计算思路:
在 GPU 上面进行前向和后向计算,将梯度传给 CPU,进行参数更新,再将更新后的参数传给 GPU。为了提高效率,可以将计算和通信并行起来,GPU 在反向传播阶段,可以待梯度值填满 bucket 后,一遍计算新的梯度一遍将 bucket 传输给 CPU,当反向传播结束,CPU 基本上已经有最新的梯度值了,同样的,CPU 在参数更新时也同步将已经计算好的参数传给 GPU,如下图所示。
2.5 混合精度:
混合精度训练是指在训练过程中同时使用 FP16(半精度浮点数)和 FP32(单精度浮点数)两种精度的技术。使用 FP16 可以大大减少内存占用,从而可以训练更大规模的模型。但是,由于 FP16 的精度较低,训练过程中可能会出现梯度消失和模型坍塌等问题。
DeepSpeed 支持混合精度的训练,可以在 config.json 配置文件中设置来启动混合精度("fp16.enabled":true)。在训练的过程中,DeepSpeed 会自动将一部分操作转化为 FP16 格式,并根据需要动态调整精度缩放因子,来保证训练的稳定性和精度。
在使用混合精度训练时,需要注意一些问题,例如梯度裁剪(Gradient Clipping)和学习率调整(Learning Rate Schedule)等。梯度裁剪可以防止梯度爆炸,学习率调整可以帮助模型更好地收敛。
三、总结
DeepSpeed 方便了我们在机器有限的情况下来训练、微调大模型,同时它也有很多优秀的性能来使用,后期可以继续挖掘。
目前主流的达模型训练方式: GPU + PyTorch + Megatron-LM + DeepSpeed
优势
存储效率:DeepSpeed 提供了一种 Zero 的新型解决方案来减少训练显存的占用,它与传统的数据并行不同,它将模型状态和梯度进行分区来节省大量的显存;
可扩展性:DeepSpeed 支持高效的数据并行、模型并行、pipeline 并行以及它们的组合,这里也称 3D 并行;
易用性: 在训练阶段,只需要修改几行代码就可以使 pytorch 模型使用 DeepSpeed 和 Zero。
参考:
1. http://wed.xjx100.cn/news/204072.html?action=onClick
2. https://zhuanlan.zhihu.com/p/513571706
作者:京东物流 郑少强
来源:京东云开发者社区 转载请注明来源
版权声明: 本文为 InfoQ 作者【京东科技开发者】的原创文章。
原文链接:【http://xie.infoq.cn/article/17293f06b6abca3f3cf08dd05】。文章转载请联系作者。
评论