Lightning-AI / Lightning-AI/pytorch-lightning

Calling `trainer.fit` twice with spawn strategies won't work as expected

Open
#18,775 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug priority: 1 strategy: ddp strategy: xla ver: 2.0.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.