Skip to content

Conversation

xmfan
Copy link
Member

@xmfan xmfan commented Oct 30, 2024

So for functional autograd + CA, most nodes are inlined in aot autograd. But user-defined callables aren't safe to make_fx unless dynamo traces through them. The AOT backward must be inlined by dynamo time. We plan to directly insert calls to the backward in the graph:

  • call prologue
  • call bwd graph
  • call epilogue

Restructuring our AOT bwd implementation will make this implementation easier.

Stack from ghstack (oldest at bottom):

Copy link

pytorch-bot bot commented Oct 30, 2024

🔗 Helpful Links

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

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

✅ No Failures

As of commit c7ee4a0 with merge base a951d99 (image):
💚 Looks good so far! There are no failures yet. 💚

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

xmfan added a commit that referenced this pull request Oct 30, 2024
xmfan added a commit that referenced this pull request Oct 30, 2024
@xmfan xmfan changed the title [aot autograd] refactor CompiledFunction.backward (1/N) [aot autograd] refactor CompiledFunction.backward: prologue (1/N) Oct 30, 2024
@xmfan xmfan added the topic: not user facing topic category label Oct 31, 2024
@xmfan xmfan marked this pull request as ready for review October 31, 2024 07:30
@xmfan xmfan requested a review from bdhirsh as a code owner October 31, 2024 07:30
@xmfan xmfan requested a review from zou3519 October 31, 2024 07:31
pytorchmergebot pushed a commit that referenced this pull request Oct 31, 2024
pytorchmergebot pushed a commit that referenced this pull request Oct 31, 2024
rahulsingh-intel pushed a commit to rahulsingh-intel/pytorch that referenced this pull request Nov 5, 2024
…torch#139331)

So for functional autograd + CA, most nodes are inlined in aot autograd. But user-defined callables aren't safe to make_fx unless dynamo traces through them. The AOT backward must be inlined by dynamo time. We plan to directly insert calls to the backward in the graph:
- call prologue
- call bwd graph
- call epilogue

Restructuring our AOT bwd implementation will make this implementation easier.

Pull Request resolved: pytorch#139331
Approved by: https://github.com/zou3519
rahulsingh-intel pushed a commit to rahulsingh-intel/pytorch that referenced this pull request Nov 5, 2024
rahulsingh-intel pushed a commit to rahulsingh-intel/pytorch that referenced this pull request Nov 5, 2024
Esquains pushed a commit to Esquains/study1 that referenced this pull request Dec 15, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants