Lightning-AI / Lightning-AI/pytorch-lightning
Calling `trainer.fit` twice with spawn strategies won't work as expected
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Since data in the spawned region is not shared with the main process, the spawn launcher saves a checkpoint of the weights before finishing that is then loaded on the main process:
https://github.com/Lightning-AI/lightning/blob/984f49f7195ddc67e961c7c498ee6e19fc0cecb5/src/lightning/pytorch/strategies/launchers/multiprocessing.py#L190-L195 https://github.com/Lightning-AI/lightning/blob/984f49f7195ddc67e961c7c498ee6e19fc0cecb5/src/lightning/pytorch/strategies/launchers/multiprocessing.py#L162-L168
This means that the optimizer states are not loaded, as well as any other state in the trainer.
This isn't a problem with calling `test/validate/predict` after `fit`.
### Solution
Since this is a silent correctness issue. We should raise an error in the short term.
The launcher can check if `fit` was called and is getting called again, and then raise a `NotImplementedError`.
In the longer term, we can save a full checkpoint that contains all the relevant data and then lift this restriction.
cc @tchaton @justusschock @awaelchli @carmocca @JackCaoG @Liyang90 @gkroiz
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/pytorch/strategies/launchers/multiprocessing.py, especially lines 162-168 and 190-195, and trace how the launcher saves and restores state around fit. Reproduce two consecutive trainer.fit calls with a spawn strategy, then verify that the second call raises NotImplementedError while test, validate, and predict after fit remain unaffected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100