已开启
【RFC】EP 通算掩盖:共享专家重叠、DualPipe 与 W/D 掩盖 #110
hedongdong创建于 4月27日
6月2日 修改了issue 的描述
6月2日 修改了issue 的描述
6月4日 关联了pull request:feat(expert_parallel): add async combine overlap for shared expert optimization
6月4日 关联了pull request:feat(expert_parallel): add async combine overlap for shared expert optimization
6月10日 关联了pull request:feat(expert_parallel): async combine a2a overlap with shared expert
7月1日 关联了pull request:fix(expert_parallel): fix async combine distributed tests using wrong init helper
7月1日 关联了pull request:fix(expert_parallel): fix async combine distributed tests using wrong init helper
EP 通算掩盖:共享专家重叠、1F1B 通算掩盖、W/D 拆分与 DualPipe
本文档描述 Expert Parallel 的四类计算与通信重叠优化:
examples/torch/pp_overlap/pp_overlap_moe_example.py,待产品化)侧重背景、职责边界、接口与契约、测试与验收,不包含具体代码实现。
1. 背景
1.1 EP 通信的时序瓶颈
标准 EP 的 forward 执行时序为:
三阶段串行执行,计算与通信完全顺序。在
ep_degree较大、机间带宽受限时,两次 all-to-all 的时延直接叠加在关键路径上,成为 MoE 层的主要耗时。反向链路上还会再叠加combine_bwd与dispatch_bwd两次 a2a,通信总量比前向更大。1.2 已有基础与缺口
已实现:
AllToAllTokenDispatcher.dispatch()与combine()的基础流程(同步 all-to-all);MoE内的shared_expert(可选FeedForward,输出与 expert 输出相加);platform.differentiable_all_to_all_single同步版本;platform.differentiable_all_to_all_single_async异步版本(已在 example 中验证);CommComputeOverlap多线程通算协调器与 A/B/C/D 同步 hook 原语(example 中使用);ScheduleInterleaved1F1B(overlap_b_f=True)调度器,能在 1F1B 稳态发出OVERLAP_B_F复合步。缺口:
shared_expert与combineall-to-all 串行执行(A 的目标);ExpertParallel把_input_splits / _output_splits / _input_shape / _permuted_indices挂在 EP 实例上,跨层只能靠"每层一个 EP 实例"绕开,跨 microbatch 不安全(B 产品化前置);2. 四类优化的定义与范围
2.1 子方向 A:共享专家与 combine 通信异步重叠
核心思想:
combine all-to-all 发起后,在等待通信结果的时间窗口内,并行执行 shared expert 的前向计算,将两者的时延进行掩盖。
时序对比:
依赖:
platform.differentiable_all_to_all_single支持异步发起,返回句柄(AsyncHandle);TokenDispatcher.combine()分解为「发起」和「等待 + 合并」两个阶段;2.2 子方向 B:1F1B 通算掩盖(BWD ↔ FWD a2a 跨线程交错)
核心思想:
ScheduleInterleaved1F1B的 1F1B 稳态把 BWD microbatchi与 FWD microbatchi+1(来自不同虚拟 chunk)打包成OVERLAP_B_F复合步,两条线程分别承载 FWD 与 BWD:CommComputeOverlap的协调器全局串行(同 group 两个 a2a 并发会触发 HCCL 数据损坏);differentiable_sync_hook把 dispatch a2a 和 combine a2a 各自夹在一对 hook 中间,hook 之间是 compute 区,coordinator 在 hook 上交接 HCCL 发射权。时序示意(一对 BWD↔FWD 配对内的一层 MoE):
sequenceDiagram participant CS as 计算流 participant CO as Coordinator participant XS as 通信流(HCCL)<br/>同 ep_group 全局串行 Note over CS,XS: 配对开始<br/>FWD microbatch i+1(MoE 层 L)<br/>BWD microbatch i(MoE 层 L+1,反向遍历) rect rgb(220, 240, 220) Note left of CS: 并发计算:<br/>BWD 上一层 expert_w 残段 CO->>XS: FWD dispatch a2a(hook A→B 内异步发射) XS-->>CO: 完成 end rect rgb(220, 240, 220) Note left of CS: 并发计算:<br/>FWD permute + expert FFN(hook B 之后的 compute 区) CO->>XS: BWD combine_bwd a2a 异步发射 XS-->>CO: 完成 end rect rgb(220, 240, 220) Note left of CS: 并发计算:<br/>BWD expert_bwd D 段 CO->>XS: FWD combine a2a(hook C→D 内异步发射) XS-->>CO: 完成 end rect rgb(220, 240, 220) Note left of CS: 并发计算:<br/>FWD 下一层 attn / qkv CO->>XS: BWD dispatch_bwd a2a 异步发射 XS-->>CO: 完成 end Note over CS,XS: 进入下一对 BWD↔FWD 配对与文档前一版「DualPipeExpertParallel」的关系:前一版用「Micro-batch N 的 combine 与 Micro-batch N+1 的 dispatch 在前向接力」描述 B,落到 HCCL/NCCL 上同 group 两 a2a 无法真正并发,需要 BWD 同时在场才能用 BWD 的计算去掩盖 FWD 的通信。example 的实现是这一观察的工程化产物,与 DeepSeek 完整 DualPipe(见 2.4)不同,没有双向流水线、没有 chunk 内 4 段细粒度切分、没有 W/D 拆分,但已能拿到主要收益。为避免名称混淆,本期内部命名采用
OverlapExpertParallel+CommComputeOverlap+ScheduleInterleaved1F1B(overlap_b_f=True)三件套,不再使用 "DualPipeExpertParallel" 这一名称。依赖:
platform.differentiable_all_to_all_single_async);CommComputeOverlap+differentiable_sync_hook("A"|"B"|"C"|"D"|"D_LAST")协调器;ScheduleInterleaved1F1B支持overlap_b_f=True/overlap_p2p=True;_input_splits等状态(或迁移至DispatchMetadata栈式对象);is_last_layer=True标记,使 D hook 退化为D_LAST(前向后无 attn、反向 combine_bwd 已自由发射,无需 rendezvous);set_device(local):current device 是 thread-local,缺失会让 HCCL 跑到 device 0 死锁。2.3 子方向 C:W/D 拆分 + deferred W 队列(本期不实施,仅定义范围)
核心思想:
反向天然由
D = ∂L/∂x(输入梯度,传给上层)+W = ∂L/∂w(权重梯度,本地消费)两部分组成。标准 autograd 把它们绑在同一个 backward node 一次性算完。W/D 拆分把它们解耦:时序示意:
flowchart TB subgraph standard[标准反向(B 已实现的状态)] direction LR s1[layer L+1<br/>expert_bwd<br/>同时算 dX+dW] --> s2[a2a comm<br/>对侧 compute<br/>已被掩盖] s2 --> s3[layer L<br/>expert_bwd<br/>同时算 dX+dW] end subgraph withwd[W/D 拆分反向] direction LR w1[layer L+1<br/>expert_bwd 仅 dX] --> w2[a2a comm 发射] w1 -.W 任务入队<br/>持有激活.-> wq[(deferred<br/>W queue)] w2 --> w4[layer L<br/>expert_bwd 仅 dX] w4 -.W 任务入队.-> wq w2 -.flush 填残余窗口.-> wq end style s2 fill:#ffe0b0 style wq fill:#c0e8c0范围(一旦 Phase 2 立项需要改的面):
expert/attn_qkv/attn_proj/ffn关键 GEMM 上提供backward_d_only(grad_out, weight) → grad_in与backward_w_only(grad_out, input_saved) → grad_w两个独立入口,绕开标准 autograd 的「原子 backward」假设。(grad_out_ref, input_saved_ref, weight_ref, target_grad_buffer)。OVERLAP_B_F复合步的 a2a 等待窗口里调度 W 任务;在 1F1B cooldown 阶段统一 drain 剩余 W。本期不实施的原因:
2.4 子方向 D:双向流水线 DualPipe(远期规划)
核心思想:
DeepSeek-V3 DualPipe 与本仓现状(含 B/C)的根本区别是 pipeline 调度:
(P - 1) × (F + B);k和反向流的 stage2P - 1 - k),来自两条对头 microbatch 流,气泡 ≈(P/2 - 1) × (F + B - 2W)。每个 rank 任一时刻同时有一个 FWD microbatch 与一个 BWD microbatch 在场,且这两个 microbatch 来自不同方向的 pipeline 流、互无数据依赖,比 1F1B 稳态下「同条流的 BWD_i + FWD_{i+1}」并行度更高。
示意:
flowchart LR subgraph baseline["标准 1F1B(单向)"] direction LR s0["rank 0<br/>stage 0"] --> s1["rank 1<br/>stage 1"] --> s2["rank 2<br/>stage 2"] --> s3["rank 3<br/>stage 3"] s3 -.反向回传.-> s2 s2 -.-> s1 s1 -.-> s0 end subgraph dualpipe["DualPipe(双向,每 rank 持 2 stage)"] direction LR d0["rank 0<br/>stage 0 (正)<br/>stage 7 (反)"] d1["rank 1<br/>stage 1 (正)<br/>stage 6 (反)"] d2["rank 2<br/>stage 2 (正)<br/>stage 5 (反)"] d3["rank 3<br/>stage 3 (正)<br/>stage 4 (反)"] d0 ==正向流==> d1 ==> d2 ==> d3 d3 ==反向流==> d2 ==> d1 ==> d0 end style dualpipe fill:#eefchunk 内 4 段交错:DualPipe 把每个 chunk 切成 ATTN / DISPATCH / MLP / COMBINE 四段,让 FWD chunk 的 comm 段与对头 BWD chunk 的 compute 段两两咬合,比 B 的 A/B/C/D 4 hook 粒度更细。
与 B/C 的关系:
-2W项就是 W 段填洞带来的;core/pipeline_parallel/引入新调度器,与现有ScheduleInterleaved1F1B并列。本期不立项的原因:
pp_size ≥ 16的大规模 MoE 才显著,需要先有明确用户场景;3. 目标与非目标
3.1 目标(本期)
子方向 A(本期唯一交付项):
platform.differentiable_all_to_all_single,支持异步发起模式,返回AsyncHandle;AllToAllTokenDispatcher.combine()分解为combine_start()和combine_wait()(或通过DispatchMetadata携带句柄);MoE.forward在combine_start()和combine_wait()之间执行 shared_expert;3.2 后续路线(按阶段排期,本期不交付)
Phase 1 — 子方向 B 产品化:
examples/torch/pp_overlap/pp_overlap_moe_example.py中的OverlapExpertParallel抬成正式 API,纳入hyper_parallel/core/expert_parallel/;DispatchMetadata或 MoE 模块,去掉「每层一个 EP 实例」的 workaround;Phase 2 — 子方向 C W/D 拆分:
ScheduleInterleaved1F1B集成;Phase 3 — 子方向 D 评估:
pp_size ≥ 16的真实用户场景时立项;3.3 非目标
4. 接口与契约
4.1 异步 all-to-all 平台接口(A 依赖)
async_op=False(默认)async_op=True(handle, output_tensor);调用方须在使用output_tensor前调用handle.wait()约束:
async_op=True时,output_tensor是已分配但未填充的 Tensor,访问需在handle.wait()后;handle.wait()是幂等的;platform/platform.py)需在抽象接口中声明此参数,PyTorch 和 MindSpore 分别实现;platform.differentiable_all_to_all_single_async,含义等价。4.2
AllToAllTokenDispatcher的分阶段 combine(A)CombinePartialResult携带 unpermute 所需的中间状态,生命周期与本次 forward 调用绑定,不持有对 dispatcher 的引用。向后兼容:保留
combine(expert_output, top_scores, metadata)同步版本,内部调用combine_start + combine_wait,对不需要 shared_expert 重叠的场景透明。4.3
MoE.forward的新时序(A)约束:
shared_expert(shared_expert is None),步骤 3~5 退化为同步combine(),无额外开销;x_flat在步骤 4 中被 shared_expert 使用,需保证在步骤 3 发起通信后未被修改(即不可原地操作)。4.4 B 参考实现接口(Phase 1 产品化目标)
example 已落地的 API(Phase 1 把它们从 example 移入
hyper_parallel/core/expert_parallel/):调度侧入口:
Phase 1 产品化要补:
OverlapExpertParallel仍把_input_splits / _output_splits / _input_shape / _permuted_indices写在自身上,跨 microbatch 不安全;改为通过DispatchMetadata显式回传给combine,与 4.2 的CombinePartialResult统一。is_last_layer自动推断:例子里靠用户手动传,产品化要让框架从 chunk 结构自动判定。differentiable_all_to_all_single_async与多线程协调器。4.5 C 接口草稿(Phase 2,本期不实现,供后续设计参考)
约束:
input_saved),生命周期由调度器管理;flush(deadline_event)在 a2a 等待窗口里调用,事件触发时停止入队;expert_w注册到DeferredWQueue.current()。4.6 D 接口草稿(远期,本期不实现)
前置条件:
4.7 平台支持矩阵
5. 测试设计
5.1 子方向 A 单元测试
combine_start+combine_waitround-trip,无 shared_expertcombine()数值完全一致handle.wait()幂等性:调用两次不报错5.2 子方向 A 分布式测试
5.3 子方向 B 产品化验收要点(Phase 1)
DispatchMetadata化后,跨 microbatch 状态不互相覆盖is_last_layer自动推断与手动指定行为一致5.4 子方向 C / D 验收
在各自后续设计文档定义,本文档不展开。指标方向:
pp_size = 16配置下下降到(P/2 - 1)量级。5.5 回归
6. 风险与开放问题
combine_start后x_flat被 shared_expert 并发访问的内存安全x_flat为 combine 发起前的只读输入,dispatch 结束后不再修改,shared_expert 仅读取,无并发写冲突differentiable_all_to_all_single的 autograd function 需与异步句柄的wait()正确绑定;backward 中wait()必须先于梯度计算_input_splits等迁出 EP 实例,统一走DispatchMetadataD_LAST标签 + 超时机制 + 失败传播;example 已踩过的坑(last layer 不 rendezvous、BWD 线程 set_device)必须复刻CombinePartialResult的内存开销(持有 unpermute 中间状态)permuted_indices相当,可接受7. 验收标准
子方向 A(本期)
子方向 B(Phase 1)
子方向 C(Phase 2)
子方向 D(评估期)
8. 参考
hyper_parallel/platform/torch/common/moe.py—MoE.forward(shared_expert 当前串行位置);hyper_parallel/core/expert_parallel/expert_parallel.py—ExpertParallel._token_dispatch、_token_combine;hyper_parallel/platform/platform.py—differentiable_all_to_all_single平台接口;hyper_parallel/core/pipeline_parallel/— Pipeline scheduler、CommComputeOverlap、ScheduleInterleaved1F1B;examples/torch/pp_overlap/pp_overlap_moe_example.py— B 子方向参考实现;docs/expert_parallel.md;