deepspeed入门教程
DeepSpeed是微软推出的大规模模型分布式训练的工具,主要实现了ZeRO并行训练算法。
本文是huggingface的DeepSpeed文档的笔记,做查询和备忘,初次学习建议结合原始文档食用。原始文档链接:
https://huggingface.co/docs/transformers/main/main_classes/deepspeed
分布式训练
目前主流的大模型分布式训练主要包括两种:
数据并行训练
模型并行训练
Deep Speed
1 基础组件
分布式训练需要掌握分布式环境中的基础配置,包括节点变化、全局进程编号、局部进程编号、全局总进程数、主节点等。这些组件都跟分布式训练紧密相关,同时组件之间也有非常大的联系,例如通信联系等。
2 通讯策略
通信策略 | 通信作用 |
---|---|
mpi | 它是一种跨界点的通信库,经常用于CPU集群的分布式训练 |
gloo | 它是一种高性能的分布式训练框架,可以支持CPU或者GPU的分布式训练 |
nccl | 它是nvidia提供的GPU专用通信库,广泛用于GPU上的分布式训练 |
我们在使用DeepSpeed进行分布式训练的时候,可以根据自身的情况选择合适的通信库,通常情况下,如果是GPU进行分布式训练,可以选择nccl。
3 Zero(零冗余优化器)
Microsoft开发的Zero可以解决分布式训练过程中数据并行和模型并行的限制。
Zero通过在数据并行过程中划分模型状态(优化器、梯度、参数),来解决数据并行成可能出现内存冗余的情况(正常数据并行训练,模型全部参数是复制在各个机器上的);同时可以在训练期间使用动态通信计划,在分布式设备之间共享重要的状态变量,这样保持计算粒度和数据并行的通信量。
措施 | 目的 | |
---|---|---|
数据并行 | 模型状态划分(优化器,梯度,参数) | 减少内存冗余(原本是参数复制在各个机器上) |
动态通讯计划 | 共享重要的状态变量 | |
Zero是用于大规模模型训练优化的技术,它的主要目的是减少模型的内存占用,让模型可以在显卡上训练,内存占用主要分为Model States和Activation两个部分,Zero主要解决的是Model States的内存占用问题。
Zero将模型参数分成三个部分:
状态 | 作用 |
---|---|
Optimizer States | 优化器在进行梯度更新的时候需要用到的数据 |
Gradient | 在反向转播过程中产生的数据,其决定参数的更新方向 |
Model Parameter | 模型参数,在模型训练过程中通过数据“学习”的信息 |
Zero的级别如下:
级别 | 作用 |
---|---|
Zero-0 | 不使用所有类型的分片,仅使用DeepSpeed作为DDP |
Zero-1 | 分割Optimizer States, 减少4倍内存,通信容量和数据并行性相同 |
Zero-2 | 分割Optimizer States和Gradients,减少8倍内存,通信容量和数据并行性相同 |
Zero-3 | 分割Optimizer States、gradients、Parametes,内存减少与数据并行度呈线性关系。例如,在64个GPU(Nd=64)之间进行拆分将产生64倍的内存缩减。通信量有50%的适度增长 |
Zero-Infinity | Zero-Infinity是Zero-3的扩展,它允许通过使用 NVMe 固态硬盘扩展 GPU 和 CPU 内存来训练大型模型 |
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