已合并
reduce scatter support tensorlist.size != world_size #44438
limuan创建于 7 天前
reduce scatter support tensorlist.size != world_size #44438
已合并
limuan成员
7 天前
7 天前
# 【合入来源】
- [x] 需求
- [ ] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】
本提案修改 `torch_npu` 的 `ProcessGroupHCCL::reduce_scatter`，使其输入支持范围与 PyTorch 社区（NCCL）保持一致。
PyTorch 2.2 的 `reduce_scatter` 对输入张量列表有形状/计数约束（列表长度须等于 `world_size`、每张量 numel 须等于输出 numel 等），2.3 起去除了这些约束。pta（torch_npu）当前的 `reduce_scatter` 与 PyTorch 2.2 实现一致：输入是 tensor list，包含多个 tensor，数量与卡数一致。这与社区后续版本不一致，需要兼容输入只有一个 tensor 的场景等用例，和社区保持一致。
核心改动：在 `reduce_scatter` 的 same_size 分支新增展平函数 `flatten_for_reduce_scatter`，替代原先复用的 `flatten_for_scatter_gather`——移除"输入张量数须等于 world_size""每张量 numel 须等于输出 numel"两条校验，保留列表长度一致与 input/output 同设备校验，并新增非空检查。Python 层 `torch.distributed.reduce_scatter` 签名不变。）
# 【资料变更】
不涉及
# 【接口变更】
不涉及
# 【功能验证】
1、基础功能验证:执行test_reduce_scatter.py用例
2、dtype*op*tensorshape(1D,2D)*input_tensor_list.size(==,>,< world_size)*all_input_numbel(==, >, < )all_output_numbel reduce_scatter 全耦合输入,reduce_scatter算子输出与gpu结果对比
3、性能验证,修改前后pta调用reduce_scatter,用时基本无变化
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]
- [ ] 代码注释完备，正确记录错误日志
- [ ] 代码实现进行了返回值、空指针等校验
- [ ] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
- [ ] PR持续集成流水线（CI）执行通过，代码检查无异常
Pull Request已成功合入, 合并人@ascend-robot
（感谢 limuan 的贡献）
atomgit-bot
7 天前 评论：
7 天前 评论：
变更摘要
此 PR 为 ProcessGroupHCCL::reduce_scatter 新增了专用的输入张量扁平化函数 flatten_for_reduce_scatter，使得 reduce_scatter 操作支持输入张量列表大小与 world_size 不一致的场景，行为与 NCCL 对齐。同时将原有对 flatten_for_scatter_gather 的调用替换为新的专用函数。
主要改动
- 新增
flatten_for_reduce_scatter函数：该函数以tensor_lists.size()作为设备数量（而非world_size），使 reduce_scatter 的输入张量列表大小不必等于world_size；同时校验输入列表与输出张量列表长度一致、对应的输入/输出张量位于同一设备。 - 替换
reduce_scatter中的扁平化调用：将原先调用flatten_for_scatter_gather(inputTensors, outputTensors, size_)改为调用flatten_for_reduce_scatter(inputTensors, outputTensors)，去除了对size_（即world_size）的依赖。
ascend-robot
7 天前 评论：
7 天前 评论：
atomgit-bot
7 天前 评论：
7 天前 评论：
代码审查
审查总结
已审查的文件：
torch_npu/csrc/distributed/ProcessGroupHCCL.cpp：发现 1 个问题
各优先级问题统计：
- P0：0
- P1：0
- P2：1
- P3：0
总体风险评估：此变更引入了专用函数 flatten_for_reduce_scatter 以支持 reduce_scatter 操作中 tensor list 大小与世界大小不同的场景，整体设计合理。唯一需要关注的是新函数中缺少对内部 vector 为空的防御性检查，可能导致未定义行为，建议在合入前修复。其余变更逻辑正确，无安全、性能或兼容性问题。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 1 |
💬 仅评论
7 天前 添加了label：ascend-cla/yes
此处折叠了79条消息 查看更多
6 天前 添加了label：ci-pipeline-passed
ascend-robot
6 天前 评论：
6 天前 评论：
流水线 PR-pipeline_pytorch#60515 [ commitID：6b087be2 ] 已完成
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | ✅ | >>> |
| Build_ARM | ✅ | >>> | |
| Build_LibTorch_x86 | ✅ | >>> | |
| Build_LibTorch_ARM | ✅ | >>> | |
| Build_X86_torchair | 🛑 | >>> | |
| Build_ARM_torchair | 🛑 | >>> | |
| patch_test | 🛑 | >>> | |
| 恶意代码检查 | Antipoison | ✅ | >>> |
| 编码安全与规范检查 | codecheck_pre-commit | ✅ | >>> |
| check_error | ✅ | >>> | |
| lintrunner | ✅ | >>> | |
| 开源片段检查 | SCA | ✅ | >>> |
| 开发者测试 | UT_ARM_A3_Part_01 | 🛑 | >>> |
| UT_ARM_A3_Part_02 | 🛑 | >>> | |
| UT_ARM_A2_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Part_02 | ✅ | >>> | |
| UT_ARM_A2_Part_03 | ✅ | >>> | |
| UT_inductor_Part_01 | 🛑 | >>> | |
| UT_inductor_Part_02 | 🛑 | >>> | |
| UT_inductor_Part_03 | 🛑 | >>> | |
| UT_inductor_Part_04 | 🛑 | >>> | |
| UT_DIST_ARM_Part_01 | ✅ | >>> | |
| UT_DIST_ARM_Part_02 | ✅ | >>> | |
| UT_DIST_ARM_Part_03 | ✅ | >>> | |
| UT_DIST_ARM_Part_04 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_02 | ✅ | >>> | |
| 流水线 | PR-pipeline_pytorch | ✅ | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
6 天前 合入了pull request
ascend-robot
6 天前 评论：
6 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
6 天前 评论：
6 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#13957 [ commitID：6b087be2 ] 已完成