计算通信并行优化
背景与挑战
在大模型训练中，ColumnParallelLinear和RowParallelLinear组件的前向与反向传播阶段常包含紧密耦合的计算与通信序列，其中计算环节主要涉及矩阵乘法（Matmul），而通信则依赖于AllReduce（在未启用序列并行时）或AllGather与ReduceScatter（在序列并行模式下）。由于这些操作间的顺序依赖性，它们往往被串行执行，导致计算与通信流程中存在空闲等待时间，进而影响整体执行效率。
解决方案
为克服上述局限，我们提出了一套计算与通信并行优化方案（Communication Over Computation，简称CoC），旨在通过将计算与通信任务细化为更小粒度的子任务，实现两者的流水线式并发执行，从而最大化资源利用率与执行效率。
Python脚本层面实现
通过进一步分割张量（分为2/4/8个子张量），借助Python脚本机制，实现各子张量间计算与通信的并行处理，尤其适用于Attention模块与MLP模块的串行执行场景。
- 要求：对于Python脚本侧实现，Matmul左矩阵的m轴尺寸需为2、4或8的倍数。这是因为切分后的矩阵块需要能够在多个处理器或设备之间均匀分布，以确保负载平衡。
- 限制：该方法不适用于计算与通信片段耗时相差较大的情况。如果两者的时间消耗差异显著，可能会导致某些进程等待其他较慢的操作完成，从而抵消了并行化的收益。
- 注意事项：当切分数较大时，可能会遇到host bound问题（即CPU成为瓶颈），这会阻碍预期性能提升的实现。因此，在设计分割策略时，必须考虑到这一点，避免过度细分任务。
融合算子实现
- MATMUL_ALL_REDUCE融合算子（先计算后通信）及其确定性计算。
量化场景下，MATMUL_ALL_REDUCE融合算子支持fp16格式下的w8A16伪量化，粒度包括per tensor、per channel与per group。
- MATMUL_REDUCE_SCATTER融合算子（先计算后通信）及其确定性计算。
- ALL_GATHER_MATMUL及ALL_GATHER_MATMUL_V2融合算子（先通信后计算，V2版本支持获取ALL_GATHER中间结果）。
使用场景
- 适用于特定模块的串行执行模式
该特性主要用于训练场景，特别适合注意力Attention模块和MLP（Multi-Layer Perceptron，多层感知机）模块以串行方式执行的情况。当计算与通信存在顺序依赖，并且这些操作在模型架构中位置毗邻时，可以利用此特性进行优化。
- 支持多种通信场景
- ALL_REDUCE：用于将所有节点上的梯度相加后广播给每个节点。
- ALL_GATHER：收集来自所有节点的数据并在每个节点上构建完整数据集副本。
- REDUCE_SCATTER：对所有节点上的数据进行部分汇总，并将结果分配给各节点。
此外，还提供是先通信还是先计算的灵活设置，这取决于具体的应用需求和模型结构。
使用方法
计算通信并行优化算法通过在训练脚本中配置环境变量来进行使能。
当前计算通信并行有两种实现方法：Python脚本使能和融合算子使能，两者选其一即可。两种方式都需要替换原Megatron框架中的ColumnParallelLinear和RowParallelLinear这两个class的forward函数，替换脚本已经根据MindSpeed指定Megatron版本进行编码和适配，位于“mindspeed/core/tensor_parallel/lcal_coc/”目录下。
计算与通信并行优化功能在MindSpeed框架内的启用，主要通过训练脚本中的环境变量配置实现，需预先安装MindSpeed。
当前，该优化提供两种实现途径即Python脚本使能与融合算子使能。根据具体需求，用户可选择以下任意一种配置方式：
- Python脚本使能
--use-ascend-coc --coc-parallel-num N # 指定张量分割的数量，N可选 2 /4 /8
- 融合算子使能融合算子使能要求安装NNAL（Ascend Transformer Boost加速库），请参考软件安装完成安装。
--use-ascend-coc --coc-fused-kernel # 当前仅支持TP（Tensor Parallelism）配置为8的特定场景
- 融合算子的环境变量拥有更高优先级，即当“ coc-parallel-num”配置值大于1且使能“coc-fused-kernel”时，前者“ coc-parallel-num”不会生效。
- 暂不兼容--use-ascend-mc2特性 。
- 暂未适配MoE模型。
使用效果
根据实际测试数据，采用计算通信并行优化能够带来显著的性能改进，例如对于LLaMA2-70B模型，性能提升约5%以上。