type
status
date
slug
summary
tags
category
icon
password
🌅
本篇介绍一下在LLMs训练中使用得比较多的一种分布式训练加速框架—deepspeed，简要整理了一下相关细节

什么是多卡并行？

我们在训练模型的时候，如果只使用一张卡进行训练，随着训练数据的增加以及模型结构的复杂化，难免会遇到OOM问题（Out of Memory)。那么这时候你会想到，我那台实验室的八卡机，刚好有四张卡都没人用，如果可以让他们都工作起来，每张卡训练一部分数据，共同优化同一个模型，那是不是就能解决我的OOM问题了？
那么恭喜你，你发明了Data Parallel,DP。具体而言，DP就是利用多张卡同时参与训练，每张卡都独立加载整个模型，并且独立进行前后向过程。通过把训练数据的一个大的 batch 分成多个小 batch，每张卡独立处理一个小 batch，最后再把各个卡上的梯度汇总整合起来（通常是求平均），在一个主卡（主进程） 中计算新的参数值，然后再把新参数同步到各个卡中，这样实现数据的并行训练，所以称之为数据并行（Data Parallel,DP ） 。 pytorch 的 Data Parallel (DP) 和 Distributed Data Parallel (DDP) 都是这种方法的实现， 其中 Data Parallel (DP) 早于 Distributed Data Parallel， Data Parallel (DP) 是 单机多卡 的实现， Distributed Data Parallel 是 多机多卡 的实现。

为什么DP不够用了

本来有了DDP这样的技术后，在处理大规模数据的时候大家都有了合理的解决方案，可以继续愉快的玩耍了，知道有一天，大模型（LLMs）横空出世了，它带来的问题是，模型太大了，导致一张卡根本无法承受一个完整的大模型进行一次训练，那这个时候怎么办呢？
首先我们需要明确，我们常说的大模型，到底是哪里“大”？
  1. 模型的参数。
  1. 前向过程中，一些中间计算结果以及激活值（即激活函数的执行结果）。
  1. 后向过程中，每个参数的梯度值。
  1. 优化器的状态。比如 adam 算法，需要为每个参数再保存一个一阶动量和二阶动量。
这些内容都是会消费显卡的显存的。

参数分割

既然现在遇到的问题是一张卡没办法装下模型的所有参数了，那么有一个很直觉的看法在于，不让一张卡装下所有参数，但在训练和推理的时候仍然可以让模型正常工作。具体怎么实现呢？答案是把参数在不参与计算的时候放到其他地方，例如CPU内存或者高速SSD中（支持NVMe的ssd，走的PCI-E总线），这就是deepspeed中使用到的offload技术。
另外一种想法就是，多张GPU卡，每张卡保存一部分，需要的时候再从其他卡同步过来。以上就是参数分割。
那么要分割，按什么逻辑去分割呢？

1.按层分割

保留每一层（Layer）为整体，不同层存储在不同的 GPU 中， 多个层（GPU）串行在一起，需要串行执行，这就是所谓的 流水线并行（Pipeline Parallel,PP）。时间效率很差， 并且如果某一层的参数量就很大并超过了单卡的显存就尴尬。

2.张量分割

把参数张量切开，切开张量分开存储很容易，但切开之后，张量计算的时候怎么办？这里可以分两种策略。 1. 张量的计算过程也可以切割，这样把一个大的张量，切分成多个小张量，每张 GPU 卡只保存一个小片段，每个小张量片段（GPU卡）独立进行相关计算，最后在需要的时候合并结果就行了。这种思路就称为 张量并行（Tensor Parallel,TP） , Megatron 就是走的这个路线。 2. 同样是把参数张量分割，每张卡只保存一个片段。但是需要计算的时候，每张卡都从其他卡同步其它片段过来，恢复完整的参数张量，再继续数据计算，这种过程被称为AllGather。Deepspeed 选取的这个策略，这个策略实现起来更简单一些。
Megatron-LM 使用 tensor parallel（张量切分）+ pipeline parallel的方式，由于省略了AllGather这一步骤，Megatron-LM的训练速度相比deepspeed其实会更快，但是 TP 并不能像 Zero3 一样把“完整参数”缩到“单层级别”，因此相应地DeepSpeed Zero3会节省更多的显存

混合精度

这一点其实也很好理解，模型运行中，大量的浮点数乘法会产生很多很小的浮点数，如果改成半精度，用2个字节的 float16 替代4个字节 float32，显存需求就降低了一半。如果用量化技术，用2个字节的 int16 或者1个字节的 int8 代替4字节的 float32则同理。虽然精度降低能带来的好处很直观，计算效率和计算代价都会大大降低，但其实精度的降低也存在显著的风险：
首先是鲁莽地对浮点数进行半精度操作，在模型训练过程中很多的极小浮点数就会丢失，造成数据下溢。同时即使不溢出，也损失了数据准确性。 模型训练时，梯度误差大，导致损失不收敛。模型推理时，误差变大，推理效果变差。
那具体什么地方什么时候可以降低精度，用混合精度来达到计算效率提升同时风险可接受的程度呢，参考论文：Mixed Precision Training：
notion image

总结

流水线并行、张量并行，把模型一次完整的计算过程（前后向）分拆到多个 GPU 上进行， 所以这两者都被称为模型并行（Model Parallel,MP）。 而如果每张卡都能进行模型一次完整前后向计算，只是每张卡处理不同的训练数据批次（batch）, 就称为数据并行（Data Parallel,DP）。 deepspeed 对参数进行了分割，每张卡存储一个片段，但在进行运算时， 每张卡都会恢复完整的参数张量，每张卡处理不同的数据批次， 因此 deepspeed 属于数据并行。