Deepspeed并行框架介绍--ZeRO去除冗余的并行方案

Deepspeed 并行框架介绍

一. 简介

Deepspeed 是微软推出的一个开源分布式工具，其集合了分布式训练、推断、压缩等高效模块。
该工具旨在提高大规模模型训练的效率和可扩展性。

它通过多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等。

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

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

二. 分布式训练方法

如今的大模型训练，离不开各种分布式的训练框架，一般来说，并行策略包含：数据并行、模型并行、流水线并行。

1. 数据并行

数据并行分为了两种模式：Data Parallel（DP）和 Distributed Data Parallel（DDP）。

Data Parallel（DP）

DP 是一种单进程多线程的并行策略，只能在单机上进行训练，步骤如下：

• 单进程控制多 GPU，即本质上是单进程多线程；
• 首先将模型加载到主 GPU 上，再复制到各个指定从 GPU；
• 将输入数据按照 Batch 维度进行拆分，各个 GPU 独立进行 forward 计算；
• 将结果同步给主 GPU 完成梯度计算和参数更新，将更新后的参数复制到各个 GPU。
  由于其是单进程控制多个 GPU，故会存在 GPU 之间负载不均衡的问题，主 GPU 负载较大。

Distributed Data Parallel（DDP）

DDP 采用 AllReduce 架构，多进程的方式，突破锁的束缚。在单机和多机上都可以使用。

负载分散在每个 GPU 节点上，通信成本（时间）是恒定的，与 GPU 数量无关，等于 V/B（参数量 / 带宽）。

DDP 不需要通过主 GPU 分发全模型的参数到每个 GPU 上。使用 ring-all-reduce 的方式进行通讯，随着 GPU 数量 N 增加，总传输量恒定。也就是理论上，随着 GPU 数量的增加，ring all-reduce 有线性加速能力。

2. 张量并行

张量并行的原理是，将张量操作划分到多个设备上，以加速计算或增加模型大小；对模型每一层的层内参数进行切分，即对参数矩阵切片，并将不同切片放到不同 GPU 上；将原本在单卡中的矩阵乘法，切分到不同卡中进行矩阵乘法。训练过程中，正向和反向传播计算出的数据通过使用 All gather 或者 All reduce 的方法完成整合。

以 transformer 为例，该策略会把 Masked Multi Self Attention 和 Feed Forward 都进行切分以并行化。利用 Transformers 网络的结构，通过添加一些同步原语来创建一个简单的模型并行实现。

张量并行适用于模型单层网络参数较大的情况。同时缺点也是十分明显：

• 若环境是多机多卡，张量并行所需的 all-reduce 通信需要跨服务器进行链接，这比单机多 GPU 服务器内的高带宽通信要慢；
• 高度的模型并行会产生很多小矩阵乘法，这可能会降低 GPU 的利用率。

3. 流水线并行

流水线原理是将不同的 layer 分配给指定 GPU 进行计算，流水线并行只需其之间点对点地通讯传递部分 activations。

具体步骤包括：

• 在流水线并行之中，一个模型的各层会在多个 GPU 上做切分。
• 一个批次（batch）被分割成较小的微批（microbatches），并在这些微批上进行流水线式执行。
• 通过流水线并行，一个模型的层被分散到多个设备上。
• 当用于具有相同 transformer 块重复的模型时，每个设备可以被分配相同数量的 transformer 层。
• 在流水线模型并行中，训练会在一个设备上执行一组操作，然后将输出传递到流水线中下一个设备，下一个设备将执行另一组不同操作。

流水线并行的方法，解决了超大模型无法在单设备上装下的难题，也解决了机器之间的通信开销的问题，使得每台机器的数据传输量跟总的网络大小、机器总数、并行规模无关。

常用的流水线方法有 G-pipe、PipeDream、virtual pipeline 等。

三. ZeRO

ZeRO（Zero Redundancy Optimizer）是一种去除冗余的并行方案，来自微软在 SC 20 上发表的论文 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. 而 Deepspeed 库最初的就是关于 ZeRO 的官方实现。也是 Deepspeed 实现 3D 并行（数据、模型、流水线）的主要模块。

1. 模型的显存占用

让我们来看看，当使用 GPU 训练模型的时候，显存中都被哪些东西占用。
从下图可以看到，GPU 需要存储优化器状态（Optimizer States）、模型参数（Model Parameters）、激活值（Activations）、梯度（Gradients）、临时缓存（Temporary Buffers）等。

模型参数只是占用其中的一部分，当使用混合精度进行训练时，模型状态（Model Parameters+Optimizer States+Gradients）会站到一大半以上。

模型参数（Model Parameters）、梯度（Gradients）会使用 FP16 精度进行存储。优化器状态（Optimizer States）是进行梯度更新时用到的数据，例如使用 Adam 优化器时，除了需要保存以 FP32 精度存储的模型参数外，还需要以 FP32 的精度存储 Variance 和 Momentum 的参数值。

2.ZeRO-1/2/3

微软提出的 ZeRO 针对并行训练的场景，对模型状态（Model Parameters+Optimizer States+Gradients）提出了 3 种不同程度的分割。目的在于将数据、以及模型本身的参数、优化器的状态、激活函数的输出值、梯度等切分并放在不同的 GPU 上，以此实现并行训练。

ZeRO 包含 3 种级别：

• ZeRO-1 : Optimizer States Partitioning（P_os）
• ZeRO-2 : Optimizer States & Gradients Partitioning（P_os+g）
• ZeRO-3 : Optimizer States & Gradients Partitioning & Parameters Partitioning（P_os+g+p）

（1）ZeRO-1

• 原理

• 只对优化器 Optimizer 进行分片（与 DDP 过程相似）
• 每个 rank（gpu）单独负责 forward 和 backward 过程，在完成 backward 后，梯度通过 AllReduce 来同步。
• 每个 rank 只负责更新当前优化器分片的部分，由于每个 rank 只有部分分片的优化器 state，所以当前 rank 会忽略其余的 state。
• 在更新优化器 state 后，通过广播或者 AllGather 的方式，确保所有的 rank 都收到最新更新过后的模型参数。

• 优点

• 因为 Adam 拥有额外的参数 m（momentum）与 v（variance），特别是 FP16 混合精度训练。
• 适合使用类似 Adam 进行优化的模型训练
• 减少了 4 倍显存，通信容量与数据并行相同

• 缺点

• 因为 SGD 只有较少的参数内存，并且由于需要更新模型参数，导致额外的通讯成本。
• 不适合使用 SGD 类似的优化器进行模型训练
• 只是解决了 Optimizer state 的冗余。

（2）ZeRO-2

• 原理

• 对优化器 Optimizer、gradients 进行分片
• Optimizer 参数被分片，并安排在不同的 rank 上
• 在 backward 过程中，gradients 在不同的 rank 上独自进行 reduce 操作（取代了 all-reduce，以此减少了通讯开销），每个 rank 独自更新各自负责的参数。
• 在更新操作之后，广播或 AllGather，保证所有的 ranks 接受到更新后的参数。

• 优点

• 减少了 8 倍显存，通信容量与数据并行相同

（3）ZeRO-3

• 原理

• 使用 AllGather 获取该层所需要的层之前过程的参数。
• 结束后释放掉不属于该 rank 分片的层的参数。
• 使用 AllGather 获取模型该层所需的前置的层的参数。
• 结束后释放掉不属于该 rank 分片的层的参数。
• AllReduce 操作可以被拆分为 Reduce 与 allgather 操作的结合。
• 模型的每一层拥有该层的完整参数，并且整个层能够直接被一个 GPU 装下。所以计算前向的时候，除了当前 rank 需要的层之外，其余的层的参数可以抛弃。
• 对优化器 Optimizer、gradients、model parameter 进行分片
• 每个 rank 计算 forward 过程
• 每个 rank 计算 backward 过程
• 使用 Reduce 对当前分片的参数的梯度进行累加。
• 让每个 rank 根据聚合的梯度，独立更新参数。

• 优点

• 内存减少与数据并行度和复杂度成线性关系。

3.ZeRO-Offload

Offload 是一种通过将数据和计算从 GPU 卸载到 CPU，以此减少训练期间 GPU 内存占用的方法。该方法提供了更高的训练吞吐量，并避免了移动数据和在 CPU 上执行计算导致的减速问题。
在单张 V100 GPU 的情况下，用 PyTorch 能训练 1.4B 的模型，吞吐量是 30TFLOPS，有了 ZeRO-Offload 加持，可以训练 10B 的模型，并且吞吐量 40TFLOPS。

切分思路

下图中左图是一个使用混合精度的模型训练和参数更新过程。包含了 4 类节点：FWD、BWD、Param update、float2half。M 表示模型的参数量，2M 表示使用 FP16（FP16 =2 Byte），4M 表示 FP32，12M 表示 3xFP32。

右图是使用 Offload 的流程，前后向（FWD/BWD) 这两个计算资源消耗较大的过程放到 GPU 上执行，参数更新和精度转换放在 CPU 上执行，也就是将 Adam 过程放在 CPU 上。

计算思路

CPU

• ① 保存着优化器状态
• ④ CPU 中的每个数据并行线程，进行优化器状态分割的更新（p update）
• ⑤ 将参数分割移回 GPU

GPU

• ② forward 和 backward 过程都在 GPU 上进行
• ③ backward 利用 reduce-scatter 计算求和均值，然后按照数据并行线程，将分割的梯度平均值卸载到 CPU 内存中（g offload）
• ⑥ 执行 all-gather 操作，收集所有更新后的参数（g swap）

在使用带有单张英伟达 V100 GPU 的机器时，可以在不耗尽显存的情况下，运行多达 130 亿个参数的模型，模型规模扩展至现有方法的 10 倍，保持有竞争力的吞吐量。

4.ZeRO-Infinity

模型参数量增长的速度，远快于 GPU 的喜爱内存增长速度，则存在内存墙问题。例如 GPT1 到 GPT3，两年时间参数量从 0.1B 增长到了 175B，而同期，NIVIDIA 则只是从 V100 32GB 更新到 A100 80GB。

ZeRO-infinity 在 ZeRO-Offload 的基础上进一步优化，除了利用 GPU 显存和 CPU 内存外，还利用了 NVMe 磁盘空间。用了这些异构存储器，ZeRO-infinity 突破了 GPU 内存壁垒。

infinity 卸载引擎通过使用 CPU 和 NVMe 内存增加了可用于存储模型参数和激活的内存量；与前几代 ZeRO 不同，infinity 引擎可以将整个模型卸载到这些位置。

以内存为中心的平铺是另一项新技术，它通过将大的模型层分解成较小的 "平铺" 来减少对内存的占用，这些平铺可以按顺序执行；这允许在不需要模型并行的情况下训练大型模型。

为了处理带宽问题，ZeRO-Infinity 引入了以带宽为中心的分区，将模型参数划分到多个数据并行进程中，还有一个重叠引擎，同时执行 NVM 到 CPU、CPU 到 GPU 以及 GPU 到 GPU 的通信。

下图比较了 3D 并行和 ZeRO-Infinity 所能达到的最大模型规模，其支持每个 NVIDIA V100 DGX-2 节点 1 万亿个参数，相比 3D 并行增加了 50 倍。

参考

• Microsoft/DeepSpeed
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• 大语言模型（LLM）分布式训练框架总结
• DeepSpeed 之 ZeRO 系列：将显存优化进行到底
• 数据并行 Deep-dive: 从 DP 到 Fully Sharded Data Parallel （FSDP）完全分片数据并行

书籍推荐：