模型并行下操作ZeRO进行显存优化

上帝也疯狂

前言

本文先简要介绍目前主流模型训练显存的优化方式，以及通信数据量与冗余显存之间的关系。最后重点介绍通过 ZeRO进行显存优化的最新工作。ZeRO是一种显存优化技术，能够通过将模型参数、梯度和优化器状态划分到多个GPU上，降低单个GPU的显存使用量，从而撑持更大规模的深度学习模型训练。

阅读本文前建议读者先了解混合精度训练、Adam优化器，以及模型训练常用显存的布局的相关概念（见文末附录）。

分布式深度学习计算

在开始介绍ZeRO前，有必要回顾一下目前分布式深度学习计算的方式。分布式深度学习计算是指将深度学习模型和数据划分到多个设备或计算节点上，操作并行计算来加速深度学习模型的训练和推理过程。常见的分布式深度学习计算包罗数据并行、模型并行、流水并行三种方式。
数据并行（data parallelism, DP）

数据并行是将大规模数据集拆分成多个小的数据集，分配给分歧的措置单元并同时进行计算。这种方式可以提高计算效率，出格是在需要措置大量数据的场所。
数据并行的主要思想是将大数据集划分成多个小数据集，然后将这些小数据集分配给分歧的措置单元，每个措置单元独登时对其分配的数据进行措置，最后将措置成果合并起来得到最终成果。
假设有 N 张卡，每张卡都保留一个模型，每一次迭代（iteration/step）都将batch数据分割成 N 个等大小的micro-batch，每张卡按照拿到的micro-batch数据独立计算梯度，然后调用AllReduce计算梯度均值，每张卡再独立进行参数更新。

1. # 假设模型有三层：L0, L1, L2
2. # 每层有两个神经元
3. # 两张卡
4. GPU0:
5. L0 | L1 | L2
6. ---|----|---
7. a0 | b0 | c0
8. a1 | b1 | c1
9. GPU1:
10. L0 | L1 | L2
11. ---|----|---
12. a0 | b0 | c0
13. a1 | b1 | c1

复制代码

模型并行（model parallelism/tensor parallelism, MP/TP）

模型并行是将大型神经网络模型分化成多个子模型，每个子模型分配给分歧的措置单元进行计算。这种方式可以使得每个措置单元只需要措置一部门模型参数和数据，从而减轻了计算承担，提高了训练效率。
在模型并行中，每个子模型负责措置整个模型的一部门参数和数据，然后将措置成果汇总到一起，以得到最终的训练成果。因此，模型并行需要高效的通信协议和通信硬件撑持，以确保措置单元之间的通信效率和正确性。
模型并行凡是用于训练大型神经网络模型，出格是在GPU等高性能计算设备上。由于每个子模型只需要措置部门参数和数据，因此可以将大型模型分化成更小的子模型，以适应更小的内存和计算资源。此外，模型并行也可以用于措置超大规模的数据集，以提高训练效率。

1. # 假设模型有三层：L0, L1, L2
2. # 每层有两个神经元
3. # 两张卡
4. GPU0:
5. L0 | L1 | L2
6. ---|----|---
7. a0 | b0 | c0
8. GPU1:
9. L0 | L1 | L2
10. ---|----|---
11. a1 | b1 | c1

复制代码

流水并行（pipeline parallelism, PP）

流水并行是将一个大型计算任务拆分成多个小的子任务，并将这些子任务在多个措置单元上同时执行。分歧于数据并行和模型并行，流水并行不是将数据或模型分割成多个部门并在措置单元间并行措置，而是将一系列计算法式分化成多个流水阶段，并在多个措置单元上同时执行，以减少总体计算时间。
在流水并行中，每个措置单元都负责执行一个或多个计算法式，并将计算成果传递给下一个措置单元。因此，流水并行需要高效的通信协议和通信硬件撑持，以确保各个措置单元之间的通信效率和正确性。
流水并行凡是用于措置需要多个计算法式的任务，例如图像措置、信号措置、视频编码等。在这些任务中，每个计算法式的计算时间相对较短，因此通过流水并行可以将计算时间缩短到最小，提高计算效率。

1. # 假设模型有8层
2. # 两张卡
3. ======================  =====================
4. |  L0 | L1 | L2 | L3 |  | L4 | L5 | L6 | L7 |
5. ======================  =====================GPU0                 GPU1
6. # 设想一下，当GPU0在进行（前向/后向）计算时，GPU1在干嘛？闲着
7. # 当GPU1在进行（前向/后向)计算时，GPU0在干嘛？闲着
8. # 为了防止”一卡工作，众卡围不雅观“，实践中PP也会把batch数据分割成
9. # 多个micro-batch，流水线执行

复制代码

相对于流水线并行和数据并行，模型并行具有以下长处：

• 撑持更大的模型规模：流水线并行和数据并行的限制凡是是 GPU 内存大小和 GPU 数量，而模型并行可以撑持更大的模型规模，因为模型可以分割成多个子模型，并分配到多个 GPU 上运行。
  
• 灵活的模型分配：模型并行可以更灵活地将模型分配给分歧的 GPU ，这意味着可以在分歧的 GPU 上运行分歧的模型子集，从而实现更好的负载平衡和性能优化。
  

模型并行下的通信数据量分析

无论是使用哪种分布式深度学习训练方式，通信数据量都是影响训练速度和效率的重要因素之一。
通信数据量指的是在分布式训练中，在分歧的计算节点或GPU之间交换的数据量。通信数据量过大会导致通信延迟增加，进而影响训练速度和效率。
在并行计算中，凡是需要进行数据的全局同步和聚合操作，此中最常用的两个操作是All-Reduce和All-Gather。
All-Reduce

All-Reduce 操作将每个节点的局部计算成果汇总到全局成果，凡是用于并行计算中的梯度求和、模型参数更新等操作。例如，当使用分布式深度学习框架进行训练时，每个计算节点将负责计算一部门训练样本的梯度，并使用 all-reduce 操作将这些梯度汇总到全局梯度，以更新模型参数。



图1

在 All-Reduce 操作中，Reduce 和 Broadcast 是 All-Reduce 的两个基本法式，它们凡是由分歧的算法实现。
Reduce 法式将每个计算节点的当地梯度进行求和或其他的归约操作，并将成果发送给一个特定的节点，凡是是根节点。Reduce 的算法有很多种，此中一种常见的算法是 ring-reduce 算法。在 ring-reduce 算法中，计算节点按照一个环形拓扑布局进行通信，每个节点将当地梯度向右传递，直到最后一个节点将梯度发送给根节点进行汇总。



图2

Broadcast 法式将归约成果发送回所有计算节点，以便更新模型参数。Broadcast 的算法也有很多种，此中一种常见的算法是 ring-broadcast 算法。在 ring-broadcast 算法中，根节点将全局成果向左传递，直到每个计算节点都收到了成果。



图3

Reduce-scatter 的实现和性质
当需要对一个大数据集进行归约操作时，如果使用单个措置器执行归约操作，可能会变得非常慢。这时就可以使用并行计算来加速归约操作。Reduce-scatter 算法是一种并行归约算法，它将大数据集分成多个部门，分配给分歧的措置器执行局部归约操作，最后将归约成果发送给其他措置器，以便计算全局归约成果。
具体来说，Reduce-scatter 算法的流程如下：

• 将数据分配给分歧的措置器。
  
• 在每个措置器上计算局部归约成果。例如，如果要对一个数组中的元素求和，则每个措置器将对它分配到的数组部门求和。
  
• 将每个措置器的局部归约成果发送给所有其他措置器。例如，如果有 4 个措置器，每个措置器将其局部归约成果发送给其他 3 个措置器。
  
• 在每个措置器上计算全局归约成果。例如，如果要对数组中的元素求和，则每个措置器将对所有局部归约成果进行求和。
  
• 将全局归约成果发送给所有其他措置器。
  
• 在每个措置器上计算 Reduce-scatter 的成果。这个法式将全局归约成果平均分配给分歧的措置器，以便计算 Reduce-scatter的成果。例如，如果有 4 个措置器，每个措置器将得到全局归约成果的 1/4。
  

AllGather




图4

All-Gather 操作将每个节点的局部数据收集到全局数据布局中，凡是用于数据并行计算中。例如，当需要在多个计算节点之间共享某些数据时，可以使用 All-Gather 操作将每个节点的数据收集到一个全局数组中。这个全局数组可以被所有节点访谒，以便进行并行计算。
在执行 All-Reduce 和 All-Gather 操作时，分歧的计算节点需要进行数据通信以交换数据。这凡是需要高效的通信协议和算法，以最小化通信开销和延迟。常见的通信协议包罗基于TCP/IP 的网络协议、InfiniBand 等高速网络协议以及 RDMA（Remote Direct Memory Access）等技术。
通信量和冗余显存之间的关系

通信量和冗余显存之间存在必然的关系，简单的说，可以通过降低通信量，从而减少冗余显存的使用。
在深度学习模型训练过程中，通信量凡是是指在模型并行化训练中，分歧计算节点之间传递的参数数据量。如果通信量较大，那么节点之间需要频繁地进行数据传输和同步，这会导致训练时间的增加，同时还会占用较多的网络带宽资源。
相反，冗余显存指的是在模型并行化训练中，由于需要将模型参数复制到每个计算节点上，因此需要额外占用显存资源。因此，较大的冗余显存使用可能会导致显存不足的问题，限制模型的规模和训练速度。
降低通信量和冗余显存的使用的技术有ZeRO技术中的模型并行交叉复制和ZeRO-Offload技术。这些技术可以将模型参数复制到每个计算节点上，并允许每个节点独登时更新模型参数，从而减少了节点之间的通信和同步。此外，ZeRO-Offload技术还可以将模型参数存储在内存中，减少了额外的显存使用。即，在保证模型训练质量的前提下，降低通信量和冗余显存的使用，提高模型训练的效率。
ZeRO

ZeRO（Zero Redundancy Optimizer）是一个由NVIDIA开发的分布式深度学习训练技术，旨在解决在大规模模型上训练时由于显存限制而导致的性能瓶颈问题。
在传统的分布式训练中，每个工作节点都必需存储完整的模型参数副本，这意味着在使用大型模型时，每个工作节点需要拥有足够的显存才能存储这些参数。而ZeRO技术通过将模型参数分成多个分片，让每个工作节点只需要存储部门参数，从而显著减少了显存占用量。
具体而言，ZeRO技术通过以下三个主要组件实现：

• ZeRO-Stage：将模型参数划分为更小的分片，每个工作节点只需存储本身所负责的参数分片。
  
• ZeRO-Offload：将一部门模型参数存储在CPU的内存中，从而释放显存空间。
  

通过使用ZeRO技术，可以大幅度提高分布式深度学习训练的效率和规模。同时，由于减少了显存占用量，还可以使用更大的批量大小，从而加速训练过程。
ZeRO-Stage

ZeRO将模型训练阶段，每张卡中显存内容分为两类：

• 模型状态（model states）: 模型参数（fp16）、模型梯度（fp16）和Adam状态（fp32的模型参数备份，fp32的momentum和fp32的variance）。假设模型参数量 Φ ，则共需要 2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字节存储，可以看到，Adam状态占比 75% 。
  
• 残剩状态（residual states）: 除了模型状态之外的显存占用，包罗激活值（activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）。
  

针对模型状态的存储优化（去除冗余），ZeRO使用的方式是分片（partition），即每张卡只存 1/N 的模型状态量，这样系统内只维护一份模型状态。

• 首先进行分片操作的是模型状态中的Adam，也就是图5中的 Pos ，这里os指的是optimizer states。模型参数（parameters）和梯度（gradients）仿照照旧是每张卡保持一份，此时，每张卡的模型状态所需显存是 4Φ+12Φ/N 字节，当 N 斗劲大时，趋向于 4ΦB ，也就是本来 16ΦB 的 1/4 。
  
• 如果继续对模型梯度进行分片，也就是图5中的 Pos+g ，模型参数仿照照旧是每张卡保持一份，此时，每张卡的模型状态所需显存是 2Φ+(2Φ+12Φ)/N 字节，当 N 斗劲大时，趋向于 2ΦB ，也便是本来 16ΦB 的 1/8 。
  
• 如果继续对模型参数进行分片，也就是图5中的 Pos+g+p ，此时每张卡的模型状态所需显存是 16Φ/N 字节，当 N 斗劲大时，趋向于 0 。
  

下面我们就分析下通信数据量，先说结论， Pos 和 Pos+g 的通信量和传统数据并行不异，Pos+g+p 会增加通信量。
传统数据数据并行在每一步（step/iteration）计算梯度后，需要进行一次AllReduce操作来计算梯度均值，目前常用的是Ring AllReduce，分为ReduceScatter和AllGather两步，每张卡的通信数据量（发送+接收）近似为 2Φ ([2])。
我们直接分析 Pos+g ，每张卡只存储 1N 的优化器状态和梯度，对于 gpu0 来说，为了计算它这 1N 梯度的均值，需要进行一次Reduce操作，通信数据量是 1/N⋅Φ⋅N=Φ ，然后其余显卡则不需要保留这部门梯度值了。实现中使用了bucket策略，保证 1N 的梯度每张卡只发送一次。
当 gpu0 计算好梯度均值后，就可以更新局部的优化器状态（包罗 1/N⋅Φ 的参数），当反向传布过程结束，进行一次Gather操作，更新 (1−1/N)Φ 的模型参数，通信数据量是 1/N⋅Φ⋅N=Φ 。
从全局来看，相当于用Reduce-Scatter和AllGather两步，和数据并行一致，使得每张卡只存了 1/N 的参数，不管是在前向计算还是反向传布，都涉及一次Broadcast操作。



图5

解决了模型状态，再来看残剩状态，也就是激活值（activation）、临时缓冲区（buffer）以及显存碎片（fragmentation）。

• 激活值同样使用分片方式，而且配合checkpointing
  
• 模型训练过程中经常会创建一些大小不等的临时缓冲区，比如对梯度进行AllReduce啥的，解决法子就是预先创建一个固定的缓冲区，训练过程中不再动态创建，如果要传输的数据较小，则多组数据bucket后再一次性传输，提高效率
  
• 显存呈现碎片的一大原因是时候gradient checkpointing后，不竭地创建和销毁那些不保留的激活值，解决方式是预先分配一块持续的显存，将常驻显存的模型状态和checkpointed activation存在里面，残剩显存用于动态创建和销毁discarded activation
  

ZeRO-Offload

ZeRO-Offload是ZeRO（Zero Redundancy Optimizer）技术的一种扩展，它使用显存作为模型参数存储和通信的中间介质，以减少模型并行化训练中的通信和同步开销。
ZeRO-Offload技术使用显存缓存将模型参数存储在显存中，这可以减少网络带宽的使用，同时还可以加速参数访谒和更新。为了最大限度地减少显存的使用，ZeRO-Offload技术使用了一种称为“按需加载”的策略。这种策略只在需要使用参数时才将其从磁盘或网络加载到显存中，而不是一次性将所有参数都加载到显存中。



图6

Offload策略
ZeRO-Offload技术的核心是使用显存缓存和显存内通信来降低通信开销。为了最大程度地操作显存并减少网络带宽的使用，ZeRO-Offload技术采用了一种称为“Offload策略”的技术。下面是ZeRO-Offload技术的Offload策略的几个关键点：

• 按需加载
  

ZeRO-Offload技术使用“按需加载”策略，只在需要使用参数时才将其从磁盘或网络加载到显存中，而不是一次性将所有参数都加载到显存中。这种策略可以最大限度地减少显存的使用，并减少网络带宽的使用。
2. 数据流水线
ZeRO-Offload技术使用“数据流水线”策略，将数据流分成多个阶段，每个阶段都使用分歧的计算资源进行措置。在模型训练期间，ZeRO-Offload技术将数据分为多个块，并将这些块分配给分歧的GPU进行计算。每个GPU只对其分配的数据块进行计算，并将计算成果传递给下一个阶段的GPU，直到所有阶段都完成为止。
3. 显存原语
ZeRO-Offload技术使用一种称为“显存原语”的通信协议，在显存中直接进行通信和同步操作，而不需要通过网络或主机内存。这种协议可以显著减少通信延迟和数据传输时间，从而提高训练效率。
4. 数据切片
ZeRO-Offload技术使用“数据切片”策略来划分模型参数，并通过显存内通信来实现分歧GPU上参数的同步。具体来说，ZeRO-Offload技术将模型参数划分为多个小块，并在每个GPU上存储一部门参数块。在训练过程中，每个GPU只对其分配的参数块进行计算，并通过显存内通信将参数块传输到其他GPU长进行同步。
总的来说，ZeRO-Offload技术的Offload策略通过按需加载、数据流水线、显存原语和数据切片等技术手段来最大化地操作显存，并降低通信开销，从而提高深度学习模型训练的效率和可扩展性。
为了找到最优的offload策略，ZeRO作者将模型训练过程看作数据流图（data-flow graph）。

• 圆形节点暗示模型状态，比如参数、梯度和优化器状态
  
• 矩形节点暗示计算操作，比如前向计算、后向计算和参数更新
  
• 边暗示数据流向
  

下图是某一层的一次迭代过程（iteration/step），使用了混合精读训练，前向计算（FWD）需要用到上一次的激活值（activation）和本层的参数（parameter），反向传布（BWD）也需要用到激活值和参数计算梯度，



图7

如果用Adam优化器进行参数更新（Param update），流程如下：



图8

下面我们为边添加权重，物理含义是数据量大小（单元是字节），假设模型参数量是 M ，在混合精度训练的前提下，边的权重要么是2M（fp16），要么是4M（fp32）。



图9

我们此刻要做的就是沿着边把数据流图切分为两部门，分布对应GPU和CPU，计算节点（矩形节点）落在哪个设备，哪个设备就执行计算，数据节点（圆形）落在哪个设备，哪个设备就负责存储，将被切分的边权重加起来，就是CPU和GPU的通信数据量。
ZeRO-Offload的切分思路如图 10 所示：



图10

图10中有四个计算类节点：FWD、BWD、Param update和float2half，前两个计算复杂度大致是 O(MB) ， B 是batch size，后两个计算复杂度是 O(M) 。为了不降低计算效率，将前两个节点放在GPU，后两个节点不单计算量小还需要和Adam状态打交道，所以放在CPU上，Adam状态自然也放在内存中，为了简化数据图，将前两个节点融合成一个节点FWD-BWD Super Node，将后两个节点融合成一个节点Update Super Node。
所以，此刻的计算流程是，在GPU上面进行前向和后向计算，将梯度传给CPU，进行参数更新，再将更新后的参数传给GPU。为了提高效率，可以将计算和通信并行起来，GPU在反向传布阶段，可以待梯度值填满bucket后，一遍计算新的梯度一遍将bucket传输给CPU，当反向传布结束，CPU基本上已经有最新的梯度值了，同样的，CPU在参数更新时也同步将已经计算好的参数传给GPU，如下图所示：



图11

总结

本文主要介绍从模型并行下操作 ZeRO 方式进行显存优化的方式，可以大大提高训练速度、降低内存使用量，并撑持大型模型训练。除了ZeRO 外还有 Megatron，Gpip 和 Mesh-TensorFlow 等分布式深度计算方式，大师也可以自行了解，也欢迎在评论区中交流探讨。
此外FlagAI也集成了ZeRO的使用方式，具体实现方式使用的是操作Deepspeed 实现ZeRO1和ZeRO2，操作bmtrain实现了ZeRO3。

FlagAI 由北京智源人工智能研究院于 2022 年 5 月开源，是大模型算法、模型及各种优化东西的一站式、高质量开源项目，旨在集成全球各种主流大模型算法技术，以及多种大模型并行措置和训练加速技术，撑持高效训练和微调，降低大模型开发和应用的门槛，提高大模型的开发效率。项目地址：https://github.com/FlagAI-Open/FlagAI

附录

混合精度训练和Adam优化器

混合精度训练（mixed precision training）和Adam优化器基本上已经是训练语言模型的标配，我们先来简单回顾下相关概念。
Adam在SGD基础上，为每个参数梯度增加了一阶动量（momentum）和二阶动量（variance）([1])。
混合精度训练，字如其名，同时存在fp16和fp32两种格式的数值，此中模型参数、模型梯度都是fp16，此外还有fp32的模型参数，如果优化器是Adam，则还有fp32的momentum和variance。



图12

模型训练常用显存

fp16
在单精度32位格式中，1位用于指示数字为正数还是负数。指数保留了8位，这是因为它为二进制，将2进到高位，其余23位用于暗示组成该数字的数字，称为有效数字。
而在双精度下，指数保留11位，有效位数为52位，从而极大地扩展了它可以暗示的数字范围和大小。半精度则是暗示范围更小，其指数只有5位，有效位数只有10位。
半精度的格式与单精度的格式类似，最左边的一位仍是符号位，指数有5位宽且以余-16（excess-16）的形式存储，尾数有10位宽，但具有隐含1。



图13

如图所示，sign为符号位，0暗示这个浮点数为正，1暗示这个浮点数为负
先介绍尾数，再说指数，fraction为尾数，有10位长，但是有隐含1，尾数可以理解为是一个浮点数小数点后的数，如1.11，尾数就为1100000000（1），最后的隐含1主要用于计算时，隐含1可能存在可以进位的情况。
exponent为指数位，有5位长，具体暗示的值有以下几种情况：

1. 当指数位全为0 ，尾数位也全为0的时，暗示的就是0
2. 当指数位全为0，尾数位不全为0时，暗示为subnormal value，非规格化浮点数，是一个非常小的数
3. 当指数位全为1，尾数位全为0时，暗示的是无穷大，此时如果符号位为0，暗示正无穷，符号位为1，暗示负无穷
4. 当指数位全为1，尾数位不全为0时，暗示的不是一个数

复制代码

其余情况下，指数位的值减去15就是其暗示的指数，如11110暗示的就是30-15=15
所以我们可以得到，半精度浮点数的值得计算方式为（-1）^sign×2^（指数位的值）×（1+0.尾数位）
备注：这里0.尾数位，暗示如尾数位为0001110001，则0.尾数位为0.0001110001
例子

1. 半精度可以暗示的最大值：
2. 0 11110 1111111111 计算方式为：
3. （-1）^0×2^(30-15)×1.1111111111 = 1.1111111111(b)×2^15 = 1.9990234375(d)×2^15 = 65504
4. 半精度可以暗示的最小值（除了subnormal  value）：
5. 0 00001 0000000000 计算方式为：（-1）^（-1）×2(1-15)=2^(-14),约等于十进制的6.104×10^(-5)

复制代码

bf16



图14

fp32
第1位暗示正负，中间8位暗示指数，后23位储存有效数位（有效数位是24位）。
第一位的正负号0代表正，1代表负。
中间八位共可暗示2(8)=256个数，指数可以是二补码；或0到255，0到126代表-127到-1，127代表零，128-255代表1-128。
有效数位最左手边的1并不会储存，因为它必然存在（二进制的第一个有效数字必定是1）。换言之，有效数位是24位，实际储存23位。



图15

1. {\displaystyle {\text{sign}}=+1}
2. {\displaystyle {\text{exponent}}=(-127)+124=-3}{\displaystyle {\text{exponent}}=(-127)+124=-3}
3. {\displaystyle {\text{fraction}}=1+2^{-2}=1.25}{\text{fraction}}=1+2^{{-2}}=1.25
4. {\displaystyle {\text{value}}=(+1)\times 1.25\times 2^{-3}=+0.15625}{\displaystyle {\text{value}}=(+1)\times 1.25\times 2^{-3}=+0.15625}

复制代码

相关论文

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
https://www.usenix.org/system/files/atc21-ren-jie.pdf
https://arxiv.org/abs/1710.03740