Leoda

1. MoE EP Overlap

4 min read

前言

关于TP = 1的bug

首先在早期我开 TP=1的时候在 moe 的 1F1B ep overlap 的时候，会出现某个 gemm 算子输入为 null 的情况，而 moe training 大部分为 tp1，所以在要想打开 --overlap-moe-expert-parallel-comm 则需要设置 "mlp": False，如下：

../Megatron-LM/megatron/core/models/gpt/fine_grained_callables.py
@internal_api
def should_free_input(name, is_moe, config):
    // ...
    free_input_nodes = {
        "mlp": False,
        "moe_combine": True,
        "moe_dispatch": not (enable_deepep or enable_hybridep)
            and (CudaGraphScope.moe_preprocess not in config.cuda_graph_scope),
    }
    // ...

以上为我在2025.10.14在 Megatron-LM 提的 issue : https://github.com/issues/created?issue=NVIDIA%7CMegatron-LM%7C1862

MoE 1F1B overlap

这里参考了 megatron 的论文内的图，核心就会原来不 overlap 的话 forward+backward 会和 a2a 纯串行，然后 baseline 内升级了两个 stream+两个 microbatch 的 overlap，这样就可以让通讯和计算 overlap。重点是 W/D 的进一步拆分后的 overlap：

  • MB2 在 backward 的 mlp 阶段，权重和激活的计算拆开，激活的 Dgrad 先计算完，那么 MB2 的 Dispatch 就可以很快执行了，让其他 rank 也可以早早开始自己 layersN-1 的链式求导。
  • MB2 的 attn 阶段同样把激活的计算拆开，一结束就可以让 pp backward的 send/recv提前并与 MB2 的权重的 attn 计算 overlap。 此时整个过程只剩下 2 个 babble；

image.png

img_v3_02135_8736069d-6ddb-4bcd-ac6c-5cff455eecfg.png

fwd与 bwd 的 8 个操作怎么拆出来的

在一个 layer 的 scheduler 初始化的时候（TransformerLayerSchedulePlan），会去 build_transformer_layer_callables，这个函数就是来拆出 5 个相位，分别是 attn, dispatch, mlp, combine,mtp。并且在 attn 和 dispatch 这俩相位做了两层 detach。如下：

第一层就是 attn→ dispatch → mlp → combine→ 这个主路径，每个箭头一个 detach。

第二层就是 attn 和 dispatch 内部产出了跳过相位的桥接张量时也需要 detach。如 attn 内的 residual 和 shared_expert_output张量跳过 dispatch 和 mlp，直接参与 combine；

拆分后，每个相位单独包成一个 ScheduleNode 类，提供每个对象的 forward 和 backward。构造的时候接收 stream，event，forward_func和 backward_func。其中 backward 是直接走的 Variable._execution_engine.run_backward（） 接口。在ScheduleNode.forward()的入口处就是上面所说的第一层 detach。

# 直接拿出了 pytorch autograd 引擎最薄的 wrapper
torch.autograd.backward(loss)          # 高层公开 API
  └─ torch.autograd.grad(...)          # 算梯度
       └─ Variable._execution_engine.run_backward(...)  # ← 就是这行
            └─ C++ autograd engine (torch/csrc/autograd/engine.cpp)
           
Variable._execution_engine.run_backward(
    tensors=outputs,          # 反向传播的起点：前向输出（即 loss / 中间输出）
    grad_tensors=output_grad, # 从下游传进来的梯度（即 ∂L/∂output），作为初始 seed
    keep_graph=False,         # 反向后释放计算图（一次性 backward）
    create_graph=False,       # 不建二阶图（非 double backward 场景）
    inputs=tuple(),           # 空元组=对所有 requires_grad 的叶子求梯度
    allow_unreachable=True,   # 允许某些输出不可达（图断了不报错）
    accumulate_grad=True,     # 梯度累加到 .grad（而不是覆盖）
)

run_backward返回 grad 在上层调度阶段通过自己手动往前传，如下：

# /Users/joker/Desktop/project/infrawaves/vccl-megatron/megatron/core/models/common/model_chunk_schedule_plan.py
# TransformerLayerSchedulePlan 的 run 函数
if b_layer is not None:
    b_grad = b_layer.mtp_post_process.backward(b_grad)   # ① 起点：b_grad = g_loss
    b_grad = b_layer.moe_combine.backward(b_grad)        # ② combine.backward(g_loss) → 返回 combine输入.grad → 赋给 b_grad
 
if b_layer is not None:
    b_grad = b_layer.mlp.backward(b_grad)                # ③ mlp.backward(b_grad) → b_grad 现在是 mlp输入.grad
 
if b_layer is not None:
    b_grad = b_layer.moe_dispatch.backward(b_grad)       # ④ dispatch.backward(b_grad) → b_grad = dispatch输入.grad
 
if b_layer is not None:
    b_grad = b_layer.attn.backward(b_grad)               # ⑤ attn.backward(b_grad) → b_grad = attn输入.grad

Graph View

Backlinks