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

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

数据并行是将大规模数据集拆分成多个小的数据集，分配给不同的处理单元并同时进行计算。这种方法可以提高计算效率，特别是在需要处理大量数据的场合。

数据并行的主要思想是将大数据集划分成多个小数据集，然后将这些小数据集分配给不同的处理单元，每个处理单元独立地对其分配的数据进行处理，最后将处理结果合并起来得到最终结果。

假设有 N 张卡，每张卡都保存一个模型，每一次迭代（iteration/step）都将batch数据分割成 N 个等大小的micro-batch，每张卡根据拿到的micro-batch数据独立计算梯度，然后调用AllReduce计算梯度均值，每张卡再独立进行参数更新。

模型并行是将大型神经网络模型分解成多个子模型，每个子模型分配给不同的处理单元进行计算。这种方法可以使得每个处理单元只需要处理一部分模型参数和数据，从而减轻了计算负担，提高了训练效率。

在模型并行中，每个子模型负责处理整个模型的一部分参数和数据，然后将处理结果汇总到一起，以得到最终的训练结果。因此，模型并行需要高效的通信协议和通信硬件支持，以确保处理单元之间的通信效率和正确性。

模型并行通常用于训练大型神经网络模型，特别是在GPU等高性能计算设备上。由于每个子模型只需要处理部分参数和数据，因此可以将大型模型分解成更小的子模型，以适应更小的内存和计算资源。此外，模型并行也可以用于处理超大规模的数据集，以提高训练效率。

流水并行是将一个大型计算任务拆分成多个小的子任务，并将这些子任务在多个处理单元上同时执行。不同于数据并行和模型并行，流水并行不是将数据或模型分割成多个部分并在处理单元间并行处理，而是将一系列计算步骤分解成多个流水阶段，并在多个处理单元上同时执行，以减少总体计算时间。

在流水并行中，每个处理单元都负责执行一个或多个计算步骤，并将计算结果传递给下一个处理单元。因此，流水并行需要高效的通信协议和通信硬件支持，以确保各个处理单元之间的通信效率和正确性。

流水并行通常用于处理需要多个计算步骤的任务，例如图像处理、信号处理、视频编码等。在这些任务中，每个计算步骤的计算时间相对较短，因此通过流水并行可以将计算时间缩短到最小，提高计算效率。

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

无论是使用哪种分布式深度学习训练方式，通信数据量都是影响训练速度和效率的重要因素之一。

通信数据量指的是在分布式训练中，在不同的计算节点或GPU之间交换的数据量。通信数据量过大会导致通信延迟增加，进而影响训练速度和效率。

在并行计算中，通常需要进行数据的全局同步和聚合操作，其中最常用的两个操作是All-Reduce和All-Gather。

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

在 All-Reduce 操作中，Reduce 和 Broadcast 是 All-Reduce 的两个基本步骤，它们通常由不同的算法实现。

Reduce 步骤将每个计算节点的本地梯度进行求和或其他的归约操作，并将结果发送给一个特定的节点，通常是根节点。Reduce 的算法有很多种，其中一种常见的算法是 ring-reduce 算法。在 ring-reduce 算法中，计算节点按照一个环形拓扑结构进行通信，每个节点将本地梯度向右传递，直到最后一个节点将梯度发送给根节点进行汇总。

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

Reduce-scatter 的实现和性质

当需要对一个大数据集进行归约操作时，如果使用单个处理器执行归约操作，可能会变得非常慢。这时就可以使用并行计算来加速归约操作。Reduce-scatter 算法是一种并行归约算法，它将大数据集分成多个部分，分配给不同的处理器执行局部归约操作，最后将归约结果发送给其他处理器，以便计算全局归约结果。

具体来说，Reduce-scatter 算法的流程如下：

All-Gather 操作将每个节点的局部数据收集到全局数据结构中，通常用于数据并行计算中。例如，当需要在多个计算节点之间共享某些数据时，可以使用 All-Gather 操作将每个节点的数据收集到一个全局数组中。这个全局数组可以被所有节点访问，以便进行并行计算。

在执行 All-Reduce 和 All-Gather 操作时，不同的计算节点需要进行数据通信以交换数据。这通常需要高效的通信协议和算法，以最小化通信开销和延迟。常见的通信协议包括基于TCP/IP 的网络协议、InfiniBand 等高速网络协议以及 RDMA（Remote Direct Memory Access）等技术。

通信量和冗余显存之间存在一定的关系，可以通过使用一些技术来降低通信量，从而减少冗余显存的使用。

在深度学习模型训练过程中，通信量通常是指在模型并行化训练中，不同计算节点之间传递的参数数据量。如果通信量较大，那么节点之间需要频繁地进行数据传输和同步，这会导致训练时间的增加，同时还会占用较多的网络带宽资源。

相反，冗余显存指的是在模型并行化训练中，由于需要将模型参数复制到每个计算节点上，因此需要额外占用显存资源。因此，较大的冗余显存使用可能会导致显存不足的问题，限制模型的规模和训练速度。

为了降低通信量和冗余显存的使用，一些技术可以被使用，如ZeRO技术中的模型并行交叉复制和ZeRO-Offload技术。这些技术可以将模型参数复制到每个计算节点上，并允许每个节点独立地更新模型参数，从而减少了节点之间的通信和同步。此外，ZeRO-Offload技术还可以将模型参数存储在显存中，减少了额外的显存使用。

因此，通过使用这些技术，可以在保证模型训练质量的前提下，降低通信量和冗余显存的使用，提高模型训练的效率。

ZeRO（Zero Redundancy Optimizer）是一个由NVIDIA开发的分布式深度学习训练技术，旨在解决在大规模模型上训练时由于显存限制而导致的性能瓶颈问题。

在传统的分布式训练中，每个工作节点都必须存储完整的模型参数副本，这意味着在使用大型模型时，每个工作节点需要拥有足够的显存才能存储这些参数。而ZeRO技术通过将模型参数分成多个分片，让每个工作节点只需要存储部分参数，从而显著减少了显存占用量。

具体而言，ZeRO技术通过以下三个主要组件实现：

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

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

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

下面我们就分析下通信数据量，先说结论， 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操作。

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

ZeRO-Offload是ZeRO（Zero Redundancy Optimizer）技术的一种扩展，它使用显存作为模型参数存储和通信的中间介质，以减少模型并行化训练中的通信和同步开销。

ZeRO-Offload技术使用显存缓存将模型参数存储在显存中，这可以减少网络带宽的使用，同时还可以加速参数访问和更新。为了最大限度地减少显存的使用，ZeRO-Offload技术使用了一种称为“按需加载”的策略。这种策略只在需要使用参数时才将其从磁盘或网络加载到显存中，而不是一次性将所有参数都加载到显存中。

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）也需要用到激活值和参数计算梯度，

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

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

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

ZeRO-Offload的切分思路如图 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，如下图所示：

本文主要介绍从模型并行下利用 ZeRO 方法进行显存优化的方法，可以大大提高训练速度、降低内存使用量，并支持大型模型训练。除了ZeRO 外还有 Megatron，Gpip 和 Mesh-TensorFlow 等分布式深度计算方式，大家也可以自行了解，也欢迎在评论区中交流探讨。

另外FlagAI也集成了ZeRO的使用方式，具体实现方式使用的是利用Deepspeed 实现ZeRO1和ZeRO2，利用bmtrain实现了ZeRO3。

混合精度训练（mixed precision training）和Adam优化器基本上已经是训练语言模型的标配，我们先来简单回顾下相关概念。

Adam在SGD基础上，为每个参数梯度增加了一阶动量（momentum）和二阶动量（variance）([1])。

混合精度训练，字如其名，同时存在fp16和fp32两种格式的数值，其中模型参数、模型梯度都是fp16，此外还有fp32的模型参数，如果优化器是Adam，则还有fp32的momentum和variance。

fp16

在单精度32位格式中，1位用于指示数字为正数还是负数。指数保留了8位，这是因为它为二进制，将2进到高位，其余23位用于表示组成该数字的数字，称为有效数字。

而在双精度下，指数保留11位，有效位数为52位，从而极大地扩展了它可以表示的数字范围和大小。半精度则是表示范围更小，其指数只有5位，有效位数只有10位。

半精度的格式与单精度的格式类似，最左边的一位仍是符号位，指数有5位宽且以余-16（excess-16）的形式存储，尾数有10位宽，但具有隐含1。

如图所示，sign为符号位，0表示这个浮点数为正，1表示这个浮点数为负

先介绍尾数，再说指数，fraction为尾数，有10位长，但是有隐含1，尾数可以理解为是一个浮点数小数点后的数，如1.11，尾数就为1100000000（1），最后的隐含1主要用于计算时，隐含1可能存在可以进位的情况。

exponent为指数位，有5位长，具体表示的值有以下几种情况：

其余情况下，指数位的值减去15就是其表示的指数，如11110表示的就是30-15=15

所以我们可以得到，半精度浮点数的值得计算方式为（-1）^sign×2^（指数位的值）×（1+0.尾数位）

备注：这里0.尾数位，表示如尾数位为0001110001，则0.尾数位为0.0001110001

例子

bf16

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位。

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

https://www.usenix.org/system/files/atc21-ren-jie.pdf

https://arxiv.org/abs/1710.03740