概述

基本思想是将模型的参数分开，放到不同的卡上进行独立计算，然后再做聚合。切分时候存在两种方式：按行切分与按列切分。

AscendSpeed提供两种张量并行的切分方式，分别是按行切分和按列切分。

按行切分：为了满足矩阵乘法规则，输入矩阵X需要按列切分X=[X1 | X2]。同时，将矩阵分块，分别放置在两张卡上，每张卡分别计算Y1=X1*A1，Y2=X2*A2。计算完成后，通过collective通信Allreduce_sum，归约其他处理器上的计算结果，可以得到最终的结果矩阵Y。
图1 按行切分
按列切分：分别将A1，A2放置在两张卡上。两张卡分别计算Y1=X*A1和Y2=X*A2。计算完成后，通过collective通信AllGather，获取其它卡上的计算结果，拼接在一起得到最终的结果矩阵Y。
图2 按列切分

通过将单卡显存无法容纳的矩阵A拆分，放置在两个处理器上，并通过多卡间通信，即可得到的最终结果。该结果在数学上与单卡计算结果上是等价的。

功能使用

通过修改输入参数--tensor-model-parallel-size的值来使能，输入值为int型，默认值为1。

训练脚本参数配置示例，加粗部分为张量并行功能配置参数示例：

python -m torch.distributed.launch $DISTRIBUTED_ARGS \
    pretrain_llama.py \
    --DDP-impl local \
    --use-contiguous-buffers-in-ddp \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 2 \
    --num-layers 30 \
    --hidden-size 6656 \
    --ffn-hidden-size 17920 \
    --num-attention-heads 52 \
    --micro-batch-size 1 \
    --global-batch-size 32 \
    --seq-length 2048 \

张量并行功能的两个切分操作定义位于“ascendspeed/mpu/layers.py”文件下的类class RowParallelLinear()与类class ColumnParallelLinear()。

行切分的操作中需要关注input是否是进行parallelize，列切分的实现中，需要关注是否对output做gather操作。

点击放大

f与g是两个共轭算子，f代表前向，g代表反向。

在列切割中：
f : forward中，直接copy输入；backward中，对梯度做AllReduce。在代码里定义为class _CopyToModelParallelRegion（同文件夹下的mappings.py）。

g: forward中，all-gather输出；backward中，对梯度做split（每张卡经过all-gather已有完整的Y了，因此以Y为起点计算梯度后，沿着列做split就可得到Y1和Y2的梯度）。在代码里定义为class _GatherFromModelParallelRegion（同文件夹下的mappings.py）。
在行切割中：
f: forward中，按列split输入；backward中，all-gather梯度。在代码里定义为class _ScatterToModelParallelRegion（同文件夹下的mappings.py）。

g: forward中，AllReduce输出；backward中，直接输出梯度，无需做任何通讯（因为经过g的foward，每块GPU上已拥有了Yi和Y，则根据图中g的backward公式可知，每块卡独立计算梯度）。在代码里定义为class _ReduceFromModelParallelRegion（同文件夹下的mappings.py）。

在“parallel_state.py”中有TP的设置与获取。

def set_tensor_model_parallel_world_size(world_size):
    """Set the tensor model parallel size"""
    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
        return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())

在实现中，TP的值的获取由以下过程来进行定义：

# Tensor model parallel size.
args.tensor_model_parallel_size = min(
    args.tensor_model_parallel_size, args.world_size)

TensorParallelism张量并行

概述

功能使用