Pull Request已成功合入, 合并人@ascend-robot
（感谢 关龙锋 的贡献）
变更摘要
此 PR 主要修复了 all_to_all_single 在 eager 模式下对多维张量的支持问题。原先的实现直接将 send_counts / recv_counts 作为 input_split_sizes / output_split_sizes 传递给 c10d.all_to_all_single，未考虑张量的非零维度元素数量，导致多维张量场景下 split sizes 计算错误。修复后，代码会根据输入张量的维度推导出 input_non_0_dim_size，并据此反向计算出正确的 split sizes，同时新增了对应的 eager 模式测试用例。
主要改动
- 修复
hcom_alltoall.py中 split sizes 的计算逻辑：移除了原本仅在recv_counts不为None时计算out_size的条件分支，改为无条件取input_tensor.size()，并引入input_non_0_dim_size（根据out_size[0]是否为零、张量是否为一维等情况分别计算），然后用send_counts/recv_counts除以input_non_0_dim_size得到正确的input_split_sizes/output_split_sizes。 - 修正
out_size[0]的取值：out_size[0]不再直接取sum(recv_counts)，而是取sum(output_split_sizes)，使其与计算后的 split sizes 保持一致。 - 新增
_test_all_to_all_single_compile_eager测试方法：在hcom_st.py中增加了针对一维张量在 eager 编译模式下的all_to_all_single正确性测试。 - 新增
_test_all_to_all_single_mulit_dims_eager测试方法：在hcom_st.py中增加了针对多维张量（形状(4, 2)）在 eager 编译模式下的all_to_all_single正确性测试，并与标准的dist.all_to_all_single结果进行一致性比对。
代码审查
审查结论
审查覆盖
| 文件 | 审查结果 |
|---|---|
python/torchair/_ge_concrete_graph/ge_converter/experimental/hcom_alltoall.py |
发现 2 个问题 (P1, P2) |
tests/smoke/hcom_st.py |
发现 1 个问题 (P3) |
问题统计
- P1: 1 个 —
input_non_0_dim_size为 0 时除零崩溃 (回归缺陷) - P2: 1 个 —
npu_all_to_all_single_npu与npu_all_to_all_single_npu_meta对零维输入的 shape 推断不一致 - P3: 1 个 — 新增 eager 测试的
dynamic参数覆盖不对称 (可选改进)
整体风险评估
中等风险。本 PR 修复了 eager 模式下 all_to_all_single 对多维 tensor 的语义错误（send_counts/recv_counts 被错误地当作 dim0 维度的分片大小传递给 c10d.all_to_all_single），核心修复逻辑正确。但新引入的除法操作缺少对 input_non_0_dim_size == 0 的保护，当输入 tensor 存在零值维度时会导致 ZeroDivisionError 崩溃，这是一个旧代码不存在的回归。此外，实际执行函数与 meta shape 推断函数在零维边界条件下的行为不一致，可能在 torch.compile 动态 shape 场景引发隐蔽的 shape mismatch 问题。建议在合入前至少修复 P1 除零问题。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 1 |
| 🟡 建议 | 1 |
⛔ 需要修改
/lgtm
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
bug fix all_to_all_single support mulit dim in eager model：
原：
out_size[0] = sum(recv_counts)
out_tensor = input_tensor.new_empty(out_size)
其中recv_counts为接受的元素个数
而实际：c10d.all_to_all_single接受的是output_split_sizes是对input[0]的切分比例，用来表示每个卡返回的元素比例是多少。而不是要返回的元素个数
就是下面举例的，入参应该是[2, 5, 6, 3]，而不是[24, 60, 72, 36] 元素个数。
这样入参1、会导致错误，output_split_sizes和input[0] 个数不匹配报错，input[0] 总共是16个。
c10d.all_to_all_single(
out_tensor,
input_tensor,
output_split_sizes=recv_counts,
input_split_sizes=send_counts,
group=pg,
async_op=False,
)
示例条件
plaintext
input_tensor shape = [16, 4, 3]
group_size = 4 # 4张卡rank0~rank3
output_split_sizes = [2, 5, 6, 3]
步骤 1：计算 input_non_0_dim_size
plaintext
out_size[0] = 16
total_elem = 16 * 4 * 3 = 192
input_non_0_dim_size = 192 // 16 = 12
等价 D1D2 = 43 = 12
含义：第 0 维取 1 行时，内部固定 12 个元素。
步骤 2：逐 rank 计算 recv_counts
公式：recv_counts[i] = output_split_sizes[i] * 12
rank0：2 × 12 = 24
rank1：5 × 12 = 60
rank2：6 × 12 = 72
rank3：3 × 12 = 36
最终
plaintext
recv_counts = [24, 60, 72, 36]
含义：all_to_all 通信时，每张卡接收的元素总量。