[inductor collectives] sink waits iterative#157708
[inductor collectives] sink waits iterative#157708IvanKobzarev wants to merge 4 commits intogh/IvanKobzarev/122/basefrom
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157708
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 9a737bb with merge base d27d361 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@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: [D77861763](https://our.internmc.facebook.com/intern/diff/D77861763) [ghstack-poisoned]
|
@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| "reorder_for_compute_comm_overlap": True, | ||
| "reorder_for_compute_comm_overlap_passes": [ | ||
| _reorder_communication_preserving_peak_memory, | ||
| sink_waits_iterative, |
There was a problem hiding this comment.
I expected the sink would happen before the reorder. I guess now I'm not sure if it actually matters what order it is in.
For the case where collectives all independent of one another, i think moving the wait can be either before or after reorder, since the wait is not an obstacle for the reorder pass.
For the case where one collective produces a temporary that feeds into another collective, i guess it could be worse to sink waits first, becuase doing so could lock the second collective in later position.
OK- i'm convinced this is right.
There was a problem hiding this comment.
As we discussed offline - I've added a ban of swaps of groups that both contain collective_ops.
Once I added - the result of [reorder, sink_waits] and [sink_waits, reorder] became the same.
torch/_inductor/comms.py
Outdated
| ) -> tuple[list[BaseSchedulerNode], dict[BaseSchedulerNode, SinkWaitInfo]]: | ||
| from torch._inductor.scheduler import GroupedSchedulerNode, init_group_node | ||
|
|
||
| def _group_name(snode, with_bufs=False): |
There was a problem hiding this comment.
should we elevate this to a util and use it from the reorder pass as well?
| gsnode = gsnodes[i] | ||
| if contains_wait(gsnode): | ||
| info = stats[gsnode.snodes[0]] = SinkWaitInfo() | ||
| for j in range(i + 1, n): |
There was a problem hiding this comment.
since a lot of stuff is boilerplate copy/paste from the reorder pass (the loop structure, the logging stuff) i wonder if we can make a reusable thing that has callbacks for the internal logic.
otoh i'd probably rather get things working nicely first and think about this later
| @@ -1613,6 +1614,7 @@ def _reorder_communication_preserving_peak_memory( | |||
| "reorder_for_compute_comm_overlap": True, | |||
| "reorder_for_compute_comm_overlap_passes": [ | |||
| _reorder_communication_preserving_peak_memory, | |||
There was a problem hiding this comment.
i think it'd be good to add a test that shows that we can move a wait (wait_a) past a collective and past another wait (wait_b), in the case that wait_b and collective are not actually depending on the wait_a
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: [D77861763](https://our.internmc.facebook.com/intern/diff/D77861763) [ghstack-poisoned]
|
@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: [D77861763](https://our.internmc.facebook.com/intern/diff/D77861763) [ghstack-poisoned]
|
@IvanKobzarev has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@pytorchbot merge |
Merge startedYour 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 |
Stack from ghstack (oldest at bottom):
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: D77861763