Fix training enablement in AOTAutograd#95975
Fix training enablement in AOTAutograd#95975ezyang wants to merge 3 commits intogh/ezyang/1872/basefrom
Conversation
Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/95975
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit cfaa39b: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]
Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]
| aot_id: int | ||
| keep_inference_input_mutations: bool | ||
| # If None, defer to config | ||
| _dynamic_shapes: Optional[bool] = None |
There was a problem hiding this comment.
Any reason not to just kill the functorch config completely, and base our decision completely off of whether the input fake tensors have a shape env?
There was a problem hiding this comment.
I have the same question here.
There was a problem hiding this comment.
I'm happy to do this in a follow up, there is just annoying test suite fixing that would need to happen
| dynamic_shapes = False | ||
| for x in full_args: | ||
| if isinstance(x, FakeTensor): | ||
| dynamic_shapes = x.fake_mode.shape_env is not None |
There was a problem hiding this comment.
Should we assert that this is all or nothing? (that either all fake tensors share the same shape env, or none of them have one?)
There was a problem hiding this comment.
Why is this change needed?
There was a problem hiding this comment.
The proximal reason for this change is, the way I decided to configure if you do AOTAutograd with dynamic shapes or not was add an entry to AOTConfig. This means that anywhere I create an AOTConfig, I need to fill in this field, if I want dynamic shapes that are not controlled by the functorch config.
In particular, when aot_module_simplified is invoked, the way we tell if we were running with dynamo dynamic shapes (because we don't have access to torch._dynamo.config, due to layering problems) is by inspecting the input arguments and checking if any of them are fake. This logic is reduxed with create_aot_dispatcher_function which also does a similar sniffing (and a further refactor would be to extract the ShapeEnv and then eliminate the equivalent logic in create_aot_dispatcher_function).
If I didn't do this change, then when we compile backwards (where we pass in all real tensors), we would not compile with dynamic shapes. I do think this part is still a little creaky though.
|
@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 |
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch/pytorch#95975 Approved by: https://github.com/ngimel, https://github.com/voznesenskym
As requested in #95975 (comment) Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]
As requested in #95975 (comment) Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: ea80ccb Pull Request resolved: #96102
As requested in #95975 (comment) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #96102 Approved by: https://github.com/Skylion007
As requested in pytorch/pytorch#95975 (comment) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch/pytorch#96102 Approved by: https://github.com/Skylion007
As requested in pytorch/pytorch#95975 (comment) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch/pytorch#96102 Approved by: https://github.com/Skylion007
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#95975 Approved by: https://github.com/ngimel, https://github.com/voznesenskym
As requested in pytorch#95975 (comment) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#96102 Approved by: https://github.com/Skylion007
Stack from ghstack (oldest at bottom):
Signed-off-by: Edward Z. Yang ezyang@meta.com