【RFC】FSDP + EP 组合支持与 Pipeline 微批次状态安全

FSDP + EP 组合支持与 Pipeline 微批次状态安全

本文档描述 HyperParallel Expert Parallel 与 FSDP（Fully Sharded Data Parallelism）的组合验证、以及 ExpertParallel 在流水线并行多微批次场景下的状态管理安全问题，不包含具体代码实现，侧重背景、职责边界、接口契约、测试与验收。

---

1. 背景

1.1 训练配置演进

大规模 MoE 模型训练通常需要同时启用多种并行策略：

• 数据并行 / FSDP：跨多个 DP rank 复制模型逻辑、分片参数存储，各 rank 处理不同数据样本；
• Expert Parallel（EP）：将专家层按专家维切分到不同设备，token 经 all-to-all 路由；
• Pipeline Parallel（PP）：将模型按层切分为多个 Stage，以微批次（micro-batch）流水执行。

三者叠加时，EP 所在的 ExpertParallel 并行样式需要在多种 hook 执行序列、参数 unshard/shard 生命周期、以及多微批次并发调用的环境中保持正确性。

1.2 现有实现

hyper_parallel/core/expert_parallel/expert_parallel.py 的 ExpertParallel 通过 distribute_module 在 GroupedExperts 子模块上注册：

• _token_dispatch（pre-hook）：执行计数 all-to-all、token all-to-all、permute，将中间状态存入 实例变量 self._input_splits、self._output_splits、self._input_shape、self._permuted_indices；
• _token_combine（post-hook）：读取上述实例变量执行 unpermute 和反向 all-to-all。

文档 docs/expert_parallel.md 已提及与 FSDP 的兼容性意图，但缺乏分布式集成测试与对状态安全的显式保证。

---

2. 已知问题

2.1 实例变量状态在多微批次场景下不安全

_token_dispatch 将本次 forward 的中间状态写入 self._* 实例变量，_token_combine 在同一次 forward 结束时消费这些变量。在单微批次场景下，两个 hook 严格配对，行为正确。

在流水线并行多微批次场景下：

• Pipeline scheduler 可能在第 N 个微批次的 _token_combine 执行之前，即发起第 N+1 个微批次的 _token_dispatch；
• 后者覆盖实例变量，导致第 N 个微批次的 _token_combine 读取到错误状态；
• 此类错误在 loss 层面可能不立即显现，属于静默数据错误。

2.2 FSDP 与 distribute_module Hook 顺序未验证

FSDP 对被包装模块注册自身的 pre-hook（参数 unshard）和 post-hook（参数 re-shard）。ExpertParallel 的 _token_dispatch / _token_combine 同样通过 pre/post hook 接入。两者叠加时：

• hook 的注册顺序决定执行顺序；
• 若 FSDP 的参数 unshard hook 在 _token_dispatch 之后执行，则 dispatch 阶段可能访问尚未 unshard 的专家权重；
• 若 FSDP 的 re-shard 在 _token_combine 之前触发，combine 结果写入时权重已被释放。

上述顺序在当前代码中均未经过明确测试。

2.3 FSDP 梯度归约与 EP Combine 反向梯度的重叠

FSDP 通过 reduce-scatter 聚合梯度，EP 的 combine 反向（即 differentiable_all_to_all_single 的反向传播）也产生针对 GroupedExperts 权重的梯度贡献。两者之间是否会在某些分片维度上产生重复归约，尚未有实测验证。

---

3. 目标与非目标

3.1 目标

• 状态安全：将 _token_dispatch 与 _token_combine 之间的中间状态从实例变量迁移为调用栈局部传递（dispatch 返回 context 对象，combine 接收该 context），从根本上消除多微批次覆写风险；
• FSDP + EP 正确性验证：提供分布式 ST 测试，在典型 FSDP + EP 配置下验证前向输出与反向梯度的数值正确性；
• Hook 顺序文档化：在代码注释或 docs/expert_parallel.md 中明确 FSDP 与 EP hook 的预期执行顺序，并在测试中断言该顺序。

3.2 非目标

• 不引入新的并行策略；
• 不修改 FSDP 自身的参数管理逻辑；
• Pipeline + EP 的性能优化（如通算掩盖）属于独立特性，不在本文档范围内。

---

4. 接口与契约

4.1 状态传递方式变更

当前：_token_dispatch 写 self._*，_token_combine 读 self._*。

目标：

• _token_dispatch 返回 (transformed_inputs, dispatch_ctx)，其中 dispatch_ctx 是携带本次 forward 所有中间状态的不可变对象（或具名元组）；
• _token_combine 接收 (module, output, dispatch_ctx, device_mesh)，不再从实例变量读取。

要求：

• dispatch_ctx 不持有对 self 的引用，生命周期与调用栈绑定；
• distribute_module 的 hook 签名需支持此状态透传机制（若现有接口不支持，需评估是否扩展 distribute_module API 或改用显式 wrapper 模式）；
• 改动对调用方（MoE.forward + ExpertParallel.apply）透明，不改变外部使用接口。

4.2 FSDP + EP 的组合约束

约束	说明
FSDP 应包裹整个 MoE，EP 应用于 MoE.experts 子模块	确保 FSDP 的参数生命周期在 EP hook 执行时已完成 unshard
EP 的 distribute_module 调用应在 FSDP 包裹之后	避免 FSDP re-wrap 覆盖 EP hook
GroupedExperts 的参数在 EP dispatch 前必须已 unshard	由 FSDP pre-hook 保证；EP pre-hook 需排列在 FSDP pre-hook 之后

上述顺序约束需在 docs/expert_parallel.md 中以使用注意事项形式明确说明，并在测试中断言。

4.3 对外 API 不变

ExpertParallel、ExpertTensorParallel、TensorParallel 的 apply(module, device_mesh) 接口签名保持不变；内部状态传递方式的变更对调用方不可见。

---

5. 测试设计

5.1 状态安全单元测试

用例 ID	描述	期望
SS-01	模拟两个微批次并发调用：在第 1 个微批次的 _token_dispatch 返回后、_token_combine 执行前，插入第 2 个微批次的 _token_dispatch	两个微批次的 combine 分别读取各自正确的 context，输出均与单独运行时一致
SS-02	验证 dispatch_ctx 的生命周期：combine 后 ctx 不被 ExpertParallel 持有	无内存泄漏，gc 可正常回收

5.2 FSDP + EP 分布式测试

用例 ID	描述	期望
FE-01	4 卡，FSDP（2 卡）× EP（2 卡），小 MoE 模型前向 + 反向	前向输出与单卡参考（相同 token 子集）数值对齐（容差 rtol=1e-3）
FE-02	同 FE-01，验证 GroupedExperts 各参数的梯度非 None 且有限	梯度正确流通
FE-03	FSDP + EP，连续 3 步训练，loss 单调下降（小规模过拟合验证）	训练可正常收敛

5.3 Hook 顺序断言

用例 ID	描述	期望
HO-01	注册 spy hook，记录 FSDP pre-hook 与 EP _token_dispatch 的执行顺序	FSDP unshard hook 先于 EP dispatch hook 执行

5.4 回归

• 现有 8 个 Ascend NPU EP 分布式测试保持通过；
• 现有 UT（51 个）保持通过。

---

6. 风险与开放问题

风险	缓解
distribute_module 当前 hook 签名不支持 context 透传	评估是否通过 functools.partial 或 wrapper class 在不修改 distribute_module 接口的前提下实现透传
FSDP 与 EP 在 hook 注册顺序上的 PyTorch 版本差异	在集成测试中锁定 PyTorch 版本，并在 CHANGELOG 中标注
FSDP reduce-scatter 与 EP combine 反向的梯度维度重叠	在 FE-02 中加入梯度值对比（与无 FSDP 的 EP baseline 对比绝对误差），如有偏差需单独分析

---

7. 验收标准

• ExpertParallel 内部不再使用实例变量传递 dispatch/combine 状态；
• SS-01 验证多微批次场景状态隔离正确；
• FE-01～FE-02 在 FSDP + EP 组合下数值正确性通过；
• docs/expert_parallel.md 补充 FSDP + EP 组合使用说明与 hook 顺序约定；
• 现有 EP 所有 UT 与 ST 无回归。

---

8. 参考

• HyperParallel：hyper_parallel/core/expert_parallel/expert_parallel.py — ExpertParallel._token_dispatch、_token_combine；
• HyperParallel：hyper_parallel/core/fully_shard/ — FSDP 参数生命周期与 hook 注册；
• HyperParallel：docs/expert_parallel.md — 现有 EP 使用说明。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 106
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/106

Start with hyper_parallel/core/expert_parallel/expert_parallel.py, especially ExpertParallel._token_dispatch and _token_combine, then inspect hyper_parallel/core/fully_shard/ and docs/expert_parallel.md. Review the existing eight Ascend NPU EP distributed tests and 51 unit tests before designing state-isolation and FSDP+EP coverage. Done means dispatch state is isolated per micro-batch, hook order and numerical correctness are tested, documentation is updated, and existing tests still pass.