并行训练框架DeepSpeed介绍

目录

并行训练框架DeepSpeed介绍

一、DeepSpeed设计概览

1. DeepSpeed简介

2. DeepSpeed软件架构

3. DeepSpeed核心技术

二、DeepSpeed使用介绍

三、DeepSpeed最佳实践

---

一、DeepSpeed设计概览

1. DeepSpeed简介

(1) DeepSpeed简介

DeepSpeed是一个由微软开发的开源深度学习优化库，皆在提高大规模模型训练的效率和可扩展性。它通过多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等

DeepSpeed还提供了一些辅助工具，如分布式训练管理、内存优化和模型压缩等，以帮助开发者更好地管理和优化大规模深度学习训练任务

此外，DeepSpeed基于PyTorch构建，只需要简单修改即可迁移。DeepSpeed已经在许多大规模深度学习项目中得到了应用，包括语言模型、图像分类、目标检测等待

(2) 深度学习模型软件体系架构

DeepSpeed作为一个大模型训练加速库，位于模型训练框架和模型之间，用来提升训练、推理等

2. DeepSpeed软件架构

(1) DeepSpeed软件架构

DeepSpeed主要包含三部分：

①APIs：提供易用的api接口，训练模型、推理模型只需要简单调用几个接口即可。其中最重要的是initializa接口，用来初始化引擎，参数中配置训练参数及优化技术等。配置参数一般保存在config.json文件中

②Runtime：运动时组件，是DeepSpeed管理、执行和性能优化的核心组件。如部署训练任务到分布式设备、数据分区、模型分区、系统优化、微调、故障检测、checkpoints保存和加载等。该组件使用python语言实现

③Ops：用C++和cuda实现底层内核，优化计算和通信，提供了一系列底层操作等

架构优势：

①可以在训练框架上进行两部分(训练和推理分开)优化

②与紧密耦合的结构比，该结构可以更好的利用整个生态，且与深度集成相比，更容易维护

③与基础设置无关，用户可以选择喜欢的平台，如Azure ML、Azure VMs等

3. DeepSpeed核心技术

(1) 混合精度训练

混合精度训练是指在训练过程中同时使用FP16(半精度浮点数)和FP32(单精度浮点数)两种精度的技术。使用FP16可以大大减少内存占用，从而可以训练更大规模的模型。在使用混合模型训练时，需要使用一些技术来解决可能出现的梯度消失和模型不稳定的问题，例如动态精度缩放和混合精度优化器等

(2) BF16

BF16和FP16都是半精度浮点数格式，主要区别如下：

BF16：对接FP32，取值范围相同

①可以表示更广的数值范围。但尾数只有7位， 所以精度较低

②更适合深度学习领域，可以减少模型大小，加速计算，同时保持足够的精度

③Ampere架构的Tensor Core支持BF16

FP16：

①数值范围窄，但尾数有10位，精度较高

②更适合图形渲染等需要高精度的场景

③大多数GPU架构都支持

(3) 混合精度优化过程

DeepSpeed提供了混合精度训练的支持，在训练过程中，DeepSpeed会自动将一部分操作转换为FP16格式，并根据需要动态调整精度缩放因子，从而保证训练的稳定性和精度

(4) 分布式模型训练

大模型在训练时往往需要大量内存来存储中间激活、权重等参数，百亿模型甚至无法在单个GPU上进行训练，使得模型训练在某些情况下非常低效和不可能。这就需要进行多卡，或者多节点分布式训练

分布式并行的核心思想是把计算和存储分散到不同的设备。大规模深度学习模型训练主要范式：

①数据并行(DP)

②模型并行(MP)

③混合并行(HP)

(5) 数据并行

在数据并行系统中，每个计算设备都有整个神经网络型的完整副本，进行迭代时，每个计算设备只分配了一个批次数据样本的子集，并根据该批次样本子集的数据进行网络模型的前向和反向计算

(6) 模型并行

模型并行往往用于解决单节点内存不足(计算节点显存)的问题。模型并行可以从计算图角度，以下两种形式进行切分：按模型的层切分到不同设备，即层间并行或算子间并行，也称之为流水线并行(Pipeline Parallelism，PP)；将计算图层内的参数切分到不同设备，即层内并行或算子内并行，也称之为张量并行(Tensor Parallelism，TP)

(7) DeepSpeed vs Megatron

目前训练超大规模语言模型技术路线：PyTorch + Megatron-LM + DeepSpeed

176B BLOOM模型使用Megatron-DeepSpeed进行训练，下表列出了在训练BLOOM时各采用了两个框架的哪些组件：

DeepSpeed 和 Megatron 都是用来训练超大AI模型的工具，但分工不同：

Megatron(英伟达开发)

核心能力：专门做模型并行(把超大模型拆到多个GPU上跑)

特点：优化Transformer结构，适合训练像GPT-3这样的巨型模型

-

DeepSpeed(微软开发)

核心能力：主打训练加速和省内存(比如ZeRO优化、梯度检查点)

特点：能优化计算、通信，还能和Megatron配合用

Megatron负责拆模型，DeepSpeed负责省内存和通信优化，合起来能训练更大的模型(比如GPT-3、MT-NLG)

(8) ZeRO(零冗余优化器)

ZeRO(Zero Redundancy Optimizer)是一种用于优化大规模深度学习模型训练的技术。它的主要目标是降低训练期间的内存(显存)占用、通信开销和计算负载，从而使用户能够训练更大的模型并更高效地利用硬件资源

ZeRO是一系列显存优化方法的统称，它分为：

①ZeRO-DP

②ZeRO-R

③ZeRO-Offload

④ZeRO-Infinity

-> 显存占用(GPU内存)越来越少，内存(CPU内存)越来越多，计算时间越来越长

(9) 显存占用分析

ZeRO首先分析了模型训练中内存(显存)主要消耗在两个方面：

Model States：模型本身相关且必须存储的内容，包括：

①Parameters：模型参数

②Gradients：模型梯度

③Optimizer States：Adam优化算法中的momentum(动量)和variance(方差)

-> 训练1个参数，必须存储数据有4个：训练参数、梯度、优化算法动量和方差

Residual States：非模型本身必须，但在训练过程中产生的内容，包括：

①Activation：激活值

②Temporary Buffers：临时存储

③Unusable Fragmented Memory：碎片化存储空间

ZeRO分别使用ZeRO-DP和ZeRO-R来优化Model States和Residual States

(10) ZeRO-DP

ZeRO-DP包括三个阶段：

(11) ZeRO-1优化器状态切分

流程：

①每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮forward和backward后，各得一份梯度

②对梯度做一次AIIReduce，得到完整的梯度G -> 平均值

③得到完整梯度G，就可以对W做更新。我们知道W的更新由Optimizer States和梯度共同决定。由于每块GPU上只保管部分Optimizer States，因此只能将相应的W(蓝色部分)进行更新

④此时，每块GPU上都有部分W没有完成更新(图中白色部分)。所以我们需要对W做一次All-Gather，从别的GPU上把更新好的部分W取回来

(12) ZeRO-2优化器状态与梯度切分

流程：

①每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮forward和backward后，算得一份完整的梯度(下图中绿色+白色)

②对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度。例如对GPU1，它负责维护G1，因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可。汇总完毕后，白色块对GPU无用，可以从显存中移除

③每块GPU用自己对应的O和G取更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次AII-Gather，将别的GPU算好的W同步到自己这来

(13) ZeRO-3优化器状态、梯度与参数切分

流程：

①每块GPU上只保存部分参数W。将一个batch的数据分成3份，每块GPU各吃一份

②做forward时，对W做一次AII-Gather，取回分布在别的GPU上的W，得到一份完整的W。forward做完，立刻把不是自己维护的W抛弃

③做backward时，对W做一次AII-Gather，取回完整的W。backward做完，立刻把不是自己维护的W抛弃

④做完backward，算得一份完整的梯度G，对G做一次Reduce-Scatter，从别的GPU上聚合自己维护的那部分梯度。聚合操作结束后，立刻把不是自己维护的G抛弃

⑤用自己维护的O和G，更新W。由于只维护部分W，因此无需再对W做任何AIIReduce操作

---

总结

ZeRO的目标是：通过“切分”模型的不同部分(参数、梯度、优化器状态)，让每块GPU只存一部分，从而节省显存

ZeRO-1：只切分优化器状态（O）
ZeRO-2：切分优化器状态 + 梯度（O + G）
ZeRO-3：切分优化器状态 + 梯度 + 参数（O + G + W）

切得越多 -> 显存越省，但通信开销越大(因为GPU之间要频繁交换数据)

ZeRO-1(切分优化器状态)：每块GPU存完整参数W和完整梯度G，但只存部分优化器状态O

更新参数时：各GPU用自己那部分O更新对应的W(其他部分不更新) -> 最后All-Gather同步所有GPU更新后的W

省显存：优化器状态(如Adam的动量、方差)被分摊了

ZeRO-2(切分优化器状态 + 梯度)：每块GPU存完整参数W，但只存部分梯度G和部分优化器状态O

计算梯度后：各GPU只保留自己负责的那块G(其他部分丢弃) -> 用这部分G和O更新对应的W -> 最后All-Gather同步所有GPU更新后的W

省更多显存：梯度也被分摊了

 ZeRO-3(切分优化器状态 + 梯度 + 参数)：每块GPU只存部分参数W、部分梯度G和部分优化器状态O

前向/反向计算时：需要临时从其他GPU收集完整的W(计算完就丢弃) -> 梯度计算后，只保留自己负责的那块G -> 最后更新自己负责的那块W(无需同步，因为W本来就是分块存储的)

省显存最多：参数、梯度、优化器状态全被分摊了，但通信最频繁

特性	ZeRO-1	ZeRO-2	ZeRO-3
切分什么	只切分优化器状态（O）	切分 O + 梯度（G）	切分 O + G + 参数（W）
每块 GPU 存什么	完整 W + 完整 G + 部分 O	完整 W + 部分 G + 部分 O	部分 W + 部分 G + 部分 O
通信操作	AllReduce（G） + AllGather（W）	Reduce-Scatter（G） + AllGather（W）	AllGather（W，临时） + Reduce-Scatter（G）
显存占用	中等（优化器状态分摊）	较小（梯度也分摊）	最小（参数也分摊）
计算速度	快（通信少）	中等	慢（通信频繁）
适用场景	优化器状态占显存大的情况	梯度占显存大的情况	超大规模模型（参数巨多）

AIIReduce(全局聚合)：所有GPU的数据合并计算(如求和)，然后每块GPU得到完整结果

AIIGather(全局收集)：每块GPU提供一部分数据，最终所有GPU拿到完整拼接的数据

Reduce-Scatter(分散聚合)：所有 GPU 的数据先合并计算(如求和)，然后结果按块分散到不同GPU

->

ZeRO：通过切分模型状态(参数/梯度/优化器)，让100B+参数的模型能在有限GPU上跑起来 -> 用来省显存+扩规模，专门解决“显存不够，大模型训不动”的问题

DeepSpeed：整合ZeRO+其他优化技术，让训练更快、更省、更稳定

---

(14) ZeRO VS 模型并行

问题：既然ZeRO都把参数W给切了，那它应该是个模型并行呀？为什么要归到数据并行里呢？

其实ZeRO是模型并行的形式，数据并行的实质

模型并行，是指在forward和backward的过程中，只需要用自己维护的那块W来计算就行。即同样的输入X，每块GPU上各算模型的一部分，最后通过某些方式聚合结果

但对ZeRO来说，它做forward和backward的时候，是需要把各GPU上维护的W聚合起来的，即本质上还是用完整的W进行计算。它是不同的输入X，完整的参数W，最终再做聚合

(15) ZeRO-R

在ZeRO-DP提高了模型状态内存效率之后，主要消耗在激活值、临时缓冲区以及无法使用的内存碎片这三个方面的剩余内存成为次要的内存瓶颈。为了解决这个问题，我们开发了ZeRO-R来进行优化：

①通过激活值分区(Partitioned Activation Checkpointing)来优化激活值内存

· Activation Checkpointing：前向计算时，只保留部分算子的激活值。反向更新时，需要其他算子的激活值时，再重新对其进行前向计算，得到其激活值 -> 重计算

②恒定临时缓冲区大小，以在内存和计算效率之间取得平衡

③根据张量的不同生命周期来管理内存，以防止内存碎片化

(16) ZeRO-Offload和ZeRO-Infinity

ZeRO-Offload通过将数据和计算卸载到CPU来实现大规模模型训练(包括模型参数、梯度和优化器状态等模型状态)，ZeRO-Offload能够在单个GPU上训练具有超高130亿参数的模型

ZeRO-Offload的做法是：

①forward和backward计算量高，因此和它们相关的部分，例如Parameters，Activation等，就全放入GPU

②update的部分计算量低，因此和它相关的部分，全部放入CPU中。例如Optimizer States和Gradients等

ZeRO-Infinity是ZeRO-3的拓展。允许通过使用NVMe固态硬盘扩展GPU和CPU内存来训练大型模型。ZeRO-Infinity需要启用ZeRO-3

---

总结

前提知识

当前大模型训练使用的精度是什么？

(1)混合精度FP16+FP32

(2)FP16处理常规操作

(3)FP32负责参数相关数据的操作

问题

数据精度越高，在训练过程中，占用的显存量大

必存参数

(1)模型参数

(2)梯度 -> 更新模型参数

(3)优化器的1阶和2阶动量 -> 用于优化梯度更新的精度

(4)相当于训练过程中1个参数需要 -> 4个对应数值进行计算，而且数据精度都是FP32级别

ZeRO(零冗余优化器)

通过显存参数转移从而降低模型对设备显存的需求量

将显存中的参数，转移到内存中进行保存

-

ZeRO-DP(数据并行)

(1)ZeRO1：数据并行场景中，将优化器相关参数，分别拆分到不同的并行计算节点中，这样可以降低每个并行中显存的占用量

(2)ZeRO2：在ZeRO1的基础上，对梯度值进行拆分，分配到不同的计算节点，进一步降低显存占用量

(3)ZeRO3：在ZeRO2的基础上，对模型参数进行差分

注意：在降低显存占用的同时，也会增加显卡间的通信量

-

ZeRO-R

(1)处理激活值、临时缓冲区的显存占用

(2)激活值的处理方式(激活重计算)：在使用时进行计算，不需要时就从显存中删除，最大限度江都显存占用量

-

ZeRO offload

(1)将梯度、优化器更新，这类计算量相对较小的计算过程，卸载到CPU侧进行计算的一种策略方式

-

ZeRO infinity

(1)基于ZeRO3的基础上，利用固态硬扩展GPU和CPU内存来训练模型

二、DeepSpeed使用介绍

三、DeepSpeed最佳实践

(1) ZeRO策略选择总体原则

速度方面(左边比右边快)：stage 0(DDP)数据并行 > stage 1 > stage 2 > stage2 + offload > stage 3 > stage3 + offload

GPU内存使用方面(右边比左边更节省GPU内存)：stage 0(DDP) < stage 1 < stage 2 < stage 2 + offload < stage 3 < stage 3 + offload

所以，当你希望在尽量使用较少数量的GPU的同时获得最快的执行速度时，可以按照以下步骤进行。我们从最快的方法开始，如果遇到GPU内存溢出(内存溢出)，然后切换到下一个速度较慢但使用的GPU内存更少的方法，以此类推

---

总结

ZeRO stage0就是基础的数据并行策略

速度对比(越靠前，训练速度越快)：stage0>stage1>stage2>stage2+offload>stage 3>stage3+offload

显存占用(越靠前，显存需求越高)：sstage0>stage1>stage2>stage2+offload>stage 3>stage3+offload

注意项：如果在推理阶段，可以ZeRO stage3进行优化