[WIP] functional autograd + compiled autograd #139098
Draft
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139098
Note: Links to docs will display an error until the docs builds have been completed. ❌ 36 New FailuresAs of commit 368259d with merge base 5c88a9f ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
zou3519
added a commit
that referenced
this pull request
Oct 28, 2024
This commit refactors autograd so that nodes can be called in a functional way. Furthermore, it refactors compiled autograd to use the new functional autograd, without any behavior changes. This is on the way to getting compiled autograd to stop tracing into autograd nodes when it constructs an FX graph out of the autograd graph. We also implement some very basic support for that, which can be toggled via `old_inline_behavior=False` in compiled_autograd.py. Functional autograd works like the following: - All torch::autograd::Node must define a `retrieve_saved(SwapSavedVariables) -> ivalue_list` API. This function takes compiled autograd's SwapSavedVariables and packs the state that is relevant to the current Node into an ivalue_list. - All torch::autograd::Node must define a `get_functional() -> std::function`. This returns a new stateless function that accepts the gradients and saved values as an ivalue_list and returns new gradients. - We developed a mechanism to bind arbitrary C++ functions that take ivalue_list to Python. This is really similar to how we bind custom ops to Python and was done in consideration of the Windows symbol limit (otherwise, we'd be binding one symbol per Node into Python). Here's an example of the new autograd generated code - https://gist.github.com/zou3519/09bb98bb0f11445bc3da063201adb818 Here's an example of the FX graph compiled autograd produces (with old_inline_behavior=False): - https://gist.github.com/zou3519/43e8106176d15d623e1377850f585c97 ghstack-source-id: c2f93cd46dd245ccc26cf1bd07f861cb18267eaf Pull Request resolved: #139098
pytorch-bot
bot
added
ciflow/inductor
module: compiled autograd
facebook-github-bot
added
the
oncall: jit
This was referenced Nov 6, 2024
Closed
Closed
fightingand
pushed a commit
to fightingand/pytorch
that referenced
this pull request
Dec 20, 2024
This commit refactors autograd so that nodes can be called in a functional way. Furthermore, it refactors compiled autograd to use the new functional autograd, without any behavior changes. This is on the way to getting compiled autograd to stop tracing into autograd nodes when it constructs an FX graph out of the autograd graph. We also implement some very basic support for that, which can be toggled via `old_inline_behavior=False` in compiled_autograd.py. Functional autograd works like the following: - All torch::autograd::Node must define a `retrieve_saved(SwapSavedVariables) -> ivalue_list` API. This function takes compiled autograd's SwapSavedVariables and packs the state that is relevant to the current Node into an ivalue_list. - All torch::autograd::Node must define a `get_functional() -> std::function`. This returns a new stateless function that accepts the gradients and saved values as an ivalue_list and returns new gradients. - We developed a mechanism to bind arbitrary C++ functions that take ivalue_list to Python. This is really similar to how we bind custom ops to Python and was done in consideration of the Windows symbol limit (otherwise, we'd be binding one symbol per Node into Python). Here's an example of the new autograd generated code - https://gist.github.com/zou3519/09bb98bb0f11445bc3da063201adb818 Here's an example of the FX graph compiled autograd produces (with old_inline_behavior=False): - https://gist.github.com/zou3519/43e8106176d15d623e1377850f585c97 ghstack-source-id: c2f93cd46dd245ccc26cf1bd07f861cb18267eaf Pull Request resolved: pytorch/pytorch#139098
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
This commit refactors autograd so that nodes can be called in a
functional way. Furthermore, it refactors compiled autograd to use
the new functional autograd, without any behavior changes.
This is on the way to getting compiled autograd to stop tracing into
autograd nodes when it constructs an FX graph out of the autograd graph.
We also implement some very basic support for that, which can be toggled
via
old_inline_behavior=Falsein compiled_autograd.py.Functional autograd works like the following:
retrieve_saved(SwapSavedVariables) -> ivalue_listAPI. This functiontakes compiled autograd's SwapSavedVariables and packs the state that
is relevant to the current Node into an ivalue_list.
get_functional() -> std::function.This returns a new stateless function that accepts the
gradients and saved values as an ivalue_list and returns new
gradients.
ivalue_list to Python.
This is really similar to how we bind custom ops to Python and was
done in consideration of the Windows symbol limit (otherwise, we'd be
binding one symbol per Node into Python).
Here's an example of the new autograd generated code
Here's an example of the FX graph compiled autograd produces (with
old_inline_behavior=False):
cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @voznesenskym @penguinwu @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @rec @xmfan