Make Dynamo inline through torch._library.custom_ops.autograd #135066
Conversation
Fixes #135057 Test Plan: - new test [ghstack-poisoned]
| @@ -172,6 +172,26 @@ def fn(): | |||
|
|
|||
| self.check_output_and_recompiles(fn) | |||
|
|
|||
| def test_graph_break_custom_op(self): | |||
There was a problem hiding this comment.
curious: does it mean we need to graph-break when custom op is used in compiled autograd graph?
(Traceable FSDP2 requires that compiled autograd backward graph has no graph break, so hoping graph-break here is not a requirement haha)
There was a problem hiding this comment.
curious: does it mean we need to graph-break when custom op is used in compiled autograd graph?
No. We are testing what happens if (1) a user has a custom op that graph breaks in the forward and (2) we use compiled autograd on this. Previously, this would raise an error. After this PR, it works
…rad" Fixes #135057 The bug was: in the situation that Dynamo graph breaks in the forward and Compiled Autograd uses Dynamo to introspect the backward, we end up running into a "Unsupported: inlining through SKIPFILES" error. The solution is to mark the entirety of this module as inlineable. Test Plan: - new test [ghstack-poisoned]
…rad" Fixes #135057 The bug was: in the situation that Dynamo graph breaks in the forward and Compiled Autograd uses Dynamo to introspect the backward, we end up running into a "Unsupported: inlining through SKIPFILES" error. The solution is to mark the entirety of this module as inlineable. Test Plan: - new test [ghstack-poisoned]
Fixes #135057 The bug was: in the situation that Dynamo graph breaks in the forward and Compiled Autograd uses Dynamo to introspect the backward, we end up running into a "Unsupported: inlining through SKIPFILES" error. The solution is to mark the entirety of this module as inlineable. Test Plan: - new test ghstack-source-id: 63bfc87 Pull Request resolved: #135066
|
@pytorchbot merge -f "unrelated failures" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…h#135066) Fixes pytorch#135057 The bug was: in the situation that Dynamo graph breaks in the forward and Compiled Autograd uses Dynamo to introspect the backward, we end up running into a "Unsupported: inlining through SKIPFILES" error. The solution is to mark the entirety of this module as inlineable. Test Plan: - new test Pull Request resolved: pytorch#135066 Approved by: https://github.com/bdhirsh, https://github.com/williamwen42, https://github.com/yanboliang
Stack from ghstack (oldest at bottom):
Fixes #135057
The bug was: in the situation that Dynamo graph breaks in the forward
and Compiled Autograd uses Dynamo to introspect the backward, we end up
running into a "Unsupported: inlining through SKIPFILES" error. The
solution is to mark the entirety of this module as inlineable.
Test Plan:
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @rec