高维张量并行

背景与挑战

在大模型训练中，张量并行（TP）通过将模型参数分配到多个设备上来减少内存占用。在训练过程中，为了更新参数梯度等信息，需要引入allreduce通信。当集群规模较大时，若TP域设置得过大，通信开销会显著增加，从而降低训练效率。

解决方案

为了提高大规模TP域的通信效率，采用了高维张量并行技术，该技术将激活值和参数同时切分到多个计算设备上，相比1D-TP，这种方法减少了通信域，降低了通信次数，从而缩短了通信时间，提升了模型训练的性能。

• 2D张量并行策略

  给定TP域的大小，通过构建多通信域，在原有的Megatron（ColumnParallelLinear、RowParallelLinear）基础上增加了一个切分维度。将原有的TP通信域分解为两个子通信域tp_x和tp_y，需满足tp = tp_x * tp_y。以MLP层为例，其实现过程如下图。

  图1 2D张量并行策略实现
  
• 分布式normalization

  在Transformer网络中，normalization会将每一层神经元的输入转换为具有相同均值和方差的形式，从而加快其收敛速度。在MLP和attention层分别进行2D张量并行时，其输入和输出分别在first-dim和last-dim上进行了tp_x和tp_y的切分。如果继续使用原LayerNorm或RMSNorm，需要先将input进行沿first-dim进行all-gather(x)和沿last-dim进行all-gather(y)操作，以确保input数据的完整性。为了提升这部分的性能，采用了分布式normalization。其处理流程如下：

  1. 计算输入的总和

     计算输入张量x在最后一个维度上的总和，公式如下：

     

  2. 分布式归约操作（All-Reduce）

     将第1步中的总和e_x在所有tp_y通信域进程中进行归约（求和），确保每个进程都拥有其通信域全局总和，公式如下：

     

     • P是分布式进程的数量。
     • 表示第p个进程中第i个元素的值。
  3. 计算输入元素的平方和

     计算输入张量每个元素的平方和，公式如下：

     

  4. 分布式归约操作（All-Reduce）

     将第3步中的平方和s_x在所有tp_y通信域进程中进行归约（求和），确保每个进程都拥有其通信域全局平方和，公式如下：

     

  5. 中心化输入数据

     将输入数据x中心化，即减去平均值。平均值μ计算如下：

     

     然后，中心化输入：

     

  6. 计算总和的平方

     计算全局总和的平方，公式如下：

     

  7. 计算归一化因子

     计算归一化因子γ，用于标准化输入数据，公式如下：

     

     • 是全局平方和的平均值。
     • 是全局总和的平方。
     • 是一个小常数，防止分母为零，增加数值稳定性。
  8. 标准化输入数据

     将中心化后的输入数据x′与归一化因子γ相乘，得到标准化后的数据 ，公式如下：

     

  9. 应用权重和偏置

     将标准化后的数据与权重向量W相乘，并根据是否存在偏置向量b来决定最终输出。

     • 如果存在偏置：

       

     • 如果不存在偏置：

       

使用场景

当TP通信域需要设置较大时，通信效率较低，需要通过分解通信域来提升其通信效率。

使用方法

在训练脚本的参数列表中加入--tp-2d，开启2D张量并行，--tp-x N1和--tp-y N2分别设置其x轴、y轴的切分大小，其中需满足tp = N1 * N2(N1 > 1, N2 > 1)。

其他优化参数，用于辅助高维张量并行特性进行通信隐藏，需要开启tp-2d时生效：

• --enable-overlap-ag-with-matmul：在linear层forward计算时，开启all-gather通信和matmul进行隐藏，以便加速。
• --enable-overlap-matmul-with-rs：在linear层forward计算时，开启matmul计算和reduce-scatter通信进行隐藏，以便加速。
• --coc-fused-kernel：在linear层forward计算时，开启计算通信融合算子，将matmul计算与all-gather、reduce-scatter都进行算子级融合，实现进一步加速。
  

  该特性不与前两个特性兼容，依赖CANN NNAL（Ascend Transformer Boost加速库）。

上述3个forward计算优化参数--enable-overlap-ag-with-matmul、--enable-overlap-matmul-with-rs、--coc-fused-kernel只能同时开启1个。

当前高维张量并行特性不与--sequence-parallel、--use-fused-rmsnorm等特性相兼容，请根据实际情况调整配置。

使用效果

在LLaMA3-405B模型训练时，当tp = 16，启用2D张量并行（tp_x = 8，tp_y = 2）的情况下，相比原Megatron 1D张量并行，性能提升了5%以上。启用coc-fused-kernel和enable-backward-overlap-ag-with-matmul通信计算融合优化后，性能进一步提升了5%以上。

在其他场景下，由于计算效率和通信组划分的差异，需要根据tp_x和tp_y的实际调优情况来配置，部分配置可能无法保证效率提升。