已合并

reduce scatter support tensorlist.size != world_size #44436

limuan创建于 14 天前从limuan/pytorch： reduce_scatter_v2.12.0 合入到 Ascend/pytorch： v2.12.0

reduce scatter support tensorlist.size != world_size #44436

已合并

limuan创建于 14 天前

从limuan/pytorch reduce_scatter_v2.12.0合入到Ascend/pytorch v2.12.0

limuan成员

14 天前

【合入来源】

需求
问题单
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 签名不变。）
各输入形态数值示例

统一 world_size=4、output=[4]、fp32、SUM，输入约定同 §2.1（rank r 展平后第 k 元 = r*10+k）。need = 16（输出 numel × world_size）；全局位置 p 的归约值在 p < have 时为 60+4p，p ≥ have 时零填充。下文列出每个 rank 的输入与输出。

① 单 tensor（input_list = [[16]]，have = 16 = need）：

rank0 IN : [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]    OUT: [60,64,68,72]
rank1 IN : [10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25]   OUT: [76,80,84,88]
rank2 IN : [20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35]   OUT: [92,96,100,104]
rank3 IN : [30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45]   OUT: [108,112,116,120]

② 张量数 < 卡数（input_list = [[4],[4]]，have = 8 < 16，尾部两 rank 零填充）：

rank0 IN:           rank1 IN:           rank2 IN:           rank3 IN:
  [0, 1, 2, 3]        [10,11,12,13]       [20,21,22,23]       [30,31,32,33]
  [4, 5, 6, 7]        [14,15,16,17]       [24,25,26,27]       [34,35,36,37]
  OUT: [60,64,68,72]  OUT: [76,80,84,88]  OUT: [0,0,0,0]      OUT: [0,0,0,0]

③ 张量数 > 卡数（input_list = [[4]]*5，have = 20 > 16，第 5 个张量整体忽略）：

rank0 IN:           rank1 IN:           rank2 IN:           rank3 IN:
  [0, 1, 2, 3]        [10,11,12,13]       [20,21,22,23]       [30,31,32,33]
  [4, 5, 6, 7]        [14,15,16,17]       [24,25,26,27]       [34,35,36,37]
  [8, 9,10,11]        [18,19,20,21]       [28,29,30,31]       [38,39,40,41]
  [12,13,14,15]       [22,23,24,25]       [32,33,34,35]       [42,43,44,45]
  [16,17,18,19]       [26,27,28,29]       [36,37,38,39]       [46,47,48,49]   ← 忽略
  OUT: [60,64,68,72]  OUT: [76,80,84,88]  OUT: [92,96,100,104] OUT: [108,112,116,120]

④ 张量数 = 卡数，每张量 > 输出（input_list = [[5]]*4，have = 20 > 16，展平后尾部 4 元忽略）：

rank0 IN:             rank1 IN:             rank2 IN:             rank3 IN:
  [0, 1, 2, 3, 4]       [10,11,12,13,14]       [20,21,22,23,24]       [30,31,32,33,34]
  [5, 6, 7, 8, 9]       [15,16,17,18,19]       [25,26,27,28,29]       [35,36,37,38,39]
  [10,11,12,13,14]      [20,21,22,23,24]       [30,31,32,33,34]       [40,41,42,43,44]
  [15,16,17,18,19]      [25,26,27,28,29]       [35,36,37,38,39]       [45,46,47,48,49]   ← 后4元忽略
  OUT: [60,64,68,72]    OUT: [76,80,84,88]     OUT: [92,96,100,104]   OUT: [108,112,116,120]

展平后每 rank 共 20 元，前 16 元参与归约，尾部 4 元（第 4 个张量的后 4 元）忽略。

⑤ 张量数 = 卡数，每张量 < 输出（input_list = [[3]]*4，have = 12 < 16，尾部 rank 零填充）：

rank0 IN:           rank1 IN:           rank2 IN:           rank3 IN:
  [0, 1, 2]          [10,11,12]          [20,21,22]          [30,31,32]
  [3, 4, 5]          [13,14,15]          [23,24,25]          [33,34,35]
  [6, 7, 8]          [16,17,18]          [26,27,28]          [36,37,38]
  [9,10,11]          [19,20,21]          [29,30,31]          [39,40,41]
  OUT: [60,64,68,72] OUT: [76,80,84,88]  OUT: [92,96,100,104] OUT: [0,0,0,0]

⑥ 2D 及多维（input_list = [[5,4]]，have = 20 > 16，展平后尾部 4 元忽略，同 ③）：

rank0 IN (shape [5,4]):  rank1 IN (shape [5,4]):  rank2 IN (shape [5,4]):  rank3 IN (shape [5,4]):
  [0, 1, 2, 3]           [10,11,12,13]           [20,21,22,23]           [30,31,32,33]
  [4, 5, 6, 7]           [14,15,16,17]           [24,25,26,27]           [34,35,36,37]
  [8, 9,10,11]           [18,19,20,21]           [28,29,30,31]           [38,39,40,41]
  [12,13,14,15]          [22,23,24,25]          [32,33,34,35]          [42,43,44,45]
  [16,17,18,19]          [26,27,28,29]          [36,37,38,39]          [46,47,48,49]   ← 忽略
  OUT: [60,64,68,72]     OUT: [76,80,84,88]      OUT: [92,96,100,104]    OUT: [108,112,116,120]

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

1、基础功能验证
2、dtypeoptensorshape(1D,2D)*input_tensor_list.size(==,>,< world_size)*all_input_numbel(==, >, < )all_output_numbel reduce_scatter算子输出与gpu结果对比

3、性能验证,修改前后pta调用reduce_scatter,用时基本无变化

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

代码注释完备，正确记录错误日志
代码实现进行了返回值、空指针等校验
PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
PR持续集成流水线（CI）执行通过，代码检查无异常

Pull Request已成功合入, 合并人@ascend-robot

（感谢 limuan 的贡献）

limuan成员

14 天前创建了 pull request，commit 785df6f7

ascend-robot成员

14 天前评论：

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=44436&platform=gitcode&codeHostingPlatformFlag=gitcode

atomgit-bot

14 天前评论：

变更摘要

本次变更旨在使 reduce_scatter 操作支持 tensor_lists.size() != world_size 的场景，对齐 NCCL 的行为。通过新增专用的 flatten_for_reduce_scatter 函数替代原先复用 flatten_for_scatter_gather 的方式，移除了对 size_（即 world_size）的依赖，仅在函数内部校验输入列表与输出列表长度一致以及对应张量位于同一设备。

主要改动

新增 flatten_for_reduce_scatter 函数: 在 ProcessGroupHCCL.cpp 中新增该函数，用于 reduce_scatter 的输入展平，不再依赖 size_ 参数，允许输入张量列表大小与 world_size 不同；同时校验输入/输出列表长度一致及设备一致性。
reduce_scatter 方法调用替换: 将 ProcessGroupHCCL::reduce_scatter 中原本调用 flatten_for_scatter_gather(inputTensors, outputTensors, size_) 替换为 flatten_for_reduce_scatter(inputTensors, outputTensors)，解除对 world_size 的强制约束。
c10d::newLikeFlat 调用适配: 新函数内部使用 c10d::newLikeFlat(tensor_lists, i) 对每个设备上的张量列表进行展平，替代原有展平逻辑中对 world_size 的遍历限制。

atomgit-bot

14 天前评论：

代码审查

✅ 未发现问题

ascend-robot成员

14 天前添加了label：ascend-cla/yes

此处折叠了76条消息查看更多

ascend-robot成员

13 天前删除了label：ci-pipeline-running

ascend-robot成员

13 天前添加了label：ci-pipeline-passed

ascend-robot成员

13 天前评论：

流水线 PR-pipeline_pytorch#60512 [ commitID：6cacec6d ] 已完成

>>>代码风格自动修复执行成功（无修复内容）

阶段	任务名	状态	详情
编译构建	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	✅	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, huangjingwei, liangsongwei, yashi999, culechan, Dring, wuyouqi1, L1919_snow, qq_52711437, WhiteNight12, nomiz, xiu_21, ffmh, wanglijun55, hss-shuai, husichao, smallsilly, lanshaozuishuai, jimmyisme1, lzy0920232, alpha-junh, Sunshine_Youngster, wei_zhuoyi, zhangyihuiben, zyw-hw, zzzkeke, rmch, yangch0324, LucciC, AACAES, renyujin, wjlflyer, senzhen-town, pengjingyou, qsc97, limuan, yule100, xiaoqi-zhou, kuhn7, chenxingying, hanye02, zichun_ye, anyrenwei, kkjocker, wangzili121, Lu_G, yvjc, puddingfjz, HandsoemLemon, bigprestigee1, huawuyi, zhenyu10, dairenjie, du-jin-hang, zou-jieyu, adelaideliu, TrHan, wanlinan, Windwindzzz, pengqihw, kisnwang, yuheng_wang, honghao_wang, jizewei, zhangguoguang, sunyu-xuan, chenrayray, hbhu_bin, liujunzhu, c_34, LiNuoh, maoyuanpeng1, zzhongmin, zhaoyu65, bellatan, jiabaolin, zhuofanshen, wencaiwen, lu_zhuge, caoshuyang, molly12, lyx324521, LQ1206, gitcode-chenjiao, cai-weiwei1989, CHDong, ogqin, yuanlipingGit, xuqinglin1, lqz2, zouwei1, chaoluoa, paradox325, jackzhang1116, yaoyao, akh, yujiacheng, dengjie0116, Hubert11111, Shine_Ws, wslhj555, longqiand, OYtao666, JiaqingQiang, luyyyy, Kingbelial, zhanghaiyu0101, wenxp1018, yanliu-luoluo, ksun_sekiro, liyong328, wgzheng, tangky, vivi_is_coding, aoiaoisola, weixin_44494597, wangmengmengwang65667, hid57809721, qq_35468730, comeonup, C547032, gcw_m5OQChA4, yao_yao_ling_xian, cnnbwcy, szqfes_12, cora_19, cann_lilin, can, shawnylee233, fanglanyue0916, hhz0, LiNuohang, taohuoquan, Jesse, WSs_321]评论有效

compile、compile_inductor、compile_torchair : 运行流水线
retry : 重试流水线所有失败子任务
retry <任务名> : 仅重试指定失败子任务
stop : 停止流水线

ascend-robot成员

13 天前合入了pull request

ascend-robot成员

13 天前评论：

流水线 pytorch_gitcode_PR_multiVersion#13959 [ commitID：6cacec6d ] 已完成