[WIP] fix reinplacing bug #152011
[WIP] fix reinplacing bug #152011
Conversation
There are two problems: 1) canonicalize_view_scatter_ops adds some new nodes into the graph. These new nodes cause the alias info on the graph to be wrong. To fix this, we try to run FakeTensorUpdater on the graph again. 2) FakeTensorUpdater's alias information is wrong. If the node was not previously seen, we need to recursively update users of the node, even if the meta["val"] looks like it is set correctly. The example is if we have `x = foo(...); y = x.view(...)`. If the user replaces `foo` with a new `bar` node and sets bar.meta["val"] correctly, then FakeTensorUpdater still needs to update y's meta["val"] to be a view of the new bar node. [ghstack-poisoned]
There are two problems: 1) canonicalize_view_scatter_ops adds some new nodes into the graph. These new nodes cause the alias info on the graph to be wrong. To fix this, we try to run FakeTensorUpdater on the graph again. 2) FakeTensorUpdater's alias information is wrong. If the node was not previously seen, we need to recursively update users of the node, even if the meta["val"] looks like it is set correctly. The example is if we have `x = foo(...); y = x.view(...)`. If the user replaces `foo` with a new `bar` node and sets bar.meta["val"] correctly, then FakeTensorUpdater still needs to update y's meta["val"] to be a view of the new bar node. ghstack-source-id: c0cdce7 Pull Request resolved: #152011
This PR needs a
|
| @@ -754,5 +754,6 @@ def tensor_with_same_storage_already_reinplaced(arg): | |||
| def reinplace_inplaceable_ops(graph: torch.fx.Graph) -> None: | |||
| with enable_python_dispatcher(): | |||
| canonicalize_view_scatter_ops(graph) | |||
| torch._inductor.fx_passes.post_grad.fake_tensor_updater.incremental_update() | |||
There was a problem hiding this comment.
The problem was:
- canonicalize_view_scatter_ops does some find and replacing of nodes.
- those new nodes have correct .meta[vals], but the users of said nodes have incorrect .meta[vals]
- so we need to update said users via FakeTensorUpdater.
Question: I don't know how to pass the FakeTensorUpdater into this pass (the GraphTransformObserver stuff is difficult to work with) so I set it as a global variable. I assume that's not what we want.
| # if "val" in node.meta and is_fake_tensor_same( | ||
| # new_fake_tensor, node.meta["val"] | ||
| # ): | ||
| # continue |
There was a problem hiding this comment.
This deletion is only important for ensuring that the alias info in the FakeTensor vals is correct. Usually there doesn't need to be correct alias info (this doesn't matter when doing functional transform), but the reinplacing pass needs the correct alias info.
Stack from ghstack (oldest at bottom):
There are two problems:
These new nodes cause the alias info on the graph to be wrong. To fix
this, we try to run FakeTensorUpdater on the graph again.
previously seen, we need to recursively update users of the node,
even if the meta["val"] looks like it is set correctly. The example
is if we have
x = foo(...); y = x.view(...). If the user replacesfoowith a newbarnode and sets bar.meta["val"] correctly, thenFakeTensorUpdater still needs to update y's meta["val"] to be a view
of the new bar node.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov