handle state tensors in training ir path #137240
Conversation
avikchaudhuri
requested review from
angelayi,
tugsbayasgalan,
ydwu4 and
zhxchen17
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/137240
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit a6338cd with merge base e80f47f ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
9737ee0 to
4fb24de
Compare
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
4fb24de to
f6c61bd
Compare
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
f6c61bd to
38b8fae
Compare
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
38b8fae to
603b89b
Compare
torch/export/_trace.py
Outdated
There was a problem hiding this comment.
In training IR, we are following the convention that we don't return mutated buffers. Could we just write it as inplace update?
There was a problem hiding this comment.
If the source code has buf = new_value can we treat it the same as buf.copy_(new_value)?
There was a problem hiding this comment.
I think this is ok as long as metadata etc remains the same (aka shape). I mean when we unlift, we just write them as copy_ anyways... So seems ok? What do you think?
torch/export/_trace.py
Outdated
There was a problem hiding this comment.
Same here: Can we make it a helper function in aot_autograd and use it here?
There was a problem hiding this comment.
This part is going to be substantially different, adding copy nodes instead of mutating buffer. I'll dedup the other part.
torch/export/_trace.py
Outdated
There was a problem hiding this comment.
Could we make this a util function in aot_autograd.py and use that here? I think this logic also exist in aot_autograd.py?
There was a problem hiding this comment.
Mostly true, although there we are seeing buffers as Functional(Fake(...)) and here just Fake(...).
There was a problem hiding this comment.
Could we filter out the other parts? What i am worried is this part is bit complicated to understand so it would be better if we have only place to change stuff. Maybe the util function can take what instance we want to assert on lol.
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
603b89b to
d8c5687
Compare
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
d8c5687 to
9140277
Compare
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
Summary: Pull Request resolved: pytorch#137240 We had attribute assignment detection and handling of registered buffer assignments when using `aot_autograd`, but not when using just `make_fx`. Fixed. Test Plan: expanded coverage of `test_state_tensors` to use `export` instead of `torch.export.export` Differential Revision: D63802576
|
This pull request was exported from Phabricator. Differential Revision: D63802576 |
avikchaudhuri
force-pushed
the
export-D63802576
branch
from
9140277 to
a6338cd
Compare
|
@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 |
4 participants
Summary: We had attribute assignment detection and handling of registered buffer assignments when using
aot_autograd, but not when using justmake_fx. Fixed.Test Plan: expanded coverage of
test_state_tensorsto useexportinstead oftorch.export.exportDifferential Revision: D63802576