常见的分布式并行策略

分布式：是指利用多个计算节点（也称工作节点，Worker）进行机器学习或者深度学习的算法和系统，旨在提高性能、并可扩展至更大规模的训练数据和更大的模型。
并行（parallel）：是指两个或者多个事件在同一时刻发生，并行的反义词是串行；要区分并发（concurrent），它是指两个或多个事件在同一时间间隔发生

看了很多资料，一直也没搞清楚分布式并行到底是个啥，直到读完 OneFlow 的博客《常见的分布式并行策略》才感觉“我好像明白了”，哈哈哈哈或….
下面就来看看被大模型训练带火的「分布式并行策略」的介绍。

为什么分布式训练越来越流行¶

近年来，深度学习被广泛应用到各个领域，包括计算机视觉、语言理解、语音识别、广告推荐等。在这些不同的领域中，一个共同的特点就是模型规模越来越大，比如 GPT-3 模型的参数量达到1750亿。即使用1024张 80 GB 的 A100，那么完整训练 GPT-3 的时长都需要1个月。
模型规模的扩大，对硬件（算力、内存）的发展提出要求。然而，因为 内存墙 的存在，单一设备的算力及容量，受限于物理定律，持续提高芯片的集成越来越困难，难以跟上模型扩大的需求。
为了解决算力增速不足的问题，人们考虑用多节点集群进行分布式训练，以提升算力，分布式训练势在必行。

常见的分布式系统

大语言模型（LLM）分布式训练框架总结

常见的并行策略¶

简单的机器堆叠并不一定会带来算力的增长。因为神经网络的训练并不是单纯的“把原来一个设备做的事情，现在分给多个设备各自做”，它不仅需要多个设备进行计算，还涉及到设备之间的数据传输，只有协调好集群中的计算与通信，才能做高效的分布式训练。
我们将以矩阵乘法的例子，解释数据并行、模型并行的区别。
先了解以下逻辑上的矩阵乘法例子:
假设神经网络中某一层是做矩阵乘法，其中的输入 x 的形状为 4×5，模型参数 w 的形状为 5×8，那么，矩阵乘法输出形状为 4×8。示意图如下：

单机单卡的训练中，以上矩阵乘法，先计算得到 out，并将 out 传递给下一层，并最终计算得到 loss，然后在反向传播过程中，得到$\frac{\partial loss}{\partial w}$，用于更新 w。
分布式训练中，依据是切分 x 还是 w 的不同，分为“数据并行”和“模型并行”策略。接下来，我们介绍常见的并行策略。

数据并行¶

所谓的数据并行，就是将数据 x 进行切分，而每个设备上的模型 w 是完整的、一致的。如下图所示，x 被按照第0维度平均切分到2个设备上，两个设备上都有完整的 w。
这样，在两台设备上，分别得到的输出，都只是逻辑上输出的一半（形状为 2×8），将两个设备上的输出拼接到一起，才能得到逻辑上完整的输出。

注意，因为数据被分发到了2个设备上，因此反向传播过程，各自设备上得到的 $\frac{\partial loss}{\partial w}$ 会不一样，如果直接使用各个设备上的梯度更新各自的模型，会造成2个设备上的 模型不一致，训练就失去了意义（到底用哪个模型好呢？）。
因此，数据并行策略下，在反向传播过程中，需要对各个设备上的梯度进行 AllReduce，以确保各个设备上的模型始终保持一致。
当数据集较大，模型较小时，由于反向过程中为同步梯度产生的通信代价较小，此时选择数据并行一般比较有优势，常见的视觉分类模型，如 ResNet50，比较适合采用数据并行。

模型并行¶

当神经网络非常巨大，数据并行同步梯度的代价就会很大，甚至网络可能巨大到无法存放到单一计算设备中，这时候，可以采用模型并行策略解决问题。
所谓的模型并行，就是每个设备上的数据是完整的、一致的，而模型 w 被切分到了各个设备上，每个设备只拥有模型的一部分，所有计算设备上的模型拼在一起，才是完整的模型。
如下图所示，w 被按照第1维度平均切分到2个设备上，两个设备上都有完整的 x。两个设备上的输出也需要通过拼接才能得到逻辑上的输出。

模型并行的好处是，省去了多个设备之间的梯度 AllReduce；但是，由于每个设备都需要完整的数据输入，因此，数据会在多个设备之间进行广播，产生通信代价。比如，上图中的最终得到的 out (4×8) ，如果它作为下一层网络的输入，那么它就需要被广播发送到两个设备上。
语言模型，如 BERT，常采用模型并行。

流水并行¶

当神经网络过于巨大，无法在一个设备上存放时，除了上述的模型并行的策略外，还可以选择流水并行。 流水并行指将网络切为多个阶段，并分发到不同的计算设备上，各个计算设备之间以“接力”的方式完成训练。
如下图，展示了一个逻辑上的4层网络（T1 至 T4)是如何做流水并行的。
4层网络被切分到2个计算设备上，其中 GPU0 上进行 T1 与 T2 的运算，GPU1 上进行 T3 与 T4 的计算。
GPU0 上完成前两层的计算后，它的输出被当作 GPU1 的输入，继续进行后两层的计算。

混合并行¶

网络的训练中，也可以将多种并行策略混用，以 GPT-3 为例，以下是它训练时的设备并行方案：
它首先被分为 64 个阶段，进行流水并行。每个阶段都运行在 6 台 DGX-A100 主机上。在6台主机之间，进行的是数据并行训练；每台主机有 8 张 GPU 显卡，同一台机器上的8张 GPU 显卡之间是进行模型并行训练。

策略选择篇

单GPU

• 显存够用： 直接用
• 显存不够：上offload，用cpu

单节点多卡

• 显存够用（模型能装进单卡）：DDP或ZeRO
• 显存不够：TP或者ZeRO或者PP
• 重点：没有NVLINK或者NVSwitch，也就是穷人模式，要用PP

多节点多卡

• 如果节点间通信速度快（穷人的万兆网肯定不算）
• ZeRO或者3D并行，其中3D并行通信量少但是对模型改动大。
• 如果节点间通信慢，但显存又少。。。用 DP+PP+TP+ZeRO-1

  常见的显存优化技术

  文章源自@Vadim Irtlach大佬在kaggle的开源notebook

1. 梯度累积（Gradient Accumulation）
2. 冻结（Freezing）
3. 自动混合精度（Automatic Mixed Precision）
4. 8位优化器（8-bit Optimizers）
5. 梯度检查点（Gradient Checkpointing）
6. Parameter-Efficient Fine-Tuning（PEFT）
7. Torch FSDP+CPU offload

这一节的内容整理自大佬总结的节省显存的技巧，除了买GPU外，可以利用这些技巧帮助节省显卡内存，方便训练更大的神经网络。

由于内容比较多，单独写成一个 blog 了，欢迎学习。
《Memory-Efficient 的 LLMs 的训练/微调/推理方法》