张量并行

问题分析

随着模型越来越庞大，其尺寸远远超出了处理器内存的限制，并且模型训练时间也变得很长。所以需要把一个模型切分，每个计算设备只负责一部分模型的存储与计算。

解决方案

张量并行将模型分成多份并存储在多个计算设备上，这样模型的权重和优化器状态可以分布在多个计算设备上，以此来克服单个计算设备显存无法容纳整个大型模型的问题。并且因为各计算设备只需要处理一部分的模型计算，训练速度也得到显著提高。这种分片策略叫做张量并行。

解决思路

参数矩阵横切

1.参数矩阵横切策略按照参数矩阵的行来切分模型，该切分策略需要将输入矩阵也进行按列切分。 2.横切策略前向时，先切分输入矩阵，对应部分的输入矩阵进入对应部分的模型进行前向计算，之后用all-reduce操作来将各部分模型计算结果累加得到最终前向计算结果。 3.横切策略反向时，可以计算得出最终输出的梯度和各部分模型的输出梯度相等，先将最终输出的梯度传到各部分模型的输出张量，再用all-gather操作将切分后的输入矩阵的梯度拼接得到最初输入矩阵的梯度。

参数矩阵纵切

1.参数矩阵纵切策略按照参数矩阵的列来切分模型，该切分策略输入矩阵无需进行切分。 2.纵切策略前向时，先将输入矩阵送入各部分模型，各部分模型分别进行前向计算得到输出结果，之后用all-gather操作来将各部分模型输出结果拼接得到最终前向计算结果。 3.纵切策略反向时，先将最终输出的梯度进行切分并将对应的部分传到对应部分模型的输出张量，之后用all-reduce操作将各部分模型的输入矩阵的梯度累加得到最初输入矩阵的梯度。

使用场景

如果用户发现训练显存占用过高或超出了处理器显存限制，或者训练时间过长，可以开启张量并行来降低单设备显存占用，加快训练速度。

使用方法

在训练脚本的参数列表中加入 --tensor-model-parallel-size N，设置张量并行的size。

使用效果

利用多个设备，降低显存占用，加快训练速度。