[scan] cloned aliased input when lowering scan to while_loop #158168
+196
−2
Conversation
This file contains hidden or 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
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158168
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 4cddc7d with merge base fc25c68 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Fixes #153679. The CSE pass accidently de-dup the zero buffer that are allocated to store the gradients of init to scan causing an aliasing among thei nputs. CSE is after dynamo and aot so went undetected by our front end safety checks. The fix follows the same spirit of skipping cse aten.empty, we just don't cse aten.full. The alternataive considered: 1. figure out a way that tell cse don't de-dup in this case. This seems tedious to do... 2. Detect the aliasing and do a clone when we decompose scan to while_loop. This is OK but in the spirit of fixing something wrong instead of fixing the original problem. [ghstack-poisoned]
How does this cause aliasing among the inputs? |
grad_init1 = torch.full(s1, s2, ...)
grad_init2 = torch.full(s1, s2, ...)
torch.ops.higher_order.scan(backward_combine_graph, (grad_init1, grad_init2), ...)The two torch.full happen to have the same argument because the forward_combine_graph's output shape is the same. CSE ffind the two torch.full have the same argument therefore are subject to cse. So it changes the graph to be: grad_init = torch.full(s1, s2, ...)
torch.ops.higher_order.scan(backward_combine_graph, (grad_init, grad_init), ...) |
Fixes #153679. The CSE pass accidently de-dups the zero buffer that are allocated to store the gradients of init to scan causing an aliasing among the inputs. CSE is after dynamo and aot so went undetected by our front-end safety checks. And it's not wrapped by auto_functionalized because there're no aliasing in auto_grad and functioanlization key becaues the zeros_like are not de-duped yet. The fix follows the same spirit of skipping cse aten.empty: we just don't cse aten.full. The alternatives include: 1. figure out a way that tell cse don't de-dup in this case. This seems tedious to do and complicates the cse function's interface. 2. Detect the aliasing and do a clone when we decompose scan to while_loop. This is OK but seems ad-hoc. Also it breaks the contract that there shouldn't be aliasing among inputs for hops if not wrapped with auto_functionalized. [ghstack-poisoned]
Fixes #153679. The CSE pass accidently de-dups the zero buffer that are allocated to store the gradients of init to scan causing an aliasing among the inputs. CSE is after dynamo and aot so went undetected by our front-end safety checks. And it's not wrapped by auto_functionalized because there're no aliasing in auto_grad and functioanlization key becaues the zeros_like are not de-duped yet. The alternatives include: 1. figure out a way that tell cse don't de-dup in this case. This seems tedious to do and complicates the cse function's interface. 2. Detect the aliasing and do a clone when we decompose scan to while_loop. This is OK but seems ad-hoc. Also it breaks the contract that there shouldn't be aliasing among inputs for hops if not wrapped with auto_functionalized. The final implemented fix after discussion with zou3519 is 2: aliasing is ok as long as there're no in place mutations, so we can fix it right before we start introducing buffers etc. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben [ghstack-poisoned]
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):
Fixes #153679.
The CSE pass accidently de-dups the zero buffer that are allocated to store the gradients of init to scan causing an aliasing among the inputs. CSE is after dynamo and aot so went undetected by our front-end safety checks. And it's not wrapped by auto_functionalized because there're no aliasing in auto_grad and functioanlization key becaues the zeros_like are not de-duped yet.
The alternatives include:
The final implemented fix after discussion with @zou3519 is 2: aliasing is ok as long as there're no in place mutations, so we can fix it right before we start introducing buffers etc.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben