大语言模型--张量并行原理及实现
原创
背景知识
NCCL
NCCL是一个Nvidia专门为多GPU之间提供通讯的通讯库，或者说是一个多GPU卡通讯的框架 ，提供了包括AllReduce、Broadcast、Reduce、AllGather、ReduceScatter等集合通讯API。NCCL屏蔽了底层复杂的细节，向上提供API供训练框架调用，向下连接机内机间的GPU以完成模型参数的高效传输。
Megatron-LM
NVIDIA Megatron-LM 是一个基于 PyTorch 的分布式训练框架，用来训练基于Transformer的大型语言模型。Megatron-LM 综合应用了数据并行（Data Parallelism），张量并行（Tensor Parallelism）和流水线并行（Pipeline Parallelism）。很多大模型的训练过程都采用它，例如bloom、opt、智源等。
torch.distributed（dist）
为运行在一台或多台机器上的多个计算节点之间的PyTorch 提供支持多进程并行性通信的原语。他能轻松地并行化在跨进程和机器集群的计算。
Group是我们所有进程的子集。
Backend进程通信库。PyTorch 支持 NCCL，GLOO，MPI。
world_size在进程组中的进程数。
Rank分配给分布式进程组中每个进程的唯一标识符。它们始终是从 0 到 world_size 的连续整数。
torch.distributed算子介绍
gather
把其它进程的数据收集到目标进程，返回一个列表
all_gather
是将所有进程的数据收集起来，再分发给它们
reduce
把所有进程的数据加起来，发送给目标进程
all_reduce
把所有节点的值加起来，再分发给所有节点。
broadcast
把某个节点的数据分发给所有节点。
scatter
把某个进程上的列表数据逐个分发给其它所有进程
LLM中支持并行的算子介绍
Embedding层
Embedding层包含两个输入，一个是word embedding(v, h)，存放的是所有词的向量，v表示词表大小。词表往往很大，并行主要考虑对word embedding的拆分。
另一个是position embedding，主要用于从word embedding中索引出对应的embedding，例如，输入数据为[0, 212, 7, 9]，数据中的每一个元素代表词序号，我们要做的就是去word embedding中的0，212，7，9行去把相应的词向量找出来。一般不太大，不用考虑并行。
word embedding切分方式：
按列切分
把word embedding按列拆分，每张卡都有完整的position embedding，根据position embedding值索引到对应位置的word embedding得到Y1和Y2，将Y1和Y2 all_gather起来就是完整的Y。
Llama2 Attention层
Attention的核心公式是：
通过下图可以看出，llama2模型中的attention层再核心公式前后各加了一层Linear层。
对于上面这个attention层的并行策略是：
1、第一个Linear按列拆分，X输入分别和W1、W2计算，得到两个输出
2、第一步的两个输出，各自进行attention核心公式计算过程，计算完也会有两个输出
3、第二个Linear按行拆分，第二步的两个输出和W1、W2计算，得到两个输出
4、两个输出进行all_reduce相加，得到最后的output。attention层的并行结束
过程如下图：
Llama2 FeedForward层
Llama2中的FeedForward计算过程如下：
down(up(X)*SiLU(gate(X)))
up、down 与 gate 是三个维度相同的 Linear 层，图示计算过程：
并行策略如下：
1、up层Linear按列拆分，X输入与之计算之后，会再每张卡上有一个输出。
2、gate层Linear按列拆分，X输入与之计算之后，同样的每张卡有一个输出，
3、每张卡的输出各自进行SiLU和矩阵乘计算
4、down层Linear按行拆分，分别与每张卡的输出计算，产生两个输出
5、两个输出进行all_reduce相加，得到最终的输出
过程如下图：
单独的Linear层
Linear主要做矩阵乘法，把一个s*h的输入和h*h`的weight做矩阵乘，得到一个s*h`的输出。过程如下图：
按列拆分
Y1和Y2使用all_gather算子汇总结果得到最终的Y
代码实现
模型加载
每个计算设备分别加载一部分权重，model的load_state_dict函数需要根据设备数，让每个设备加载对应的权重。除此之外，layer和算子的定义也要根据设备数，支持只计算对应的部分。
Linear汇总结果
如上文，Attention层最后一个Linear、MLP层最后一个Linear都需要汇总结果，需要使用all_reduce算子。
ppl.pmx/torch_function/RowParallelLinear.py at master · openppl-public/ppl.pmx (github.com)
单独的Linear需要使用all_gather汇总结果
ppl.pmx/torch_function/ColumnParallelLinear.py at master · openppl-public/ppl.pmx (github.com)
参考文献：
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。