分布式训练-模型并行（张量并行）-- 原理

最新推荐文章于 2024-06-09 15:03:53 发布

她的我

最新推荐文章于 2024-06-09 15:03:53 发布

阅读量1.2k

点赞数 21

文章标签：算法 llama 分布式 pytorch

本文链接：https://blog.csdn.net/weixin_42503655/article/details/135597408

版权

本文详细探讨了张量并行在Transformer模型中的应用，特别是MLP层和Attention层的并行策略，以及如何通过列切割和行切割进行GPU拆分，同时介绍了AllReduce在模型训练中的作用。

摘要由CSDN通过智能技术生成

0 引言

本文介绍张量并行原理，只作为学习总结记录，错误请指出。

1 原理

流水线并行和张量并行都可以看作是模型并行的一种，只是对模型切分的维度不同，流水线并行可以看作是层间并行，将模型不同的层放到不同的GPU上，张量并行看作是层内并行，是对层内具体的矩阵运算进行拆分。
在这里插入图片描述

1.1 拆分方式

x为输入，A为权重参数，y为输出，按列拆分是保持x不变，将A按列拆分，分别计算后再拼接在一起。按行拆分是将x按列拆分，A按行拆分，分别计算后，再相加得到最终的输出。
在这里插入图片描述
在两张GPU上的具体的并行操作如图所示，思想就是将进行运算的矩阵拆分到不同的GPU上，然后将每个GPU上的运算结果进行汇总（GPU之间通信）。

1.2 Transformer 各模块并行

1.2.1 MLP层

MLP层比较简单，就是将输入 X 先升维后降维，增加模型的表达能力，GELU是激活函数，A和B分别为两个线性层。
在这里插入图片描述

在MLP层中，先对A采用“列切割”，然后对B采用“行切割” 。f 和 g 分别表示两个算子，每个算子都包含一组forward + backward 操作

f 的 forward 计算：把输入X拷贝到两块GPU上，每块GPU即可独立做forward计算。
g 的 forward 计算：每块GPU上的forward的计算完毕，取得Z1和Z2后，需要相加，因此在 GPU间做一次AllReduce，进行通信。
g 的 backward 计算：需要损失函数对 Z 的梯度拷贝到两块GPU上，两块GPU就能各自独立做梯度计算。
f 的 backward 计算：中间层的梯度计算完毕，需要传递到 X 做梯度计算, 则此时两块GPU做一次AllReduce，把各自支路对输入矩阵X 的梯度和相加即可。

MLP层做forward时产生一次AllReduce，做backward时产生一次AllReduce。一共进行了两次AllReduce。

1.2.2 Attention 层

Attention 层是多头注意力机制，在具体实现多头注意力机制时，主要是先通过三个线性映射，将输入 x 的维度映射到 d，得到Q、K、V，然后对 d 进行拆分，每个头的维度变为 d / num_head。
因此，Attention 的多头计算是为张量模型并行量身定做的，因为每个头上都可以独立计算，最后再将结果concat起来。也就是说，可以把每个头的参数放到一块GPU上。
在这里插入图片描述
对三个参数矩阵Q，K，V，按照“列切割”，每个头放到一块GPU上，做并行计算。对线性层B，按照“行切割”。切割的方式和MLP层基本一致，其forward与backward原理也一致，在实际应用中，并不一定按照一个head占用一块GPU来切割权重，我们也可以一个多个head占用一块GPU，这依然不会改变单块GPU上独立计算的目的。

self-attention层在forward中做一次AllReduce，在backward中做一次AllReduce。一共进行了两次AllReduce。

1.2.3 映射层

映射层是将文字序列转化为模型能够处理的数字，后面会学习总结一篇关于大模型输入输出的文章，对于输入X，word embedding的过程，就是等于用token的序号去word embedding中查找对应词向量的过程。例如，输入数据为[0, 212, 7, 9]，数据中的每一个元素代表词序号，我们要做的就是去word embedding中的0，212，7，9行去把相应的词向量找出来。
在这里插入图片描述

X为文字序列：[君不见黄河之水天上来，奔流到海不复回。]

经过序列化：['BOS','君','不','见','黄','河','之','水','天','上','来','，'奔','流','到', '海', '不', '复', '回', '。', 'EOS']

经过索引化：['BOS','10','3','67','89','21','45','55','61','4','324','565' ,'789','6567','786' ...]

根据词表进行向量化：
'BOS'-> [p_{00},p_{01},p_{02},...,p_{0d-1}]
'10' -> [p_{10},p_{11},p_{12},...,p_{1d-1}]
'3'  -> [p_{20},p_{21},p_{22},...,p_{2d-1}]
...
'EOS'-> [p_{n0},p_{n1},p_{n2},...,p_{nd-1}]

输出层中，同样有一个word embedding，把输入再映射回词表里，得到每一个位置的词。一般来说，输入层和输出层共用一个word embeding。
在backward的过程中，我们在输出层时会对word embedding计算一次梯度，在输入层中还会对word embedding计算一次梯度。在用梯度做word embedding权重更新时，我们必须保证用两次梯度的总和进行更新。
若模型输入层和输出层在不同的GPU上时，我们就要保证在权重更新前，两块GPU上的word embedding梯度做了一次AllReduce。
在这里插入图片描述