Megatron权重更新通信隐藏

背景与挑战

在大模型的训练过程中，数据并行策略被普遍采用。在梯度更新阶段，传统方法要求数据并行组内的通信操作必须在反向传播计算完全结束后才启动，这种串行执行模式导致计算与通信流程中存在明显的空闲等待期，从而降低了整体执行效率。

解决方案

为解决上述问题，引入了计算与通信任务的并行执行策略，通过流水线技术来实现计算与通信的流水掩盖，有效提升资源利用率。

• 仅启用分布式优化器（--use-distributed-optimizer）
  在前向与反向计算结束后，将进行独立的通信阶段，包括梯度的reduce-scatter、权重计算以及权重的all-gather，如下图所示。获取更新后的权重后，系统将进入下一轮的前向计算阶段。

  图1 仅启用分布式优化器时的计算流和通信流
  

• 同时启用分布式优化器与梯度通信重叠（--use-distributed-optimizer + --overlap-grad-reduce）
  对梯度的reduce-scatter过程与反向计算过程并行，从而避免了额外的reduce-scatter时间，显著提高了计算与通信的并行效率，如下图所示。

  图2 同时启用分布式优化器和梯度通信重叠时的计算流和通信流
  
• 启用分布式优化器、梯度通信重叠与参数聚合重叠（--use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather）
  在上述基础（启用分布式优化器与梯度通信重叠）上，进一步启用参数聚合重叠，如下图所示。all-gather操作与下一轮的前向计算并行执行，从而避免了独立的all-gather时间段。

  图3 启用分布式优化器、梯度通信重叠、参数聚合重叠时的计算流和通信流
  

  通过对比不同流程，可以观察到在启用--overlap-param-gather后，计算与通信完全并行，极大地提高了计算与通信的并行效率，从而显著提升了模型训练的效率。

使用场景

该特性适用于采用数据并行策略的训练场景，特别是当通信开销不可忽视时。“权重更新通信隐藏”通过减少通信延迟对整体训练时间的影响，显著提升了模型训练速度。

使用方法

• 要启用权重更新通信隐藏功能，需在训练配置中加入以下参数。

  --overlap-param-gather

• 确保同时开启了以下两个参数。

  --use-distributed-optimizer
  --overlap-grad-reduce

• Megatron原生版本存在生成顺序有误的问题，该异常导致下一轮的前向计算提前。开启该特性后，为了修复此异常，attention层init的顺序会更正为先创建linear_qkv再创建linear_proj。
• Legacy下，--overlap-param-gather暂不支持和--reuse-fp32-param一起使用。

使用效果

权重更新通信隐藏通过实现计算与通信任务的并行执行，提高大模型训练中的资源利用率和训练效率。在使用LLaMA2-70B中应用权重更新通信隐藏技术后，端到端性能提升约2%。