已合并
migrate distributed ops to the new dispatch flow #736
migrate distributed ops to the new dispatch flow #736
已合并
rongyue创建于 6月1日
rongyue
rongyue
6月1日

What type of PR is this?

/kind refactor


What does this PR do / why do we need it:
将多个分布式算子实现迁移到新的集中式三阶段派发流程,同时在派发器中集中处理 aclop 打包参数的标准化。

改动内容

  1. 派发器层面重构(_op_dispatch.py)
  • aclop 标准化集中处理:_normalize_aclop_args 现在在上游 _dispatch_layout_infer 中统一调用,解包后的结果通过 _packed_call 关键字参数向下传递。消除了 _with_layout_infer、_with_layout_infer_with_tuple_expand 和 _with_layout_infer_with_shape 中重复的解包逻辑。
  • 提取 _call_op_impl 静态方法:统一了算子调用路径(同时处理 aclop 打包调用和普通调用),减少代码重复。
  • 将 multinomial 加入通用随机派发集合:torch.multinomial 现在通过通用 RNG 路径派发,不再需要独立的 parallel_multinomial.py。
  1. 算子层面迁移(11 个算子重构)
  • 每个算子现在遵循统一的三阶段模式:preprocess 负责将 构建缓存数据,infer_layout从缓存数据推导输出布局,get_expand_impl 可选地包装算子调用以注入分布式逻辑。涉及算子: parallel_cumsum.py, parallel_elementwise.py, parallel_embedding.py, parallel_expand_dims.py, parallel_flatten.py, parallel_gather.py, parallel_repeat_interleave.py, parallel_reshape.py, parallel_slice.py, parallel_squeeze.py, parallel_tuple_elementwise.py
  1. 移除的算子
  • parallel_multinomial.py 及其 YAML 配置 —— 已迁移至通用 RNG 派发,不再需要专用 DistributedOp。
  1. 测试更新
  • 对应 12 个 UT 文件全部更新为测试新的三阶段协议。
  • 移除 test_parallel_multinomial.py。
  • 更新 test_op_dispatch.py 以覆盖派发器变更。

Which issue(s) this PR fixes:

Fixes #94


Test Plan and Test result:What scenarios were tested, and what were the verification results(Function, performance, reliability, etc.):
测试场景

  1. 验证集中式 _normalize_aclop_args 正确解包。
  2. 验证 multinomial 在移除专用算子后正确通过通用 RNG 路径派发。
  3. 针对每个迁移算子,验证 preprocess 正确将 DTensor 输入转为 local tensor 并缓存布局和元数据、infer_layout(cache_values) 推导出的输出布局与迁移前一致、get_expand_impl 返回正确的扩展包装函数(或在不需扩展时返回 None)。

测试结果

  • 相关 UT 用例:全部通过。
  • 相关 ST 用例:已有 ST 覆盖的算子均通过。其中 parallel_slice 和 parallel_tuple_elementwise 当前无对应 ST 用例,仅验证了 UT。

Self-checklist:(请自检,在[ ]内打上x,我们将检视你的完成情况,否则会导致pr无法合入)

likedislike
Pull Request已成功合入, 合并人@liuchongming74
(感谢 rongyue 的贡献)
rongyuerongyue
6月1日 创建了 pull request,commit 4fa35dea
rongyuerongyue
6月1日 关联了issue:【RFC】分布式算子类迁移指导
MindSpore-BotMindSpore-Bot成员
6月1日 添加了label:mindspore-cla/yes
司小南(机器人)司小南(机器人)成员
6月1日 添加了label:pr-check-pass
MindSpore-BotMindSpore-Bot成员
6月1日 添加了label:no-pass-all-review
此处折叠了267条消息 查看更多
liuchongming74liuchongming74成员
6月22日 通过审查
MindSpore-BotMindSpore-Bot成员
6月22日 删除了label:no-pass-all-review
MindSpore-Bot
MindSpore-Bot成员
6月22日 评论:

The ci-pipeline-passed label is expired. Please retest again.

likedislike
MindSpore-Bot
MindSpore-Bot成员
6月22日 评论:

The ci-pipeline-passed label is expired. Please retest again.

likedislike
liuchongming74liuchongming74成员
6月22日 合入了pull request,合并节点 SHA:2455275a28b18352dfa0e1d266bf078cab30b8ee