问题分析

在大规模模型训练中，ColumnParallelLinear和RowParallelLinear组件的前向与反向传播阶段常包含紧密耦合的计算与通信序列，其中计算环节主要涉及矩阵乘法（Matmul），而通信则依赖于AllReduce（在未启用序列并行时）或AllGather与ReduceScatter（在序列并行模式下）。由于这些操作间的顺序依赖性，它们往往被串行执行，导致计算与通信流程中存在空闲等待时间，进而影响整体执行效率。

解决方案

为克服上述局限，我们提出了一套计算与通信并行优化方案，旨在通过将计算与通信任务细化为更小粒度的子任务，实现两者的流水线式并发执行，从而最大化资源利用率与执行效率。

解决思路

Python脚本层面实现

通过进一步分割张量（分为2/4/8个子张量），借助Python脚本机制，实现各子张量间计算与通信的并行处理，尤其适用于Attention模块与MLP模块的串行执行场景。需注意，此方法对Matmul左侧矩阵的m轴尺寸有特定要求，须为分割数的整数倍；且在计算与通信片段耗时差异较大或矩阵分割数量过多时，可能遭遇主机瓶颈，影响预期性能提升。

融合算子实现

依托MTE远端内存访问能力，以融合大Kernel的形式，内部细化计算与通信任务，实现两者的流水线并行执行。目前支持以下融合算子：

MATMUL_ALL_REDUCE融合算子（先计算后通信）及其确定性计算。
MATMUL_REDUCE_SCATTER融合算子（先计算后通信）及其确定性计算。
ALL_GATHER_MATMUL、ALL_GATHER_MATMUL_V2融合算子（先通信后计算，V2版本支持获取ALL_GATHER中间结果）。
量化场景下，MATMUL_ALL_REDUCE融合算子支持fp16格式下的w8A16伪量化，粒度包括per tensor、per channel与per group。

使用场景

计算与通信并行优化技术适用于Attention模块与MLP模块的串行执行，特别是在计算通信存在顺序依赖与位置毗邻的条件下。

使用方法

计算通信并行优化算法通过在训练脚本中配置环境变量来进行使能。

当前计算通信并行有两种实现方法：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 Tensor Boost）。可通过以下两种方式安装ATB：
  - 自动安装：设置环境变量后在MindSpeed目录下执行如下命令
```
ENABLE_ATB=1 pip install -e MindSpeed
```
  - 手动安装：将ATB编译后的output文件夹复制至“MindSpeed/mindspeed/”目录下，并重命名为atb。
- 配置如下参数激活计算通信并行特性。
```
--use-ascend-coc
--coc-fused-kernel # 截至当前版本，专精于支持TP（Tensor Parallelism）配置为8的特定场景
```

融合算子的环境变量拥有更高优先级，即当“ coc-parallel-num”配值大于1且使能coc-fused-kernel时，前者“ coc-parallel-num”不会生效。

使用效果

采用COC融合算子优化后，在BLOOM 7B、BLOOM 176B与LLAMA2 70B模型的训练中，分别实现了端到端性能提升约3.20%、5.47%与7.85%。精度相对误差控制在±2%范围内，确保了模型训练的准确性。

注意事项

鉴于机制层面的潜在冲突，计算通信并行优化与TP重计算通信优化特性目前无法同时启用。我们建议用户依据实际工作负载与性能需求，独立选用上述任一特性进行配置，以确保最佳的运行状态与资源利用效率。