Lightning-AI / Lightning-AI/pytorch-lightning
Remove the `optimizer_to_device` logic if possible
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Outline & Motivation
The trainer uses a function `optimizer_to_device` here:
https://github.com/Lightning-AI/pytorch-lightning/blob/631911c00413ad028e2887d83eb264cb4822097e/src/lightning/pytorch/strategies/strategy.py#L160-L161
In #19955 an issue was raised that the function moved the "step" parameter in the optimizer state to the CUDA device, causing device-to-host syncs during optimizer.step() because the "step" tensor was expected to remain on CPU. #20019 fixed this with special treatment of that key. However, good arguments were made in #19955 that this `optimizer_to_device` shouldn't even be necessary in the first place (https://github.com/Lightning-AI/pytorch-lightning/issues/19955#issuecomment-2197353178).
### Pitch
Remove `optimizer_to_device` and show that it is redundant by running the tests. We will still need a `optimizer_to_cpu` for teardown.
### Additional context
_No response_
cc @justusschock @awaelchli @borda
Contributor guide
Research direction
Start by reading optimizer_to_device in src/lightning/pytorch/strategies/strategy.py and the context from issues #19955 and #20019. Run the relevant tests to verify that removing optimizer_to_device is safe while retaining optimizer_to_cpu for teardown; done means the test suite passes without the device-sync issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100