Skip to content

[Auto Parallel] Add co_shard spmd_rule for matmul - #75095

Merged
From00 merged 11 commits into
PaddlePaddle:developfrom
ooooo-create:co_shard_for_matmul_
Sep 22, 2025
Merged

[Auto Parallel] Add co_shard spmd_rule for matmul#75095
From00 merged 11 commits into
PaddlePaddle:developfrom
ooooo-create:co_shard_for_matmul_

Conversation

@ooooo-create

@ooooo-create ooooo-create commented Sep 4, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator

PR Category

Auto Parallel

PR Types

New features

Description

elementwise 相关逻辑在 #75036 中

Add co_shard spmd_rule for matual

  1. Matmul 的 SPMD 规则，按照 batch 维度，非规约维度和规约维度三种，先对 batch 维度进行合并，然后单独处理解析出来的左非规约 和 右非规约，如果被同一个 mesh dim 切分，保留维度大的切分，之后处理两边的规约轴。
  2. 适配 matmul_grad 隐藏 bug，ReduceGradBroadCastDims 用来处理被广播轴，被切分就属于 partial_dims，但是只处理了增加维度的广播，遗漏了大小为 1 时候的 batch 维度广播，添加处理逻辑
  3. 修改部分单测的 process_mesh 的 shape 设置，比如 32 不能乘除 3
  4. GetAxesSizes 现在会判断张量维度是否能进行广播，CI 上报错 k 轴不相同，所以修改了 test/auto_parallel/static_reshard_api_cross_mesh.py，简单看了一下，之前没有报错，是因为这个检查是在 kernel 中，而不在 infer_meta 中
    • 现在 GetAxesSizes 使用了 with_broadcast 会判断同一个 axes 对应的能不能广播，并获取最终广播之后的维度（TODO: 矩阵乘中有些位置是要相同的，比如 k 轴，后面拓展一下函数）
    • 如果 with_broadcast=False，会选择到维度最小作为 axes size
    • 如果是在静态图下，获取size这个逻辑会不会有问题，可能是未知的 -1

@paddle-bot

paddle-bot Bot commented Sep 4, 2025

Copy link
Copy Markdown

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

@paddle-bot paddle-bot Bot added the contributor External developers label Sep 4, 2025
@ooooo-create ooooo-create changed the title [Auto Parallel] Add co_shard spmd_rule for matual [Auto Parallel] Add co_shard spmd_rule for matmul Sep 5, 2025
@codecov-commenter

codecov-commenter commented Sep 5, 2025 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.72131% with 6 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@25f987d). Learn more about missing BASE report.

Files with missing lines Patch % Lines
paddle/phi/infermeta/spmd_rules/utils.cc 95.38% 6 Missing ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             develop   #75095   +/-   ##
==========================================
  Coverage           ?   96.72%           
==========================================
  Files              ?        2           
  Lines              ?      183           
  Branches           ?        0           
==========================================
  Hits               ?      177           
  Misses             ?        6           
  Partials           ?        0

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ooooo-create
ooooo-create marked this pull request as ready for review September 5, 2025 09:50
@ooooo-create

ooooo-create commented Sep 12, 2025

Copy link
Copy Markdown
Collaborator Author

/re-run all_failed

@ooooo-create

ooooo-create commented Sep 12, 2025

Copy link
Copy Markdown
Collaborator Author

/re-run all-failed

From00
From00 approved these changes Sep 22, 2025
View reviewed changes

@From00 From00 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@From00
From00 merged commit c9be274 into PaddlePaddle:develop Sep 22, 2025
74 of 75 checks passed
@ooooo-create
ooooo-create deleted the co_shard_for_matmul_ branch September 29, 2025 09:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@From00 From00 From00 approved these changes
@LiYuRio LiYuRio Awaiting requested review from LiYuRio LiYuRio is a code owner
@ForFishes ForFishes Awaiting requested review from ForFishes ForFishes is a code owner
@zhiqiu zhiqiu Awaiting requested review from zhiqiu

Assignees

No one assigned

Labels

contributor External developers

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ooooo-create @codecov-commenter @From00