计算通信并行优化
背景与挑战
在大规模模型训练中，ColumnParallelLinear和RowParallelLinear组件的前向与反向传播阶段常包含紧密耦合的计算与通信序列，其中计算环节主要涉及矩阵乘法（Matmul），而通信则依赖于AllReduce（在未启用序列并行时）或AllGather与ReduceScatter（在序列并行模式下）。由于这些操作间的顺序依赖性，它们往往被串行执行，导致计算与通信流程中存在空闲等待时间，进而影响整体执行效率。
解决方案
为克服上述局限，我们提出了一套计算与通信并行优化方案（Communication Over Computation，简称CoC），旨在通过将计算与通信任务细化为更小粒度的子任务，实现两者的流水线式并发执行，从而最大化资源利用率与执行效率。
Python脚本层面实现
通过进一步分割张量（分为2/4/8个子张量），借助Python脚本机制，实现各子张量间计算与通信的并行处理，尤其适用于Attention模块与MLP模块的串行执行场景。需注意，此方法对Matmul左侧矩阵的m轴尺寸有特定要求，须为分割数的整数倍；且在计算与通信片段耗时差异较大或矩阵分割数量过多时，可能遭遇主机瓶颈，影响预期性能提升。
融合算子实现
- 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，多层感知器）模块的串行执行，特别是在计算通信存在顺序依赖与位置毗邻的条件下。
使用方法
计算通信并行优化算法通过在训练脚本中配置环境变量来进行使能。
当前计算通信并行有两种实现方法：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
- 融合算子使能融合算子使能要求安装ATB（Ascend Transformer Boost），请参考软件安装完成安装。
--use-ascend-coc --coc-fused-kernel # 截至当前版本，专精于支持TP（Tensor Parallelism）配置为8的特定场景
- 融合算子的环境变量拥有更高优先级，即当“ coc-parallel-num”配值大于1且使能“coc-fused-kernel”时，前者“ coc-parallel-num”不会生效。
- 鉴于机制层面的潜在冲突，计算通信并行优化与TP重计算通信优化特性目前无法同时启用。建议用户依据实际工作负载与性能需求，独立选用上述任一特性进行配置，以确保最佳的运行状态与资源利用效率。
使用效果
采用CoC融合算子优化后，在BLOOM 7B、BLOOM 176B与LLAMA2 70B模型的训练中，分别实现了端到端性能提升约3.20%、5.47%与7.85%。精度相对误差控制在±2%范围内，确保了模型训练的准确性。