MoE All2All 重叠

MoE All2All 重叠#

LoongForge 为 MoE 模型提供面向通信的优化。通过添加相应的 Megatron 启动参数，专家并行（EP）中的 All-to-All 通信可以与计算重叠，实现最佳训练吞吐量。

1. 1F1B A2A 重叠#

在 MoE 训练中，EP 引起的 All-to-All 通常是主要瓶颈之一。受 DSv3 DualPipe 启发，本框架在时间轴上将不同微批次的前向/反向计算与 EP All-to-All 交织进行。代价高昂的 EP 通信因此被隐藏在计算之后，大大减少了其对整体吞吐量的影响。

特性#

  • 通过微批次重叠隐藏 All-to-All 通信

  • 拆分权重梯度和激活梯度传输，以获得更好的流水线并行重叠

  • 支持 FP8 训练

使用方法#

添加以下启动参数：

--overlap-moe-expert-parallel-comm \
--delay-wgrad-compute

前提条件：必须启用 Interleave 1F1B 调度。

为获得最佳重叠效果，建议：

export CUDA_DEVICE_MAX_CONNECTIONS=32

（这可能会略微影响 TP 重叠；请根据 TP 还是 EP 占据通信主导来选择。）

2. 细粒度激活卸载#

在长上下文训练中，激活内存随序列长度快速增长，很快成为限制因素。常见的补救方法是将张量并行（TP）与全层重计算结合以压缩激活占用。 然而，1F1B A2A 重叠策略依赖于相邻批次的逐模块交织，使得传统的全层重计算不兼容。

为解决此问题，框架引入了模块级选择性重计算加细粒度激活卸载，在保留重叠调度的同时近似全层重计算的内存节省（见图）。 offload_stream

特性#

  • 激活卸载与重新加载隐藏在计算之后

  • 张量级卸载粒度

  • 支持 FP8 训练

使用方法#

启用模块级选择性重计算：

--recompute-granularity selective \
--recompute-modules a2a_overlap_attn a2a_overlap_post_attn a2a_overlap_mlp

启用张量级激活卸载：

--fine-grained-activation-offloading \
--offload-tensors dispatched_input pre_mlp_layernorm_output

此外，将每个进程绑定到其 GPU 所在的 NUMA 节点，以提高 D2H/H2D 带宽：

--bindpcie