Create context managers before entering any with ExitStack #18716
Conversation
⚡ Required checks status: All passing 🟢Groups summary🟢 pytorch_lightning: Tests workflow
These checks are required after the changes to 🟢 pytorch_lightning: Azure GPU
These checks are required after the changes to 🟢 pytorch_lightning: Benchmarks
These checks are required after the changes to 🟢 fabric: Docs
These checks are required after the changes to 🟢 lightning_fabric: CPU workflowThese checks are required after the changes to 🟢 lightning_fabric: Azure GPU
These checks are required after the changes to 🟢 mypy
These checks are required after the changes to 🟢 installThese checks are required after the changes to Thank you for your contribution! 💜
|
|
Should be safe to merge, I'll handle the docs failure in e.g. #18718 |
| module_sharded_ctx = self.module_sharded_context() | ||
| stack = ExitStack() | ||
| if not self.zero_stage_3: | ||
| stack.enter_context(super().module_init_context(empty_init=empty_init)) |
There was a problem hiding this comment.
for completeness, it would probably also be better to apply it to this line right?
There was a problem hiding this comment.
I don't think this one matters because in the super() call, all the ctxmanagers are instantiated before any is entered
| # These operations are applied to each submodule 'bottom up' in the module hierarchy. | ||
| stack.enter_context(torch.device("meta")) | ||
| elif _TORCH_GREATER_EQUAL_1_13: | ||
| stack.enter_context(_EmptyInit(enabled=bool(empty_init))) |
There was a problem hiding this comment.
why isn't it applied in other places here?
There was a problem hiding this comment.
torch.device wont raise an exception but I missed this _EmptyInit
3 participants
What does this PR do?
ExitStackwill only__exit__its context managers after it has ben__enter__ed. This means thatctx1.__exit__never gets called.The solution is to reorder the code
Fixes #18705.
This is unreleased code. No need to mention it in the CHANGELOG
cc @Borda @carmocca @justusschock @awaelchli