已合并
feat: add input validation and test cases to all_to_all_vc #44783
limuan创建于 6 天前
feat: add input validation and test cases to all_to_all_vc #44783
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 limuan 的贡献）
atomgit-bot
6 天前 评论：
6 天前 评论：
变更摘要
本 PR 为 all_to_all_vc 新增完整的输入校验逻辑与配套测试用例。all_to_all_vc 是基于 HCCL HcclAlltoAllVC 的可变长度 all-to-all 集合通信接口，由全局 [rankSize][rankSize] 的 send_count_matrix（行主序，matrix[i][j] 表示 rank i 发送给 rank j 的元素数）驱动，无需 displs，数据按目标/源 rank 连续排布。该 PR 打通了从 Python 层到 HCCL 层的完整调用链，并新增了 6 类参数校验与正/负向测试，重点保障非法输入（矩阵尺寸不符、负计数、dtype 不匹配、0 维张量、发送/接收缓冲区过小）在进入集合通信前同步报错、避免挂起。
主要改动
- 新增
all_to_all_vc全链路实现：torch_npu/distributed/distributed_c10d.py新增_all_to_all_vcPython 封装（将矩阵展平为 1-D int 列表并调用group.alltoallvc），torch_npu/distributed/__init__.py以all_to_all_vc别名导出；ProcessGroupHCCL.cpp实现_alltoallvc_inner（核心集合通信逻辑）与_alltoallvc（经npu_custom_dist::wrap_alltoallvc_inner分发的入口，便于图模式捕获），并在ProcessGroupHCCL.hpp声明。 - 新增输入参数校验：
_alltoallvc_inner校验矩阵尺寸必须等于rankSize*rankSize、所有计数非负（防止转uint64时负数静默回绕）、输入输出 dtype 一致、不支持 0 维标量张量、input.numel()不小于发送行和、output.numel()不小于接收列和，并输出明确的错误信息。 - 打通 HCCL 符号加载与算子注册：
HcclCompile.h新增HcclAlltoAllVC符号加载与hcclAlltoAllVC封装；HcclOps.cpp注册wrap_alltoallvc_innerop 及TORCH_LIBRARY_IMPL实现；Init.cpp为ProcessGroupHCCL绑定alltoallvc方法（支持opts与 GIL 释放）。 - 新增测试用例：
test/distributed/test_all_to_all_vc.py新增正/负向用例，正向覆盖 2/4 rank、{float32, int32, float16}与同步/异步模式的正确性（值编码为i*10+j），负向覆盖发送/接收缓冲区过小、负计数、dtype 不匹配、矩阵尺寸错误、标量张量 6 类校验场景，并约定每个 rank 传入相同非法输入以保证同步报错不挂起。
ascend-robot
6 天前 评论：
6 天前 评论：
atomgit-bot
6 天前 评论：
6 天前 评论：
6 天前 添加了label：stat/needs-squash
此处折叠了136条消息 查看更多
pengjingyou
5 天前 评论：
5 天前 评论：
/lgtm
5 天前 添加了label：approvedlgtm
5 天前 合入了pull request
ascend-robot
5 天前 评论：
5 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#14226 [ commitID：ac6917e5 ] 已完成
【合入来源】
【修改方案】
为 torch_npu 新增
all_to_all_vc（variable-count all-to-all）集合通信算子，底层 backed by HCCLHcclAlltoAllVC，对应需求"ProcessGroupHCCL 模块接入 HcclAlltoAllVC 算子；python 层 distributed_c10d.py 封装 all_to_all_vc 接口对外使用"。与已有的
alltoall/alltoall_base（走HcclAlltoAllV，带 displs）不同，all_to_all_vc由全局[rankSize][rankSize]的sendCountMatrix驱动：matrix[i][j]= rank i 发给 rank j 的元素个数。每个 rank 传同一份矩阵，发送自己的行、接收自己的列；send/recv buffer 按 rank 连续摆放，无需 displs。对外路径：用户
from torch_npu.distributed import all_to_all_vc, torch_npu.distributed.all_to_all_vc。【资料变更】
资料需更新,增加torch_npu.distributed.all_to_all_vc(output, input, matrix) 接口描述:
已有issue跟踪:https://gitcode.com/Ascend/pytorch/issues/4185
【接口变更】
新增接口:torch_npu.distributed.all_to_all_vc(output, input, matrix)
【功能验证】
【CheckList】