Add functional collective all_to_all_single and support it in Inductor #110195

ezyang · 2023-09-28T03:25:57Z

Stack from ghstack (oldest at bottom):

-> Add functional collective all_to_all_single and support it in Inductor #110195

Copy of #106655 from yf225
rebased on top of item() support changes

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler

Copy of #106655 from yf225 rebased on top of item() support changes [ghstack-poisoned]

pytorch-bot · 2023-09-28T03:25:59Z

🔗 Helpful Links

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

📄 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! (2 Unrelated Failures)

As of commit bd0c738 with merge base cf1b494 ():

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

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

ezyang · 2023-09-28T03:27:42Z

cc @yf225

ezyang · 2023-09-28T03:28:17Z

Hmm I think I didn't do the merge update completely correctly

… in Inductor" Copy of #106655 from yf225 rebased on top of item() support changes cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

yf225 · 2023-09-28T19:26:17Z

test/distributed/test_inductor_collectives.py

+            compiled_fn = torch.compile(example, fullgraph=True, dynamic=True)
+            code = run_and_get_triton_code(compiled_fn, *inputs, **trs)
+            FileCheck() \
+                .check("fun_col_impl._all_to_all_single") \


should we assert that generated code has the unbacked symints at the right places?

it was a pain manually updating the FileChecks when I wobbled the code lol. Can we come up with a less brittle way to test this

perhaps a regex matching the prefix i but not the actual variable number? I can add it after this PR

… in Inductor" Copy of #106655 from yf225 rebased on top of item() support changes cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

Copy of #106655 from yf225 rebased on top of item() support changes ghstack-source-id: 3ae9f68 Pull Request resolved: #110195

torch/_inductor/ir.py

ezyang · 2023-09-28T23:56:05Z

torch/distributed/_functional_collectives.py


    return [mk_out_tensor(t) for t in inputs]

+# TODO: think, this isn't actually dynamic?!


I was surprised here

torch/fx/experimental/symbolic_shapes.py

Skylion007 · 2023-10-02T14:43:07Z

torch/_inductor/ir.py

        return set()

+    def codegen_unbacked_symbol_defs(self, wrapper):
+        symbols_to_define = self.get_unbacked_symbol_defs()


Hmm, if it wasn't for the logging here, I would advise rewriting this as set intersection / difference ops.

ezyang · 2023-10-05T13:16:29Z

@pytorchbot merge

pytorchmergebot · 2023-10-05T13:18:45Z

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

torch/_inductor/ir.py

… in Inductor" Copy of #106655 from yf225 rebased on top of item() support changes cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

pytorchmergebot · 2023-10-05T13:29:43Z

Merge failed

Reason: New commits were pushed while merging. Please rerun the merge command.

Details for Dev Infra team

Raised by workflow job

… in Inductor" Copy of #106655 from yf225 rebased on top of item() support changes cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

Copy of #106655 from yf225 rebased on top of item() support changes ghstack-source-id: e02044d Pull Request resolved: #110195

ezyang · 2023-10-05T13:48:35Z

@yf225 because the kernel is no longer "dynamic", it can be implemented exactly symmetrically to all the other collectives. I've updated the implementation accordingly.

ezyang · 2023-10-05T23:09:44Z

@pytorchbot merge

pytorchmergebot · 2023-10-05T23:11:32Z

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

Add functional collective all_to_all_single and support it in Inductor

108d0b6

Copy of #106655 from yf225 rebased on top of item() support changes [ghstack-poisoned]

ezyang requested review from H-Huang, awgu, fduwjj, fegin, kiukchung, kwen2501, mrshenli, rohan-varma, wanchaol, wz337 and zhaojuanmao as code owners September 28, 2023 03:25

ezyang requested a review from d4l3k as a code owner September 28, 2023 03:25

ezyang mentioned this pull request Sep 28, 2023

Add support for item() and nonzero() codegen in Inductor #109893

Closed

pytorch-bot bot added the release notes: fx release notes category label Sep 28, 2023

github-actions bot added module: inductor ciflow/inductor labels Sep 28, 2023

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh, miladm, voznesenskym and wconstab September 28, 2023 03:26

yf225 reviewed Sep 28, 2023

View reviewed changes

ezyang added a commit that referenced this pull request Sep 28, 2023

Add functional collective all_to_all_single and support it in Inductor

ffd342b

Copy of #106655 from yf225 rebased on top of item() support changes ghstack-source-id: 3ae9f68 Pull Request resolved: #110195

ezyang commented Sep 28, 2023

View reviewed changes

torch/_inductor/ir.py Outdated Show resolved Hide resolved

ezyang commented Sep 28, 2023

View reviewed changes

torch/fx/experimental/symbolic_shapes.py Outdated Show resolved Hide resolved

ezyang mentioned this pull request Sep 29, 2023

Inductor does not check input SymInt invariant on GraphModules passed in #104102

Closed

Skylion007 reviewed Oct 2, 2023

View reviewed changes

Skylion007 approved these changes Oct 2, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 5, 2023

pytorchmergebot added the merging label Oct 5, 2023

ezyang commented Oct 5, 2023

View reviewed changes

torch/_inductor/ir.py Outdated Show resolved Hide resolved

pytorchmergebot removed the merging label Oct 5, 2023

ezyang added a commit that referenced this pull request Oct 5, 2023

Add functional collective all_to_all_single and support it in Inductor

8c7c1b3

Copy of #106655 from yf225 rebased on top of item() support changes ghstack-source-id: e02044d Pull Request resolved: #110195

pytorchmergebot added the merging label Oct 5, 2023

pytorchmergebot added Merged and removed merging labels Oct 5, 2023

pytorchmergebot closed this in f274c7b Oct 5, 2023

facebook-github-bot deleted the gh/ezyang/2375/head branch October 9, 2023 14:23

yf225 mentioned this pull request Oct 23, 2023

Add regex matching to Inductor all2all collective unit tests #111776

Closed


		return [mk_out_tensor(t) for t in inputs]

		# TODO: think, this isn't actually dynamic?!

Add functional collective all_to_all_single and support it in Inductor #110195

Add functional collective all_to_all_single and support it in Inductor #110195

Uh oh!

Conversation

ezyang commented Sep 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 28, 2023 • 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/110195

✅ You can merge normally! (2 Unrelated Failures)

Uh oh!

ezyang commented Sep 28, 2023

Uh oh!

ezyang commented Sep 28, 2023

Uh oh!

yf225 Sep 28, 2023

Choose a reason for hiding this comment

Uh oh!

ezyang Sep 28, 2023

Choose a reason for hiding this comment

Uh oh!

yf225 Oct 2, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ezyang Sep 28, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Skylion007 Oct 2, 2023

Choose a reason for hiding this comment

Uh oh!

ezyang commented Oct 5, 2023

Uh oh!

pytorchmergebot commented Oct 5, 2023

Merge started

Uh oh!

Uh oh!

pytorchmergebot commented Oct 5, 2023

Merge failed

Uh oh!

ezyang commented Oct 5, 2023

Uh oh!

ezyang commented Oct 5, 2023

Uh oh!

pytorchmergebot commented Oct 5, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ezyang commented Sep 28, 2023 •

edited

Loading

pytorch-bot bot commented Sep 28, 2023 •

edited

Loading