Add distributed ops OneHotExt · Pull Request !123 · MindSpore/hyper-parallel

1. Add distributed ops OneHotExt
---
### 相关的Issue

### 原因（目的、解决的问题等）
1. 新增 OneHotExt 算子的分布式支持，使用算子替换流程，原因如下：
OneHotExt 支持 num_classes=-1 时，MindSpore 底层会自动推导 depth。但在分布式场景下，当 indices 发生切分（shard）时，各 rank 仅持有局部数据，若仍按单卡逻辑推导 depth，会出现：depth 推导不一致：各 rank 可能得到不同的 local max + 1，导致 OneHotExt 输出最后一维（depth 维）在不同 rank 上不一致。后续通信失败或结果错误：由于输出 shape 不一致，分布式 allgather/concat 等会出现维度不匹配，或最终聚合结果与单卡不一致。因此需要在分布式层面通过算子替换机制接管 OneHotExt：当 num_classes=-1 且 indices 被切分时，先做 全局 allreduce(max) 得到全局最大值，再推导统一的 depth，确保所有 rank 的 OneHotExt 输出 shape 一致，最终聚合正确。

通信算子 dtype 限制踩坑：MindSpore 通信 allreduce 在当前 Ascend 后端对 int64 输入支持受限（报错提示 InnerCommAllReduce 不支持 Int64），而 OneHotExt 的 indices 在底层要求使用 int64，过程中需要做类型转换。

### 描述（做了什么，变更了什么）
1. 新增 OneHotExt 算子实现：添加 OneHotExtDistributedOp 类。
（1）infer_layout：推导输出 layout，插入 one-hot 新维度，且该维度始终为 replicated（"None" / -1）。
（2）get_expand_impl：在满足条件时（num_classes == -1 且 indices layout 存在分片轴）返回展开实现，替换原生 OneHotExt 执行：
1）在本地计算 local_max = ReduceMax(indices)
2）为规避通信不支持 int64 的限制，将 scalar max 转为 int32（并做 int32 上界检查）
3）对所有参与分片的 mesh 轴逐轴执行 allreduce(max) 得到 global_max
4）depth = global_max + 1，再调用原始 OneHotExt（func(...)）完成 one-hot

2. 新增 OneHotExt 算子注册 (YAML配置文件)。
3. 新增 OneHotExt 算子UT测试用例。
4. 新增 OneHotExt 算子ST测试用例。
5. 修改_op_dispatch.py 的 _with_layout_infer。

### 测试用例（新增、改动、可能影响的功能）
1. 新增 OneHotExt UT 测试用例，覆盖：
（1）数据并行 / 全复制场景下的 infer_layout 推导
（2）多维 indices 的约束校验（实现覆盖范围）
（3）1D indices 下不同 axis 的 layout 推导
（4）axis 取值范围校验
（5）非 indices 输入必须 replicated 的约束验证

2. 新增 OneHotExt ST 测试用例：
（1）test_one_hot_ext_data_parallel_1d_int64_1
功能：验证 OneHotExt 在 1D indices 数据并行场景下的正确性（indices 为 int64）。
配置：
输入 layout：layout("dp",)
输出 layout：layout("dp", "None")
参数：tensor shape=(64,), num_classes=32，indices dtype=int64
mesh：Layout((2, 2, 2), ("dp", "cp", "mp"))
验证：并行输出经 local_to_global 聚合后与 standalone 输出一致。
（2）test_one_hot_ext_data_parallel_2d_int64_2
功能：验证 OneHotExt 在 2D indices（dp 切分首维）场景下正确性（indices 为 int64）。
配置：
输入 layout：layout("dp", "None")
输出 layout：layout("dp", "None", "None")
参数：tensor shape=(16, 32), num_classes=64，indices dtype=int64
mesh：同上
验证：并行聚合后的输出与 standalone 一致。
（3）test_one_hot_ext_replicate_all_int64_3
功能：验证 OneHotExt 在全复制场景下正确性（indices 为 int64）。
配置：
输入 layout：layout("None",)
输出 layout：layout("None", "None")
参数：tensor shape=(64,), num_classes=32，indices dtype=int64
mesh：同上
验证：并行聚合后的输出与 standalone 一致。
（4）test_one_hot_ext_3d_data_parallel_int64_4
功能：验证 OneHotExt 在 3D indices 的 dp-only 场景正确性（indices 为 int64）。
配置：
输入 layout：layout("dp", "None", "None")
输出 layout：layout("dp", "None", "None", "None")
参数：tensor shape=(8, 16, 12), num_classes=32，indices dtype=int64
mesh：同上
验证：并行聚合后的输出与 standalone 一致。
（5）test_one_hot_ext_auto_depth_skewed_distribution_int64_5
功能：验证 num_classes=-1 自动推导在 dp 切分下的正确性，且必须依赖 global allreduce(max) 才能通过。
关键点：刻意构造 skewed distribution，使只有一个 dp shard 持有全局最大值，避免“错误实现（local max）也误通过”。
配置：
输入 layout：layout("dp",)
输出 layout：layout("dp", "None")
参数：tensor shape per-rank=(64,), num_classes=-1，indices dtype=int64
mesh：同上
验证：
所有 rank 的输出 local_shape[-1] 必须等于 global_max + 1；
local_to_global 后的全局输出与 standalone（对齐同一全局 indices）一致。
影响功能：覆盖分布式实现中 get_expand_impl 的替换路径（自动深度推导 + allreduce(max)）。
（6）test_one_hot_ext_auto_depth_all_same_local_max_int64_6
功能：验证 num_classes=-1 在 dp 切分下的“easy case”：各 shard local max 相同，仍应得到正确 depth。
配置：
输入 layout：layout("dp",)
输出 layout：layout("dp", "None")
参数：tensor shape per-rank=(64,), num_classes=-1，indices dtype=int64
mesh：同上
验证：
输出 depth == local_max + 1；
local_to_global 后输出与 standalone 一致。
影响功能：补充覆盖 auto depth 在无需依赖 skewed 的情况下仍正确。

MindSpore / hyper-parallel关闭 .gitee-modal { width: 500px !important; }

搜索帮助

MindSpore / hyper-parallel
关闭