[simple_fsdp][inductor_collectives] rewrite reorder_collectives, sink_waits_iterative #158062

IvanKobzarev · 2025-07-10T20:21:54Z

Stack from ghstack (oldest at bottom):

-> [simple_fsdp][inductor_collectives] rewrite reorder_collectives, sink_waits_iterative #158062

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov

Differential Revision: D78159013

[ghstack-poisoned]

pytorch-bot · 2025-07-10T20:21:58Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158062

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit e42de67 with merge base d5af0ec ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

trunk / linux-jammy-rocm-py3.10 / test (default, 1, 2, linux.rocm.gpu.2) (gh) (similar failure)
test_nn.py::TestNNDeviceTypeCUDA::test_avg_pool_large_tensor2_cuda

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

inductor / cuda12.8-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
torchrec_dlrm
pull / linux-jammy-cuda12.8-py3.10-gcc11-sm89 / test (default, 5, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (trunk failure)
inductor/test_codecache.py::TestFxGraphCache::test_flex_attention_caching

This comment was automatically generated by Dr. CI and updates every 15 minutes.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]

preserving peak memory ghstack-source-id: cc1038b Pull Request resolved: #158062

IvanKobzarev · 2025-07-11T11:19:13Z

@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: 03ea238 Pull Request resolved: #158062

IvanKobzarev · 2025-07-11T14:04:41Z

@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

IvanKobzarev · 2025-07-11T14:41:42Z

@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

ghstack-source-id: 5b8d32d Pull Request resolved: pytorch/pytorch#158062

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

IvanKobzarev · 2025-07-14T17:50:12Z

@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

test/distributed/test_inductor_collectives.py

wconstab · 2025-07-14T19:14:44Z

torch/_inductor/comms.py


    # Dicts to keep track of "next" and "previous" as double-linked structure during grouping
-    _prev: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]] = {}
-    _next: dict[BaseSchedulerNode, Optional[BaseSchedulerNode]] = {}
+    _prev: dict[Optional[BaseSchedulerNode], Optional[BaseSchedulerNode]] = {}


what's the point of storing the 'prev' node of a None? or was this just a workaround for linter complaints?

Yes, only linter complaints.

torch/_inductor/comms.py

wconstab

lgtm. thanks!!

preserving peak memory ghstack-source-id: 30187c8 Pull Request resolved: pytorch/pytorch#158062

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: 0b68cca Pull Request resolved: #158062

IvanKobzarev · 2025-07-15T09:58:31Z

@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: 307e4f6 Pull Request resolved: #158062

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: 57d656c Pull Request resolved: #158062

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: a8fe869 Pull Request resolved: #158062

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: 1a5ee37 Pull Request resolved: #158062

IvanKobzarev · 2025-07-17T05:08:16Z

@pytorchbot merge

pytorchmergebot · 2025-07-17T05:10:47Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-07-17T05:10:56Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-docs / build-docs-python-false

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

IvanKobzarev · 2025-07-17T11:20:31Z

@pytorchbot merge

pytorchmergebot · 2025-07-17T11:22:22Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-07-17T11:28:07Z

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x 54576f90272e2713df746c57a5b2d881f5d5bd3e returned non-zero exit code 1

Auto-merging benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
CONFLICT (content): Merge conflict in benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
Auto-merging torch/_inductor/dependencies.py
error: could not apply 54576f90272... [inductor_collectives] rewrite reorder_collectives, sink_waits_iterative
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".
hint: Disable this message with "git config set advice.mergeConflict false"

Details for Dev Infra team

Raised by workflow job

…tives, sink_waits_iterative" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D78159013](https://our.internmc.facebook.com/intern/diff/D78159013) [ghstack-poisoned]

preserving peak memory ghstack-source-id: fe170c3 Pull Request resolved: #158062

IvanKobzarev · 2025-07-17T16:12:27Z

@pytorchbot merge

pytorchmergebot · 2025-07-17T16:14:24Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-07-17T19:02:40Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-jammy-rocm-py3.10 / test (default, 1, 2, linux.rocm.gpu.2)

Details for Dev Infra team

Raised by workflow job

IvanKobzarev · 2025-07-17T19:51:52Z

@pytorchbot merge

pytorchmergebot · 2025-07-17T19:53:43Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

DEBUG sink waits

e3dc0fe

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: inductor oncall: distributed Add this issue/PR to distributed oncall triage queue labels Jul 10, 2025

IvanKobzarev mentioned this pull request Jul 10, 2025

[WIP] Rewrite reorder collectives/sink_wait to preserve peak memory #158063

Closed

Update on "DEBUG sink waits"

c89be61

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]

IvanKobzarev added a commit that referenced this pull request Jul 11, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

a52607c

preserving peak memory ghstack-source-id: cc1038b Pull Request resolved: #158062

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 11, 2025

IvanKobzarev added a commit that referenced this pull request Jul 11, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

a11bd6f

preserving peak memory ghstack-source-id: 03ea238 Pull Request resolved: #158062

IvanKobzarev changed the title ~~DEBUG sink waits~~ [simple_fsdp][inductor_collectives] rewrite reorder_collectives, sink_waits_iterative Jul 11, 2025

IvanKobzarev added the topic: not user facing topic category label Jul 11, 2025

IvanKobzarev mentioned this pull request Jul 11, 2025

[inductor_collectives] brute force preserve peak memory for sink_waits #158119

Open

superiwan pushed a commit to superiwan/pytorch that referenced this pull request Jul 14, 2025

DEBUG sink waits

a52c0fa

ghstack-source-id: 5b8d32d Pull Request resolved: pytorch/pytorch#158062

IvanKobzarev requested a review from wconstab July 14, 2025 14:36

wconstab reviewed Jul 14, 2025

View reviewed changes

test/distributed/test_inductor_collectives.py Show resolved Hide resolved

wconstab reviewed Jul 14, 2025

View reviewed changes

torch/_inductor/comms.py Show resolved Hide resolved

wconstab approved these changes Jul 14, 2025

View reviewed changes

Git-Hub-Chris pushed a commit to Git-Hub-Chris/PyTorch that referenced this pull request Jul 15, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

db35754

preserving peak memory ghstack-source-id: 30187c8 Pull Request resolved: pytorch/pytorch#158062

IvanKobzarev added a commit that referenced this pull request Jul 15, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

f92e3f4

preserving peak memory ghstack-source-id: 0b68cca Pull Request resolved: #158062

IvanKobzarev added a commit that referenced this pull request Jul 16, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

0bbb83a

preserving peak memory ghstack-source-id: 307e4f6 Pull Request resolved: #158062

IvanKobzarev added a commit that referenced this pull request Jul 16, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

47aadf6

preserving peak memory ghstack-source-id: 57d656c Pull Request resolved: #158062

IvanKobzarev added a commit that referenced this pull request Jul 16, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

edd2d30

preserving peak memory ghstack-source-id: a8fe869 Pull Request resolved: #158062

IvanKobzarev added a commit that referenced this pull request Jul 16, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

54576f9

preserving peak memory ghstack-source-id: 1a5ee37 Pull Request resolved: #158062

pytorchmergebot added the merging label Jul 17, 2025

pytorchmergebot removed the merging label Jul 17, 2025

pytorchmergebot added the merging label Jul 17, 2025

pytorchmergebot removed the merging label Jul 17, 2025

IvanKobzarev added a commit that referenced this pull request Jul 17, 2025

[inductor_collectives] rewrite reorder_collectives, sink_waits_iterative

6c82ebe

preserving peak memory ghstack-source-id: fe170c3 Pull Request resolved: #158062

pytorchmergebot added the merging label Jul 17, 2025

pytorchmergebot removed the merging label Jul 17, 2025

pytorchmergebot added the merging label Jul 17, 2025

pytorchmergebot closed this in eeb0783 Jul 17, 2025

pytorchmergebot removed the merging label Jul 17, 2025

[simple_fsdp][inductor_collectives] rewrite reorder_collectives, sink_waits_iterative #158062

[simple_fsdp][inductor_collectives] rewrite reorder_collectives, sink_waits_iterative #158062

Conversation

IvanKobzarev commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158062

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

IvanKobzarev commented Jul 11, 2025

Uh oh!

IvanKobzarev commented Jul 11, 2025

Uh oh!

IvanKobzarev commented Jul 11, 2025

Uh oh!

IvanKobzarev commented Jul 14, 2025

Uh oh!

Uh oh!

wconstab Jul 14, 2025

Choose a reason for hiding this comment

Uh oh!

IvanKobzarev Jul 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

IvanKobzarev commented Jul 15, 2025

Uh oh!

IvanKobzarev commented Jul 17, 2025

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge started

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge failed

Uh oh!

IvanKobzarev commented Jul 17, 2025

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge started

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge failed

Uh oh!

IvanKobzarev commented Jul 17, 2025

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge started

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge failed

Uh oh!

IvanKobzarev commented Jul 17, 2025

Uh oh!

pytorchmergebot commented Jul 17, 2025

Merge started

Uh oh!

Uh oh!

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

IvanKobzarev commented Jul 10, 2025 •

edited

Loading

pytorch-bot bot commented Jul 10, 2025 •

edited

Loading