DeepSpeed

DeepSpeed 深入解析

混合精度训练

在训练过程中，DeepSpeed 支持同时使用 FP16 和 FP32 两种精度的数据类型。具体而言，在进行前向传播（forward）和反向传播（backward）时，数据类型会转换为 FP16，而在参数更新阶段则转换为 FP32。这种方法可以有效提高训练效率，同时保持较高的计算精度。关于混合精度训练的详细信息，可以参考 2.3.7 章。

ZeRO 零冗余优化器

ZeRO的显存分类

ZeRO 优化器通过显存的精细管理来提高模型训练的效率，其显存分类如下：

模型状态（model states）：包括模型参数（FP16）、模型梯度（FP16）和 Adam 状态（FP32 的模型参数备份，FP32 的 momentum 和 FP32 的 variance）。假设模型参数量为 $Φ$ ，则共需要存储空间为 $2 Φ + 2 Φ + (4 Φ + 4 Φ + 4 Φ) = 4 Φ + 12 Φ = 16 Φ$ 字节。可以看到，Adam 优化器状态占据了显存的 75%，而更新参数时还需加上 FP32 的梯度。
剩余状态（residual states）：包括激活值（activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）。其中，激活值可以通过 activation checkpointing 来大幅减少。

ZeRO 三个阶段以及显存占用分析

数据通信量分析

DeepSpeed ZeRO 主要采用数据并行的方法。

传统数据并行

在传统的数据并行中，每一步（step/iteration）计算梯度后，需要进行一次 AllReduce 操作来计算梯度均值。常用的方法是 Ring AllReduce，分为 ReduceScatter 和 AllGather 两步，每张卡的通信数据量（发送+接受）近似为 $4 Φ$ 。

ZeRO1、2阶段（优化器状态分区、梯度分区）

在 ZeRO1 和 ZeRO2 阶段，与传统数据并行类似。每张卡只存储 $\frac{1}{N}$ 的优化器状态和梯度。对于 gpu_0 来说，为了计算它这 $\frac{1}{N}$ 的梯度均值，需要进行一次 Reduce 操作，通信数据量是 $2 \cdot \frac{1}{N} \cdot Φ \cdot N = 2 Φ$ （发送+接收）。实现中使用了 bucket 策略，保证 $\frac{1}{N}$ 的梯度每张卡只发送一次。当 gpu_0 计算好梯度均值后，就可以更新局部的优化器状态（包括 $\frac{1}{N} \cdot Φ$ 的参数），当反向传播过程结束，进行一次 Gather 操作，更新 $(1 - \frac{1}{N}) \cdot Φ$ 的模型参数，通信数据量是 $2 \cdot \frac{1}{N} \cdot Φ \cdot N = 2 Φ$ （发送+接收）。全局来看，相当于用 Reduce-Scatter 和 AllGather 两步，与数据并行一致。