计算通信并行优化算法

基本原理

在大模型中存在许多计算与通信相邻的片段，尤其是Linear模块。计算通信并行优化算法旨在降低将此类相邻的且有结果依赖的计算通信片段的执行用时。该算法当前支持三种通信接口，All-Reduce，All-Gather和Reduce-Scatter，各支持先计算后通信或者先通信后计算两种应用场景。根据模型需求，该算法对ColumnParallelLinear以及RowParallelLinear两个类的forward函数进行替换，从而将原模型中对应的模块用计算通信并行接口替代。

具体地，该算法通过对张量并行（TP）切分后的张量在其所在的卡上沿着m轴进一步切分，使其得以在对结果有依赖的、相邻的计算和通信过程中形成并行流水线。一般地，矩阵乘计算中的m轴由模型的sequence length和batch size决定，即该切分只作用于不同数据间、同一数据的不同token间，对hidden size等单一token处理流程中涉及的维度没有影响，因此不会对模型本身的参数进行任何切分。

更多地，该算法会根据用户指定的切分份数来对左矩阵m轴进行相应数量的切分，然后交错执行计算和通信从而达到掩盖通信时间的效果。需要注意的是，计算通信并行的切分份数与TP并行数量是独立的两个超参，无制约关系，只需要保证TP切分后的张量的m轴大小可以被此算法的切分数量整除即可。该算法的并行与掩盖原理如同下图所示。

使用场景

对于有ParallelAttention与ParallelMLP模块且串行执行它们的大模型，可以按需选择此算法来降低这两个模块前后集合通信的执行时间。需注意如果您使用的大模型对ColumnParallelLinear或RowParallelLinear两个类进行过自定义操作，需要酌情选择修改计算通信并行模块中对应的代码或不使用此优化算法。

操作步骤

1. （必选）修改环境变量，开启或关闭计算通信并行。

   在运行时可以通过环境变量来控制是否开启计算通信并行（CC）。CC_PARALLEL_NUM 代表在进行计算通信并行的时候对张量进行切分的份数，目前支持在1, 2, 4, 8 四个数字中进行选择。其中 1代表不进行计算通信并行。

   export CC_PARALLEL_NUM=1/2/4/8

2. （必选）修改环境变量，开启辅助功能。

   当前PyTorch框架默认的多流内存复用机制，会导致被通信流依赖的延迟释放，内存复用不及时。在torch_npu中提供了一个多流内存复用的增强功能，提早释放被通信流占用的内存，使得内存能有效利用。开启内存复用，需设置如下环境变量。

   export MULTI_STREAM_MEMORY_REUSE=1

3. （可选）根据Matmul Shape自定义CC切分份数。

   有时候会出现一部分 Matmul 进行计算通信并行有收益，而一部分没有收益的情况。这时候我们可以手动根据 shape 来关闭那部分 Matmul的计算通信并行。当然，如果希望让一部分 Matmul 使用非 CC_PARALLEL_NUM 设置的值作为其切分份数，也可以通过这个方法配置。该配置同样通过 ./user_config.py 中的 cc_cfgs 字典中对应的字段进行设置。如果不配置，则使用如下方式进行设置：

   'customized_cc': { }

   如果需要给 profiling（原始不开CC的profiling） 中某个 [m, k] * [k, n] 的 Matmul 算子设置自定义的 CC 切分份数，则可以通过如下方式进行设置：

   'customized_cc':{
       "[16384, 5120, 1920]":8,
       "[16384, 1920, 5120]":1
   }

   上面的设置代表对于 [16384, 5120] * [5120, 1920] 这个 Matmul 在做计算通信并行的时候切 8份，而对 [16384, 1920] * [1920, 5120] 这个 Matmul 则不做计算通信并行。