大模型训练~数据并行

大模型场景里巨大的存储和GPU间通讯量是系统设计时需要考虑的重点，本文递进介绍了三种主流数据并行的实现方法：DP、DD皮、ZeRo。

当模型太大，一块GPU放不下时，流水线并行将模型的不同层放到不同的GPU上，通过切割mini-batch实现对训练数据的流水线处理，提升GPU计算通讯比。同时通过re-materialization机制降低显存消耗。

但在实际应用中，流水线并行并不特别流行，主要原因是模型能否均匀切割，影响了整体计算效率，这就需要算法工程师做手调。因此，今天我们来介绍一种应用最广泛，最易于理解的并行范式：数据并行。

数据并行的核心思想是：在各个GPU上都拷贝一份完整模型，各自吃一份数据，算一份梯度，最后对梯度进行累加来更新整体模型。理念不复杂，但到了大模型场景，巨大的存储和GPU间的通讯量，就是系统设计要考虑的重点了。在本文中，我们将递进介绍三种主流数据并行的实现方式：

• DP（Data Parallelism）：最早的数据并行模式，一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡

• DDP（Distributed Data Parallelism）：分布式数据并行，采用Ring AllReduce的通讯方式，实际中多用于多机场景

• ZeRO：零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式，旨在降低存储。

数据并行（DP）

整体架构

一个经典数据并行的过程如下：

• 若干块计算GPU，如图中GPU0~GPU2；1块梯度收集GPU，如图中AllReduce操作所在GPU。

• 在每块计算GPU上都拷贝一份完整的模型参数。

• 把一份数据X（例如一个batch）均匀分给不同的计算GPU。

• 每块计算GPU做一轮FWD和BWD后，算得一份梯度G。

• 每块计算GPU将自己的梯度push给梯度收集GPU，做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义。

• 梯度收集GPU聚合完毕后，计算GPU从它那pull下完整的梯度结果，用于更新模型参数W。更新完毕后，计算GPU上的模型参数依然保持一致。

• 聚合再下发梯度的操作，称为AllReduce。

前文说过，实现DP的一种经典编程框架叫“参数服务器”，在这个框架里，计算GPU称为Worker，梯度聚合GPU称为Server。在实际应用中，为了尽量减少通讯量，一般可选择一个Worker同时作为Server。比如可把梯度全发到GPU0上做聚合。需要再额外说明几点：

• 1个Worker或者Server下可以不止1块GPU。

• Server可以只做梯度聚合，也可以梯度聚合+全量参数更新一起做

在参数服务器的语言体系下，DP的过程又可以被描述下图：

通讯瓶颈与梯度异步更新

DP的框架理解起来不难，但实战中确有两个主要问题：

• 存储开销大。每块GPU上都存了一份完整的模型，造成冗余。关于这一点的优化，我们将在后文ZeRO部分做讲解。

• 通讯开销大。Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时，Server的带宽将会成为整个系统的计算效率瓶颈。

我们对通讯开销再做详细说明。如果将传输比作一条马路，带宽就是马路的宽度，它决定每次并排行驶的数据量。例如带宽是100G/s，但每秒却推给Server 1000G的数据，消化肯定需要时间。那么当Server在搬运数据，计算梯度的时候，Worker们在干嘛呢？当然是在：

人类老板不愿意了：“打工系统里不允许有串行存在的任务！”，于是梯度异步更新这一管理层略诞生了。

上图刻画了在梯度异步更新的场景下，某个Worker的计算