[annotation] Override metadata on regenerated node in functional mode#166200
[annotation] Override metadata on regenerated node in functional mode#166200
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/166200
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (10 Unrelated Failures)As of commit ea81061 with merge base 8daef35 ( BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| torch._from_functional_tensor(args[0].elem) | ||
| ].proxy.node | ||
| with fx_traceback.set_current_replay_node(curr_node): | ||
| torch._sync(a) |
There was a problem hiding this comment.
This is probably what's breaking the tests. Why is it inconvenient to annotate later during the normal replay?
yushangdi
force-pushed
the
annotation_replay
branch
4 times, most recently
from
26813e9 to
a5292b6
Compare
yushangdi
force-pushed
the
annotation_replay
branch
2 times, most recently
from
a5aa81f to
57fa5b0
Compare
torch/fx/proxy.py
Outdated
| if "seq_nr" in replay_node.meta: | ||
| annotation_log.debug("seq_nr from replay_node") | ||
| new_seq_nr = replay_node.meta["seq_nr"] | ||
| node.meta["is_functional_regenerated"] = True |
There was a problem hiding this comment.
do we actually need this info? I guess it doesn't hurt to include
There was a problem hiding this comment.
this is not used now. I just think it could be helpful for debugging later. e.g. if we see that the seq_nr/stack_trace is not consistent with some other metadata on the node, and if we see "is_functional_regenerated", we can confirm that this is due to regeneration.
torch/fx/proxy.py
Outdated
| if "custom" in replay_node.meta: | ||
| node.meta["custom"] = replay_node.meta.get("custom") | ||
| if "stack_trace" in replay_node.meta: | ||
| node.meta["stack_trace"] = replay_node.meta.get("stack_trace") |
There was a problem hiding this comment.
it seems a bit inconsistent to manually set the node.meta["stack_trace"] in this code path, which only runs when replay_node is set. Is there a shared codepath we can use for setting the stack_trace metadata?
There was a problem hiding this comment.
In particular, looking up at line 189 it seems like this will give us inconsistent values between node.meta['stack_trace'] and node.stack_trace
There was a problem hiding this comment.
node.stack_trace is the same as node.meta['stack_trace']. it's a property that does "return self.meta.get("stack_trace", None)". I changed to use node.stack_trace now to be more consistent with the other code, but it's the same as setting node.meta.
I put it at the bottom so 1) the metadata can override whatever we copied from _COPY_META_FIELDS above, and 2) the metadata overriding logic are all in the same code block.
yushangdi
force-pushed
the
annotation_replay
branch
2 times, most recently
from
df8b94a to
a2c7b42
Compare
yushangdi
force-pushed
the
annotation_replay
branch
from
a2c7b42 to
ea81061
Compare
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
5 participants
Fixes #165810
If we regenerate a node during functionalization, we override the "stack_trace", "custom", and "seq_nr" metadata of the regenerated node with the node meta of the original node.
cc @ezyang @EikanWang @jgong5 @wenzhe-nrv