AI知识库

53AI知识库

学习大模型的前沿技术与行业应用场景


分布式训练:DeepSpeed 与 Zero 数据并行
发布日期:2024-05-07 07:58:47 浏览次数: 2940 来源:AINLP


前言

随着大模型时代的到来,数据规模,计算量,模型参数等各方面都有巨大的增长,这给训练大模型带来了一些挑战:
  • 显存的挑战:当前常见的GPU 显存如 A100 是 80G 显存,而如 175B 的 GPT-3 模型需要 175B * 4bytes 即 700GB 模型参数空间,连加载都困难更不要说训练了。
  • 计算的挑战:175B模型所需要的计算量也是非常庞大的,并且叠加预训练数据量的增大,所需的计算量与BERT时代完全不可同日而语。
本文主要从大模型时代的分布式训练出发,主要讨论以数据并行为核心的 DeepSpeed Zero 优化。

三大并行策略

并行技术可以整体划分为两类:数据并行(DP)流水线并行(PP)以及张量并行(TP)

1 数据并行(DP)

  • 原理:数据并行简单来说,就是每个GPU worker都有一个完整模型的副本,然后获取 per_gpu_batch_size 的数据,每个worker在自己那批数据上计算梯度。最后所有worker 对梯度汇总进行平均最后更新模型权重。
  • 举例:假设有 8 张卡,per_gpu_batch_size=4,那么 global batch size=8*4=32,此时每张卡都保存一个模型,然后每张卡通过 4 个数据来分别计算梯度,之后计算梯度均值,最后每张卡再进行参数更新。
数据并行是在小模型时代(BERT)最常用的并行方法,在模型不是非常大的情况下,优先使用数据并行。

2 流水线并行(PP)与张量并行(TP)

当模型过大时,一张卡无法容纳一个模型时,就会考虑使用流水线并行以及张量并行,二者的区别在于对模型参数切分方向不同。

2.1 流水线并行

模型并行此处也可以称为流水线并行,详细的流水线解释可以参见:https://zhuanlan.zhihu.com/p/613196255。
  • 思想:将模型的不同层放到不同的 GPU 上。比如:一个6层的trasformer ,那么就可以将前3层放在一个GPU上,后3层放在另一个GPU上。如下图所示:

  • 缺点:每个GPU在处理相同batch时都需要等待前一个GPU完成计算才能开始处理,这样容易产生GPU空闲时间,即流水线气泡。为了改善这种情况,通常使用流水线并行(PP)来通过让不同的batch数据的计算重叠。

2.2 张量并行

  • 思想:对层内进行分割,实际上是将Tensor 矩阵运算分配到不同的 GPU 上。如下图就是将每一层分为两部分:

  • 优点:张量并行存储效率较高,
  • 缺点:每次前向和后向计算时会引入额外通信开销。

 

DeepSpeed 与 Zero

什么是Deepspeed?

Deepspeed 现在已经是大模型标配了,其是微软发布的大规模分布式训练工具,专门用来训练大模型。
Deepspeed Zero 是 Deepspeed 的核心组件,本质上是一种显存优化的数据并行方案。

Deepspeed Zero 的核心思想

混合精度训练阶段,ZERO 将训练中占显存的参数分为两部分:
  • 模型状态:模型参数(fp16/bf16,2A),模型梯度(fp16/bf16,2A),Adam状态(fp32的模型参数备份,fp32的momentum和fp32的variance,4A+4A+4A)。假设模型参数量为 A,则共需要:2A + 2A + (4A + 4A + 4A) = 16A 字节存储,如一个7B的模型,模型状态就需要 16 * 7 = 112GB
  • 剩余状态:除了模型状态之外的显存占用,包括激活值(activation)、各种临时缓冲区(buffer)以及无法使用的显存碎片(fragmentation)。


zero的优化思想为:针对模型状态部分,Zero 使用的方法是分片,即每张卡只存
的模型状态量,N是GPU个数 。但是本质上 ZERO 仍然属于一种数据并行方案。
一个常用的结论如下:
  • 速度: Zero 1 > Zero 2 > Zero 2 + offload > zero 3 > zero 3 + offload
  • 显存:阶段 0 (DDP) < 阶段 1 < 阶段 2 < 阶段 2 + 卸载 < 阶段 3 < 阶段 3 + 卸载

 

Deepspeed 的优化细节

  1. Deepspeed Zero Stage 1

针对 Adam 状态进行分片,此时每张卡模型状态所需显存变成了。我们以单机8卡,7B模型为例,此时模型状态部分所需显存 降低到 4 * 7 + (12*7)/8 = 38.5GB
  1. Deepspeed Zero Stage 2

在Zero-1 的基础上,针对模型梯度再次分片,此时每张卡模型状态所需显存变成了  同样以单机8卡,7B模型为例,此时模型状态部分所需显存降低到 2*7+(14*7)/8=26.25GB
  1. Deepspeed Zero Stage 3

在Zero-2的基础上,对模型参数也进行分片,此时每张卡模型状态所需显存变成了
同样以单机8卡,7B模型为例,此时模型状态部分所需显存降低到 16*7/8=14GB
  1. Zero-offload

核心思想:显存不足,内存来补,用时间来换空间。Zero-offload 核心是在 Zero-2的基础上将Adam状态,梯度转移到 CPU内存。一般情况下不推荐。

一点补充

注意,Deepspeed 优化的是模型参数,梯度,Adam状态所占的显存,而实际上,正如我在 大模型:训练时GPU显存不足怎么办中提到的,激活值同样占用着巨大的显存,此外还有一些其他的显存碎片。
以 qwen 1.5 为例,per_gpu_batch_size=1,seq_length=2048,精度bf16 时所需显存为 28.5G(粗略估计),此时即使用 Deepspeed zero 3(非常慢) ,都需要 28.5+14=42.5GB 显存。
这意味着在大模型时代中,如果不引入流水线并行和张量并行的情况下,V100-32G 和 A100-40G 对 7B 的模型训练都非常吃力。

最后

在大多数情况下,我们往往会使用 Deepspeed Zero 2 ,主要是考虑到训练效率的问题。
Deepspeed 是大模型时代必会的工具,推荐大家熟练使用。
最后给大家推荐一个常用的微调仓库,我本来想自己写一套,后来发现开源模型发展太快了,所以我现在也是无脑掉开源框架:https://github.com/hiyouga/LLaMA-Factory

参考

https://zhuanlan.zhihu.com/p/513571706


53AI,企业落地应用大模型首选服务商

产品:大模型应用平台+智能体定制开发+落地咨询服务

承诺:先做场景POC验证,看到效果再签署服务协议。零风险落地应用大模型,已交付160+中大型企业

联系我们

售前咨询
186 6662 7370
预约演示
185 8882 0121

微信扫码

与创始人交个朋友

回到顶部

 
扫码咨询