Merged
feat: add distributed ops AllGatherMatmul and MatmulReduceScatter (MC2) #656
feat: add distributed ops AllGatherMatmul and MatmulReduceScatter (MC2) #656
Merged
hedongdongAdd on May 18
hedongdongCollaborators
May 18

What type of PR is this?

/kind feature


What does this PR do / why do we need it:

为 HyperParallel 接入 Matrix Computation & Communication通算融合算子的分布式调度逻辑,包含:

  • mindspore.ops.all_gather_matmul(AllGatherMatmul)
  • mindspore.ops.matmul_reduce_scatter(MatmulReduceScatter)

CANN MC2 算子在内核内部自行完成 AllGather / ReduceScatter 通信,不需要 HyperParallel 插入额外通信原语。HyperParallel 的作用是:

  1. DTensor 输入解包:通过 preprocess() 的 to_local() 提取本地 tensor 后再传给 CANN。
  2. 输出 DTensor Layout 推断:通过 infer_layout() 基于输入分片方式推断输出 tensor 的分布状态,使下游算子能正确组合分布式计算图。

AllGatherMatmul

参数 说明
x1 (m, k) m 维被 CANN 内部 AllGather 消耗,输出 m 变为全局(Replicate)
x1 k 维 支持 Replicate 或 Shard;k 被切分时输出带 Partial(sum),调用方需 AllReduce
x2 (k, n) 或 (n, k) trans_x2=False 时 k 在 dim 0,trans_x2=True 时 k 在 dim 1;k 维须与 x1 k 维 placement 匹配
gather_output True 时额外返回 AllGather 后的 x1 全局张量;False 时 CANN 返回空 tensor

输出 layout 推断:

  • output dim 0 (m):始终 Replicate(-1),因 AllGather 消耗了 m 维分片
  • output dim 1 (n):继承 x2 的 n 维 placement
  • k 被切分时,output 带 Partial(sum)(与 LinearDistributedOp contract_dim 分片语义一致)
  • gather_output=False 时 gather_out layout 强制全 Replicate,避免空 tensor 的维度越界

MatmulReduceScatter

参数 说明
x1 (m, k) k 维须 Shard(TP),m 维可为 Replicate(纯 TP)或 Shard(DP)
x2 (k, n) 或 (n, k) k 维须与 x1 k 维 layout 一致;n 维可任意

输出 layout 推断:

  • output dim 0 (m):ReduceScatter 将 k 的 TP 分片转化为 m 分片;若 x1 m 维有 DP 分片则联合分片(tuple tensor_map)
  • output dim 1 (n):继承 x2 的 n 维 placement
  • 无 Partial 状态:CANN ReduceScatter 已在内部完成 k 方向的 sum reduce + m 方向 scatter

测试覆盖

测试类型 文件 内容
UT test_parallel_all_gather_matmul.py (15 cases) normalize_args / YAML 注册 / infer_layout 正例(1D+2D mesh、trans_x2、k 分片带 Partial)/ 错误处理(Partial 输入、k-placement 不匹配、m 维多 mesh 联合分片)
UT test_parallel_matmul_reduce_scatter.py (12 cases) normalize_args / YAML 注册 / infer_layout 正例 / 错误处理
ST all_gather_matmul_shard_in_python.py (7 cases) x1 Shard(0)+x2 Replicate / x1 Shard(0)+x2 Shard(1) / trans_x2=True / gather_output=False / 2D (dp,tp) mesh / 4 卡 k 切分 (mp,tp) / 8 卡 mnk 全切 (mp,np,tp)
ST matmul_reduce_scatter_shard_in_python.py (5 cases) tp_basic / trans_x2=True / large_m / 2D (dp,tp) mesh / 8 卡 mnk 全切 (mp,np,tp)

Which issue(s) this PR fixes:

Fixes #152


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

UT(单元测试)

  • tests/ut/core/shard/ops/test_parallel_all_gather_matmul.py — 15 cases,覆盖参数归一化、YAML 注册、infer_layout(replicate/shard/trans_x2/2D-mesh/k 分片带 Partial)、错误处理(Partial 输入、k-placement 不匹配、m 维多 mesh 联合分片)
  • tests/ut/core/shard/ops/test_parallel_matmul_reduce_scatter.py — 12 cases,覆盖参数归一化、YAML 注册、infer_layout(TP/DP+TP/n-shard/trans_x2/联合 tuple m 分片)、错误处理(Partial 输入、k Replicate、k layout 不一致)

ST(系统测试 - distributed, MindSpore)

  • tests/mindspore/st/shard/ops/test_all_gather_matmul_shard_in_python.py — 3 组:group1(4×2 卡)、group2(2×4 卡,含 k 分片场景)、group3(1×8 卡,mnk 三轴全切)
  • tests/mindspore/st/shard/ops/test_matmul_reduce_scatter_shard_in_python.py — 3 组:group1(3×2 卡)、group2(1×4 卡)、group3(1×8 卡,mnk 三轴全切)

验证结果

  • 所有 27 条 UT 通过(mock 平台,无 GPU/NPU 依赖)
  • 所有 AGM group1/group2 ST 在 Ascend NPU 上通过(float16 精度,atol/rtol 满足)
  • k 分片场景使用 float32 numpy 参考值,atol=0.1/rtol=0.02(与 MRS 一致,因 CANN 跨 rank partial sum 累加顺序与单次全 k matmul 路径不同)
  • pylint 通过

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

likedislike
Pull Request has been successfully merged by MindSpore-Bot.
(Thank you for hedongdong's contribution.)
HhedongdongCollaborators
May 18 Created pull request, commit 54554447
MindSpore-BotMindSpore-BotCollaborators
May 18 Added label: mindspore-cla/yes
司小南(机器人)
司小南(机器人)Collaborators
May 18 Comments:

@david-he91, 当前/check-pr未通过,原因如下:

以下Pull Request描述检查项未通过:
存在不符合模板的选项: 设计:AllGatherMatmul / MatmulReduceScatter 分布式 layout 推导逻辑已覆盖所有合法分片组合
存在不符合模板的选项: 测试:UT 25 cases + ST 9 cases,覆盖 normalize / preprocess / infer_layout 及分布式场景
存在不符合模板的选项: 接口:通过 YAML 注册,对上层调用方无破坏性变更
存在不符合模板的选项: 约束文档:trans_x2、gather_output 语义、k 维约束已在 docstring 中注明
部分检查项缺失 请重新使用模板
模板中'Test Plan and Test Result' 信息为空,请补充对应信息。

以下issue检查项未通过:
Pull Request未关联issue

请修改好上述检查错误后,重新使用/check-pr触发检查。

likedislike
司小南(机器人)司小南(机器人)Collaborators
May 18 Added label: pr-check-fail
MindSpore-BotMindSpore-BotCollaborators
May 18 Added label: no-pass-all-review
76 messages are collapsed here。 See more
YyangzhenzhangCollaborators
May 19 Approval passed
MindSpore-Bot
MindSpore-BotCollaborators
May 19 Comments:

Notice

The PR needs 2 assignees to review. 1 does not review. if all are passed, please comment /check-pr to try merge the PR. 😄

likedislike
liuchongming74liuchongming74Collaborators
May 19 Approval passed
MindSpore-BotMindSpore-BotCollaborators
May 19 Removed label: no-pass-all-review
MindSpore-BotMindSpore-BotCollaborators
May 19 Merged pull request, merge node SHA: 51c236bd893d04e7d4325c38d1682b062ab5b78d