LLM：并行化训练

llama3.1论文称使用了：张量并行（TP）、流水线并行（PP）、上下文并行（CP）和数据并行（DP）。

并行化可以先粗分为：
1.数据并行
2.模型并行：张量并行，流水线并行
3.上下文并行（llama3.1）
4.deepspeed （专门提出）

所以主要就是 数据并行（DP）、张量并行（TP）、流水线并行（PP）。还有deepspeed。

1.数据并行

数据并行是最常见的并行形式，因为它很简单。在数据并行训练中，数据集被分割成几个碎片，每个碎片被分配到一个设备上。这相当于沿批次（Batch）维度对训练过程进行并行化。每个设备将持有一个完整的模型副本，并在分配的数据集碎片上进行训练。在反向传播之后，模型的梯度将会聚合（All Reduce），以便在不同设备上的模型参数能够保持同步。

DP（Data Parallelism）：早期数据并行模式，一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡。
DDP（Distributed Data Parallelism）：分布式数据并行，采用Ring AllReduce的通讯方式，多用于多机多卡场景。
可见：https://www.bilibili.com/video/BV1mm42137X8/?share_source=copy_web&vd_source=29af710704ae24d166ca951b4c167d53

2.张量并行

张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。
以一般的矩阵乘法为例，假设我们有 C = AB。我们可以将B沿着列分割成 [B0 B1 B2 … Bn]，每个设备持有一列。然后我们将 A 与每个设备上 B 中的每一列相乘，我们将得到 [AB0 AB1 AB2 … ABn] 。此刻，每个设备仍然持有一部分的结果，例如，设备(rank=0)持有 AB0。为了确保结果的正确性，我们需要收集全部的结果，并沿列维串联张量。通过这种方式，我们能够将张量分布在设备上，同时确保计算流程保持正确。
也可以对这两个矩阵都进行切分。

3.流水线并行

流水线并行的核心思想是，模型按层分割成若干块，每块都交给一个设备。
在前向传播过程中，每个设备将中间的激活传递给下一个阶段。
在后向传播过程中，每个设备将输入张量的梯度传回给前一个流水线阶段。

该方案有个很大的缺陷是GPU的利用率太低，当一个GPU进行计算的时候，其他层的GPU都是闲置的。为了解决这个缺陷，谷歌提出了GPipe 流水线并行（Pipeline model parallesim ），下图来自GPipe官方论文。

方案引用了一个chunks的超参,将每一个mini-batch数据划分到chunks个micro-batches上。

4.deepspeed

官方文档：https://www.deepspeed.ai/

模型训练时的大概流程

正向传播时（FWD），上一层fp16的激活值和fp16的参数参与了计算，得到这一层的fp16激活值
反向传播时（BWD），本层的激活值和参数参与计算，得到fp16的梯度
参数更新时，fp16的梯度以及fp32的参数副本，momentum和variance参与计算，最终算出更新后的fp32参数、momentum和variance，然后将fp32的参数转化为fp16进行本层的参数更新

训练过程中的显存占用分析:

训练大模型时通常会采用AdamW优化器，并用混合精度训练来加速训练，基于这个前提分析显存占用。

来自：https://zhuanlan.zhihu.com/p/624740065

来自：https://zhuanlan.zhihu.com/p/639872915

按我的理解：fp16的模型参数，fp16的梯度，fp32的优化器状态：模型参数、动量、方差。最终就是2+2+4*(1+1+1)=16。

推理过程中的显存占用分析

deepspeed详解

可以看到baseline的参数估计就是我理解的这种，条件是提到的那些，k是优化器状态，2是模型参数，2是梯度。

ZeRO技术优化重点便是优化Model States 显存占用。在整个训练过程中， States只是会在对应算子计算时候才会被用到，optimizer states只在最终做update时才用到。数据并行中，gradients只在最后做AllReduce和updates时才用到。参数W只在做forward和backward的那一刻才用到，ZeRO技术针对以上做了以下优化：

Stage 0: 禁用所有类型的分片，仅使用 DDP (Distributed Data Parallel) ，对应下图的Baseline情况；

Stage 1: 把优化器状态(optimizer states) 分片到每个数据并行的工作进程(每个GPU)下，对应下图Pos

Stage 2: 把优化器状态(optimizer states) + 梯度(gradients) 分片到每个数据并行的工作进程(每个GPU)下，对应下图Pos + Pg的情况；

Stage 3: 把优化器状态(optimizer states) + 梯度(gradients) + 模型参数(parameters) 分片到每个数据并行的工作进程(每个GPU)下，Pos + Pg +Pp的情况；

Optimizer Offload: 在Stage2的基础上，把梯度和优化器状态下沉到CPU内存或硬盘上

Param Offload: 在Stage3的基础上，把模型参数下沉到CPU内存或硬盘上

deepspeed的实现需要设置一个config文件，在训练时一般使用zero2。

参考：
https://zhuanlan.zhihu.com/p/598714869

https://zhuanlan.zhihu.com/p/688873027
https://zhuanlan.zhihu.com/p/685472786
https://zhuanlan.zhihu.com/p/624740065
【动画理解Pytorch 大模型分布式训练技术 DP，DDP，DeepSpeed ZeRO技术】
https://zhuanlan.zhihu.com/p/639872915