[dynamo, nested graph breaks] fix nested graph breaks with error_on_graph_break #170135
Conversation
…raph_break [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/170135
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 8bbf3e0 with merge base 93bc987 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Lucaskabela
left a comment
Lucaskabela
left a comment
There was a problem hiding this comment.
Few thoughts on the changes
| torch._dynamo.graph_break() | ||
| return x + 4 | ||
|
|
||
| @torch._dynamo.error_on_graph_break(True) |
There was a problem hiding this comment.
For completeness, can we have a test covering True in the outer and True in the inner? This should be covered but want to make sure our unit tests covers it
| start_point: Optional[int] | ||
| is_leaf_tracer: bool | ||
| # Does this function make no inlined function calls? | ||
| is_leaf_function: bool |
There was a problem hiding this comment.
nit: Maybe we should just rename this to the more general is_leaf since this tracer may be for nn module (i.e the function part of the name is not so important)
There was a problem hiding this comment.
I want to avoid this because there's another concept of "leaf frame" for nested graph breaks that refers to the innermost frame; maybe we should avoid the term "leaf" for this attr completely
|
|
||
| def inline_call_(self) -> VariableTracker: | ||
| parent = self.parent | ||
| parent.is_child_tracer_active = True |
There was a problem hiding this comment.
Do we also need to set the parent's is_leaf to False?
There was a problem hiding this comment.
I think it's set somewhere else but this might be a better place to do it, yeah
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @Lucaskabela @jataylo