已合并
Optimize the patch for FSDP #34631
yc_233创建于 4月28日
Optimize the patch for FSDP #34631
已合并
Y
yc_233成员
4月28日
4月28日
# 【合入来源】
> **如有社区issue,请关联issue链接**\
> **请勿携带内部流程信息(需求链接、问题单、内部issue等)**
- [x] 需求
- [ ] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
社区 issue:https://gitcode.com/Ascend/pytorch/issues/1788
# 【修改方案】
> 请描述修改内容的具体实现,涉及哪些组件之间进行交互,可以用1、2、3、...进行罗列\
> 如果是需求或者重构类的PR,需要补充详细设计文档(说明上下游组件关系、时序图、类图、DFX能力等内容)
采用“上层复用原生,下层补齐后端语义”的设计:
1. torch-npu 删除 FSDP2 collectives patch:
- 不再覆写 `_get_param_all_gather_inputs`
- 不再覆写 `torch.ops.fsdp.all_gather_copy_in`
- 不再覆写 `FSDPParamGroup.finalize_backward`
- `fully_shard()` 入口只保留 NPU 侧增强 patch,例如内存缓存和 recompute/prefetch 状态管理
2. op-plugin 在 foreach copy 算子中对齐同设备 D2D fast path:
- fast path 判定显式使用 PyTorch foreach 工具函数:
- `at::native::can_use_fast_route(self, src)`
- `check_tensor_dtype_support_base(src)`
- `can_use_fast_route(self, src)` 本身会检查 `self/src` dtype 一致,因此 dtype 不同不会进入 NPU fast path,保持当前 NPU 不支持跨 dtype fast path 的既有语义。
- 命中 fast path 时调用:
```cpp
split_and_exec_npu_cmd_copy(self, src, /*non_blocking=*/true);
```
3. fallback 路径保持调用方传入语义:
- 未命中 fast path 时,`memcpyBatch(self, src, non_blocking)` 不变
- slow path `foreach_tensor_copy_list_kernel_slow_(self, src, non_blocking)` 不变
- `DO_COMPATIBILITY` 回退仍使用原始 `non_blocking`
# 【资料变更】
> 请确认是否涉及资料变更。如涉及,需要在PR中体现,并简要说明修改内容。如不涉及,需填写“不涉及”
不涉及。
# 【接口变更】
> 请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及,需要详细说明接口以及对应的变更内容,同时需要在资料中体现。如不涉及,需填写“不涉及”
不涉及。无新增客户可见接口;FSDP 对外使用方式不变。`_foreach_copy_` 的 NPU 后端实现由 op-plugin 承接,不新增 Python 侧公开接口。
# 【功能验证】
> 说明测试场景,测试方法。如果本次测试方式与常规单元测试不同,请详细说明您的测试步骤\
> 新增/变更内容是否已新增/适配UT测试用例看护,并补充测试自验证截图
测试场景:
1. FSDP collectives 覆写删除后,`fully_shard` 继续复用上游 PyTorch 原生 `finalize_backward`、`_get_param_all_gather_inputs`、`all_gather_copy_in` 主流程;
2. NPU 侧 `_foreach_copy_` 在 op-plugin 中优先走 D2D fast path,不满足条件时回退到已有 batch copy / slow path;
3. 跨 dtype NPU D2D copy 场景在异步排队后结果正确。
测试方法:
- op-plugin 用例:
- `python test/test_v2r1_ops/test_foreach_copy.py`
- `python test/test_v2r2_ops/test_foreach_copy.py`
- torch-npu 侧建议验证 FSDP fully-shard 训练主路径,功能科通过仓内测试用例。
执行结果:
- 已完成代码检查与 push hook 校验通过。
- NPU 环境 UT 结果请以当前 CI/转测执行结果为准。
UT看护:
- 已适配 op-plugin foreach copy 用例,在跨 dtype copy 断言前增加 `torch.npu.synchronize()`,覆盖新增异步 D2D fast path 行为。
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]
- [x] 代码注释完备,正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签,如:feat、fix、refactor、docs、test等
- [x] PR持续集成流水线(CI)执行通过,代码检查无异常


Pull Request已成功合入, 合并人@ascend-robot
(感谢 yc_233 的贡献)ascend-robot
4月28日 评论:
4月28日 评论:
ascend-robot
4月28日 评论:
4月28日 评论:
Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-Ascend/pytorch | ✅ liujunzhu, chengpeng25 (2/2) | ✅ liujunzhu (1/1) |
| test | ✅ liujunzhu, wjlflyer, kisnwang, chengpeng25 (4/2) | ✅ liujunzhu, wjlflyer (2/1) |
| torch_npu/distributed | ✅ liujunzhu, wjlflyer, kisnwang, chengpeng25 (4/2) | ✅ liujunzhu, wjlflyer (2/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
zhenyu10, thanks for your pull request. All authors of the commits have signed the CLA. 👍


4月28日 添加了label:ascend-cla/yes
ascend-robot
4月28日 评论:
4月28日 评论:
当前仓库存在以下 保护分支 :
| Protected Branch | Version | Release |
|---|---|---|
| master | ||
| v2.7.1 | ||
| v2.9.0 | ||
| v2.10.0 | ||
| v2.11.0 | ||
| v2.8.0 |
评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支(创建同步 PR):
a) 如果当前 PR 是 Open 状态,同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged,将立即执行同步操作
注意:
- /sync 命令可以指定同步到多个分支,仅最后一个 /sync 命令生效
- 如果创建的同步 PR 不正确,可通过向同步 PR 的源分支提交轻量级 PR 完善,或使用 /close 命令关闭


此处折叠了396条消息 查看更多
liujunzhu
5月25日 评论:
5月25日 评论:
/approve


5月25日 添加了label:approvedlgtm
ascend-robot
5月25日 评论:
5月25日 评论:
5月25日 合入了pull request
ascend-robot
5月28日 评论:
5月28日 评论:
The repo or branch is not access to PR-cooperate, please check the current repo https://gitcode.com/Ascend/pytorch.git, branch=v2.11.0

