Lightning-AI / Lightning-AI/pytorch-lightning

Remove the `optimizer_to_device` logic if possible

Open
#20,165 3 comments 1 reaction 0 assignees View on GitHub
checkpointing performance refactor
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.