Add tests to aid in debugging dynamo + aot_autograd dedup #90130
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90130
Note: Links to docs will display an error until the docs builds have been completed. ❌ 5 FailuresAs of commit 88173a0: The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
github-actions
bot
requested review from
SherlockNoMad,
albanD,
anjali411,
antoniojkim,
bdhirsh,
miladm and
wconstab
voznesenskym
changed the title
test/functorch/test_aotdispatch.py
Outdated
| ) | ||
|
|
||
| @patch('functorch._src.aot_autograd.AOT_COUNTER', new_callable=itertools.count) | ||
| @patch("functorch._src.config.debug_assert", True) |
…grad dedup" This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors. cc mlazos soumith yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
test/dynamo/test_aot_autograd.py
Outdated
| for out in pytree.tree_flatten(outs)[0]: | ||
| if isinstance(out, torch.Tensor) and out.requires_grad: | ||
| out.sum().backward(retain_graph=True) | ||
| grads = [inp.grad for inp in pytree.tree_flatten(inps)[0]] | ||
| for inp in pytree.tree_flatten(inps)[0]: | ||
| inp.grad = None |
There was a problem hiding this comment.
This is a simpler way to do this
| for out in pytree.tree_flatten(outs)[0]: | |
| if isinstance(out, torch.Tensor) and out.requires_grad: | |
| out.sum().backward(retain_graph=True) | |
| grads = [inp.grad for inp in pytree.tree_flatten(inps)[0]] | |
| for inp in pytree.tree_flatten(inps)[0]: | |
| inp.grad = None | |
| flat_outs = pytree.tree_flatten_only(torch.Tensor, outs)[0] | |
| flat_inps = pytree.tree_flatten_only(torch.Tensor, inps)[0] | |
| grads = torch.autograd.grad((o for o in flat_outs if o.requires_grad), flat_inps) | |
| return outs, grads |
There was a problem hiding this comment.
Also, don't we already have this helper function defined, why do we need another copy?
|
|
||
| x = torch.randn(3, 3, requires_grad=True) | ||
| y = torch.randn(3, 3, requires_grad=True) | ||
| z = torch.randn(3, 3, requires_grad=False) |
There was a problem hiding this comment.
nit: it would have been better to have more descriptive name for y and z; e.g., y_true and y_false or y_grad and y_nograd, to make it more clear which ones is requires grad True and which one is False.
|
I mentioned this in earlier review; these tests are straight up copy pastes of already existing tests in the test suite. Lines of code in the test suite are a liability, not an asset, since the more lines of test code you have, the harder it is to refactor APIs (since you now have to update all the test code.) Please make more effort to deduplicate these tests. |
One is, and one is different enough and needs some dynamo test imports that we'd probably need to keep it. But for the aot_dispatch test, 100% agreed.
Of course :) I agree! Higher level - does the reasoning in the tests for why we have these two seem sound? I was hoping more to get a take from you on that. |
|
The test itself is fine. I'm imagining a new type of test which is done in test/dynamo, but the model in question is simple enough that you can also run it through directly with aot_module_simplified. Or maybe something like OpInfo but ProgramInfo, where we have sample programs that can be run through various levels of the stack. |
…grad dedup" This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors. cc mlazos soumith yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
…grad dedup" This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors. cc mlazos soumith yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
voznesenskym
changed the title
This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors. cc mlazos soumith yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
|
i'm unsubsubscribing, rerequest review when ready |
This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors. cc mlazos soumith yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Stack from ghstack (oldest at bottom):
This adds a pair of equivalent tests, both of which invoke aot_autograd with the same inputs, more or less. One through dynamo, one directly. Where the direct call fails, dynamo passes due to guards. This pair of tests protects that behavior and sets us up for future refactors.
cc @mlazos @soumith @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire