Fix decomp behaviour in export training IR #134801
Conversation
[ghstack-poisoned]
tugsbayasgalan
requested review from
angelayi,
avikchaudhuri,
ydwu4 and
zhxchen17
as code owners
[ghstack-poisoned]
|
@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
Subset of changes in #132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export. Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D62000525 |
|
This pull request was exported from Phabricator. Differential Revision: D62000525 |
torch/_export/utils.py
Outdated
| fake_val = node.meta["val"] | ||
| if fake_val is not None and isinstance(fake_val, torch.Tensor): | ||
| fake_vals.append(fake_val) | ||
|
|
There was a problem hiding this comment.
I think this makes sense, but could you explain what prompted the change? The previous implementation seemed equivalent, and only required one loop right?
| node.meta[k] = v | ||
| _populate_param_buffer_metadata_to_new_gm( | ||
| params_buffers_to_node_meta, gm, export_graph_signature | ||
| ) |
| return gm, new_graph_signature | ||
|
|
||
|
|
||
| def _remove_unneccessary_copy_op_pass( |
There was a problem hiding this comment.
noob question: why is this op special?
There was a problem hiding this comment.
This is because ep.module() adds copy_ nodes in the end to update the buffers. When we retrace, we functionalize these nodes and they will show up as extra nodes in the end. We actually don't need it because aot_export_module will take care of returning extra updated buffers.
Subset of changes in #132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export. Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D62000525 |
Pull Request resolved: #134801 @imported-using-ghimport Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525/) ghstack-source-id: 240844924
Subset of changes in #132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export. Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D62000525 |
Pull Request resolved: #134801 @imported-using-ghimport Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525/) ghstack-source-id: 240909553
Subset of changes in #132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export. Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525) [ghstack-poisoned]
|
@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| # When aot_export lifts the params, we lose metadata (e.g. source_fn_stack, stack_trace) | ||
| # from the param nodes as they are treated as fresh inputs | ||
| # Therefore, we manually extract them before calling into aot_export | ||
| # params_buffers_to_node_meta = _collect_param_buffer_metadata(gm_torch_level) |
|
@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@pytorchbot merge -f 'Landed internally' (Initiating merge automatically since Phabricator Diff has merged, using force because this PR might not pass merge_rules.json but landed internally) |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
This pull request was exported from Phabricator. Differential Revision: D62000525 |
Pull Request resolved: pytorch/pytorch#134801 @imported-using-ghimport Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525/) ghstack-source-id: 15a3e01
Subset of changes in pytorch#132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export. Differential Revision: [D62000525](https://our.internmc.facebook.com/intern/diff/D62000525) Pull Request resolved: pytorch#134801 Approved by: https://github.com/avikchaudhuri, https://github.com/angelayi
6 participants
Stack from ghstack (oldest at bottom):
Subset of changes in #132901, can't land the previous one because it is too complicated. Rest of the change will be implemented as follow up after export design meeting. This part just makes the training IR -> inference IR decomp to have the same path as normal export.
Differential Revision: D62000525