已开启
[Feature]: HyperGraph 静态图支持 Expert Parallel #357
zhenyu创建于  8月31日
zhenyu成员
8月31日 创建

需求背景

HyperGraph 已具备 Forward、Loss 和 Backward 联合图捕获及 Pass 管线骨架,但当前静态图路径尚未支持 Expert Parallel。

AutoModels 动态 EP 使用数据相关的 token dispatch、专家分片和 AllToAll 通信。若重新实现一套静态 EP,不仅会重复维护路由与专家计算逻辑,也容易造成动态图与静态图语义不一致。

本需求目标是复用 AutoModels 已有的动态 EP 路由、专家分片和专家权重能力,将 EP 通信和计算稳定捕获到 HyperGraph 静态图中。首阶段仅支持 EP-only 拓扑,TP、CP、PP 和 FSDP 不在本次范围内。

开发设计

  • 在 hyper_parallel/compile 中增加动态 EP capture adapter。
  • 保留 AutoModels 对 router、expert sharding 和 expert weights 的所有权,不新增独立 EP 实现。
  • 使用固定容量 dispatch/combine buffer,避免数据相关 token 数转为 Python 值导致图结构变化。
  • 使用显式 autograd AllToAll,保证正向与反向通信均可捕获和重放。
  • 扩展 ParallelConfig,增加 EP degree、拓扑和 process group 校验。
  • 在 Pass 管线中增加 EP evidence pass,检查联合图中存在 AllToAll 通信节点,避免出现“配置开启但 EP 未入图”。
  • 提供 Qwen3-MoE 静态 EP 训练入口。
  • ep_degree=1 时保持原有路径不变。

已知限制和风险:

  • 当前只支持 world_size == ep_degree 的 EP-only 拓扑。
  • 固定容量 buffer 会增加临时显存占用,暂不作为性能优化方案。
  • TP、CP、PP、FSDP 与静态 EP 的组合留待后续独立适配和验证。
  • 当前验证模型为 Qwen3-MoE。

验收标准

  • 2 rank Qwen3-MoE 可完成静态图 Forward、Backward 和梯度回传。
  • 捕获的联合图中可以检测到 EP AllToAll 通信节点。
  • 静态 EP 与动态图参考路径的 Loss 最大绝对误差不超过 1e-6。
  • 静态 EP 与动态图参考路径的梯度最大绝对误差不超过 1e-6。
  • 非法并行度、未初始化 process group、world size 不匹配及专家分片不匹配能够提前报错。
  • 固定容量 AllToAll 的正反向和 capture adapter 恢复逻辑具有硬件无关单元测试。

验证计划

已执行:

  • 静态 EP 单元测试:14 passed, 5 subtests passed。
  • 2 rank 静态 EP 验证连续执行两次,均通过。
  • 捕获的 collective 数量:5。
  • Loss 最大绝对误差:4.76837158203125e-07。
  • 梯度最大绝对误差:0.0。

本需求不包含性能收益验收;性能与混合并行组合将在后续需求中单独验证。

likedislike