大模型通信笔记2

数据并行

传统的数据并行是让每个GPU分别进行FWD和BWD，然后把梯度进行聚合操作，然后再下发给每个GPU，称为All Reduce。

缺点

  • 存储开销大。每块GPU上都存了一份完整的模型，造成冗余
  • 通讯开销大。Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时，Server的带宽将会成为整个系统的计算效率瓶颈。

异步梯度更新

  • Worker不等待梯度更新，用旧的参数进行下一轮训练，可能会延迟一步更新梯度，整体收敛速度变慢，但是提升通讯计算比。
  • 延迟步数会指定

分布式数据并行

核心目标是降低通信压力，因此要将Server的通信压力转到Worker上，最简单的就是Ring-AllReduce。

Ring Allreduce

核心思路是实现Reduce Scatter和All-Gather。GPU每次之和前后两个GPU通信，1卡给2卡发a号数据，2给3发b号，以此类推。三次更新后每张卡都有1个号的完整的数据。

之后在进行All-Gather，依旧环形通信，把每个部分全聚合的都发给下一个，然后依此类推，3轮通信就可以覆盖所有。

显存开销

数据并行中，每个卡都存储了所有参数，怎么办？

在实际存储中，分为两部分存储：

  • 模型状态：包括参数，优化器，梯度等
  • 驻留数据：包括activation，碎片内存和缓冲区等。

优化措施

混合精度训练

对于参数，activation，梯度，都使用fp16，对于参数（多存一份）和优化器使用fp32。
模型必存数据为$K\phi$,那么最终总存储数据为$K\phi + 4\phi$
实际上，activation大小和batch有关，而且是可以抛弃的。

ZeRO-DP

第一步：优化器分割

每张卡只存储一部分优化器参数，在数据并行中，先通过AllReduce得到完整梯度，每个卡都更新自己的一部分梯度和参数，然后再AllGather。产生单卡通讯量$\phi$。

第二步：梯度分割

经过FWD和BWD后，对梯度进行Reduce-Scatter，保证每张卡都有自己一份聚合梯度，用分割的优化器和梯度进行更新相应的W，然后再AllGather参数进行更新

第三步：参数分割

FWD时，先All Gather一次参数，用完即弃。
BWD时，再All Gather一次参数，用完即弃
用自己的梯度进行一次All Gather得到完整梯度
更新参数，无需通信。

ZeRO-R

通过对驻留数据进行优化来实现显存使用减少和通信负载降低。

activation

每块GPU上只维护部分的activation，需要时再聚合。或者重新计算。

Buffer

通过使用固定大小的Buffer，降低通信次数，减少碎片信息发送，提高带宽利用率

碎片内存整合

ZeRO-Offload

见论文，主要是把显存的优化器参数卸载到CPU内存。