Megatron系列——张量并行

本文整理自bilibili Zomi视频

1、行切分和列切分

注意：

（1）A按列切分时，X无需切分，split复制广播到A1和A2对应设备即可。最后Y1和Y2需要拼接下，即All Gather

（2）A按行切分时，X需要按列切分split，最终Y1和Y2需要做All Reduce Sum操作

（3）多次连续矩阵乘，可以把All Gather和Split省略，这个操作是不必要的。

1.1 反向传播

备注：

（1）这个视频中有人提到为什么需要对X进行求导，可以看做X为上一层或上一个模块的输出，要对上一层参数进行求导，就必须要对X求导。

（2）反向传播时，注意是两个设备上X的导数求和

行并行

2、MLP

备注：

（1）经过了两次矩阵乘，A按列切分，算出的Y1和Y2也为按列切分了，故B按行切分即可。

（2）f为copy，g为all reduce

3、Attention

备注：

（1）这里的Q，K，V是按头的个数进行列切分