已合并
reduce scatter support tensorlist.size != world_size #44440
limuan创建于 8 天前
reduce scatter support tensorlist.size != world_size #44440
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 limuan 的贡献）
ascend-robot
8 天前 评论：
8 天前 评论：
atomgit-bot
8 天前 评论：
8 天前 评论：
变更摘要
此 PR 为 ProcessGroupHCCL::reduce_scatter 新增了专用的 flatten 函数 flatten_for_reduce_scatter，使其支持输入 tensor list 大小与 world_size 不一致的场景，对齐 NCCL 的行为。同时增加了输入/输出 tensor 之间的设备一致性校验。
主要改动
- 新增
flatten_for_reduce_scatter函数：与原有的flatten_for_scatter_gather不同，新函数不再依赖size_（world_size），而是以tensor_lists.size()作为设备数量进行 flatten，从而支持 tensor list 大小不等于 world_size 的情况。 - 修改
reduce_scatter中的 flatten 调用：将same_size分支中的flatten_for_scatter_gather(inputTensors, outputTensors, size_)替换为flatten_for_reduce_scatter(inputTensors, outputTensors)，不再传入size_参数。 - 增加输入/输出设备一致性检查：新函数中校验了
tensor_lists[i].front().get_device()与outputTensors[i].get_device()是否一致，不一致时抛出错误。 - 增加输入/输出长度一致性检查：新函数中校验了
tensor_lists.size()与outputTensors.size()是否相等，不相等时抛出错误。
atomgit-bot
8 天前 评论：
8 天前 评论：
代码审查
审查总结
已审查变更文件：
torch_npu/csrc/distributed/ProcessGroupHCCL.cpp— 3 个发现
| 优先级 | 数量 |
|---|---|
| P0 | 0 |
| P1 | 0 |
| P2 | 2 |
| P3 | 1 |
整体风险评估：中等。新增的 flatten_for_reduce_scatter 函数在核心逻辑上是合理的（去掉了 world_size 约束以对齐 NCCL 行为），但遗漏了对空内层列表的防御性检查（.front() UB），且下游拷贝循环仍以 inputTensors[0].size() 作为所有设备的循环上界，与本次变更允许不同内层列表大小的语义不一致。建议在合入前修复这两个 P2 问题。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 2 |
💬 仅评论
8 天前 添加了label：ascend-cla/yes
此处折叠了190条消息 查看更多
6 天前 添加了label：approvedlgtm
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#14005 [ commitID：e730aa3e ] 已完成
【合入来源】
【修改方案】
本提案修改
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、dtypeoptensorshape(1D,2D)*input_tensor_list.size(==,>,< world_size)*all_input_numbel(==, >, < )all_output_numbel reduce_scatter 全耦合输入,reduce_scatter算子输出与gpu结果对比
3、性能验证,修改前后pta调用reduce_scatter,用时基本无变化
【CheckList】