[dtensor] remove torchgen function schema and parse manually #90106

wanchaol · 2022-12-03T04:10:25Z

Stack from ghstack (oldest at bottom):

This PR get rids of torchgen FunctionSchema parsing and parse
it manually, it should resolve torchgen package issue and also
provide some perf wins when running DTensor eagerly

This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly [ghstack-poisoned]

pytorch-bot · 2022-12-03T04:10:28Z

🔗 Helpful Links

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

📄 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 our office hours

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

❌ 1 Failures

As of commit fa891ef:

The following jobs have failed:

manywheel-py3_7-cuda11_6-test / build

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

This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly ghstack-source-id: df45231ebe398f97bdfe75de05adecae1faed1ab Pull Request resolved: #90106

awgu · 2022-12-03T04:21:44Z

torch/distributed/_tensor/dispatch.py

-        self.is_out_variant = (
-            schema_kind == SchemaKind.out  # pyre-ignore [16] pyre bad at enum
-        )
+        # simple analysis of function schema to determine


Conceptual: Does this PR regress the computation of self.is_inplace and self._is_out_variant in some cases because you are using a heuristic instead of SchemaKind?

I am curious how you think about the tradeoff of getting rid of the torchgen dependency versus possibly having incorrect values for is_inplace and is_out_variant in some cases (if at all).

Good point, from the current operator set registered, i think it should work fine as all of the op signature naming following the same convention for inplace and out variant (_ and .out), so I think this heuristic is relatively safe to use and should not regress existing registered operators. However, it does rely on operators following the same convention of naming, which I think it might get changed. Given that ATen op set is relatively stable at this point, I think we can go with this approach as it's simple and very fast without parsing FunctionSchema, if ever op signature gets changed, the heuristic need to be adapted to accommodate.

awgu

I am not sure I have the context to approve, so I will defer to someone else.

awgu · 2022-12-05T21:00:17Z

torch/distributed/_tensor/dispatch.py

+                out_dt = cast(dtensor.DTensor, kwargs[arg.name])
+                out_dt._spec = cast(DTensorSpec, output_specs[spec_idx])
+                out_dts.append(out_dt)
+                spec_idx += 1
        return tuple(out_dts) if len(out_dts) > 1 else out_dts[0]


General comment: This return assumes len(out_dts) >= 1 or else there will be an index out of bounds error when accessing out_dts[0]. I am guessing that this len(out_dts) >= 1 invariant always holds, but I just wanted to point this out. Maybe you want an assert?

I'll land this as is as I don't want to trigger another full CI run, the assert is added in the follow up PR #90241

nvm, might still need to change sth, will update.

awgu

This looks good to me, so I will stamp to unblock. Feel free to wait for others' review.

fduwjj

Overall this looks good to me and do we have an estimate of potential perf win?

This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly [ghstack-poisoned]

wanchaol · 2022-12-06T03:15:34Z

Overall this looks good to me and do we have an estimate of potential perf win?

I don't yet as I didn't profile this, but my suspection is that it could result in around 20-30 perf wins on DTensor execution, feel free to try it out, I'll focus on making other parts be faster if possible.

wanchaol · 2022-12-06T05:38:12Z

@pytorchbot merge

pytorchmergebot · 2022-12-06T05:41: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 · 2022-12-06T05:41:25Z

Merge failed

Reason: 2 additional jobs have failed, first few of them are: linux-binary-manywheel ,linux-binary-manywheel / manywheel-py3_7-cuda11_6-test / build

Details for Dev Infra team

Raised by workflow job

wanchaol · 2022-12-06T05:42:41Z

@pytorchbot merge -f "failure not related"

pytorchmergebot · 2022-12-06T05:44:54Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

…#90106) This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly Pull Request resolved: pytorch#90106 Approved by: https://github.com/awgu

) This is a reland of #89845 with nothing changed. This should avoid the internal breakage now that `DTensor` does not import `torchgen` (#90106). Pull Request resolved: #90562 Approved by: https://github.com/fduwjj

[dtensor] remove torchgen function schema and parse manually

d06beba

This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly [ghstack-poisoned]

wanchaol requested review from mrshenli, pritamdamania87, zhaojuanmao, rohan-varma, H-Huang, awgu and kwen2501 as code owners December 3, 2022 04:10

wanchaol added the release notes: distributed (dtensor) release notes category label Dec 3, 2022

awgu reviewed Dec 3, 2022

View reviewed changes

wanchaol requested a review from fduwjj December 3, 2022 04:41

wanchaol added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category and removed release notes: distributed (dtensor) release notes category labels Dec 3, 2022

awgu reviewed Dec 5, 2022

View reviewed changes

wanchaol requested a review from aazzolini December 5, 2022 23:13

awgu approved these changes Dec 5, 2022

View reviewed changes

fduwjj reviewed Dec 5, 2022

View reviewed changes

wanchaol mentioned this pull request Dec 6, 2022

[dtensor] handle the case where output of op is Optional[Tensor] #90241

Closed

Update on "[dtensor] remove torchgen function schema and parse manually"

fa891ef

This PR get rids of torchgen FunctionSchema parsing and parse it manually, it should resolve torchgen package issue and also provide some perf wins when running DTensor eagerly [ghstack-poisoned]

pytorchmergebot added the Merged label Dec 6, 2022

pytorchmergebot closed this in 2c2cce7 Dec 6, 2022

awgu mentioned this pull request Dec 9, 2022

[Reland][FSDP] Another fix for DTensor, use_orig_params=True #90562

Closed

facebook-github-bot deleted the gh/wanchaol/225/head branch June 8, 2023 19:06

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[dtensor] remove torchgen function schema and parse manually #90106

[dtensor] remove torchgen function schema and parse manually #90106

wanchaol commented Dec 3, 2022 •

edited

pytorch-bot bot commented Dec 3, 2022 •

edited

awgu Dec 3, 2022

wanchaol Dec 3, 2022

awgu left a comment

awgu Dec 5, 2022

wanchaol Dec 6, 2022

wanchaol Dec 6, 2022

awgu left a comment

fduwjj left a comment

wanchaol commented Dec 6, 2022

wanchaol commented Dec 6, 2022

pytorchmergebot commented Dec 6, 2022

pytorchmergebot commented Dec 6, 2022

wanchaol commented Dec 6, 2022

pytorchmergebot commented Dec 6, 2022

[dtensor] remove torchgen function schema and parse manually #90106

[dtensor] remove torchgen function schema and parse manually #90106

Conversation

wanchaol commented Dec 3, 2022 • edited

pytorch-bot bot commented Dec 3, 2022 • edited

🔗 Helpful Links

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

❌ 1 Failures

awgu Dec 3, 2022

Choose a reason for hiding this comment

wanchaol Dec 3, 2022

Choose a reason for hiding this comment

awgu left a comment

Choose a reason for hiding this comment

awgu Dec 5, 2022

Choose a reason for hiding this comment

wanchaol Dec 6, 2022

Choose a reason for hiding this comment

wanchaol Dec 6, 2022

Choose a reason for hiding this comment

awgu left a comment

Choose a reason for hiding this comment

fduwjj left a comment

Choose a reason for hiding this comment

wanchaol commented Dec 6, 2022

wanchaol commented Dec 6, 2022

pytorchmergebot commented Dec 6, 2022

Merge started

pytorchmergebot commented Dec 6, 2022

Merge failed

wanchaol commented Dec 6, 2022

pytorchmergebot commented Dec 6, 2022

Merge started

wanchaol commented Dec 3, 2022 •

edited

pytorch-bot bot commented Dec 3, 2022 •

edited