Lightning-AI / Lightning-AI/pytorch-lightning
Introduce training and eval modes for the strategy
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Some of our strategies need to be set up differently depending on whether the model is being trained or just evaluated (`.fit` vs `.test`).
Some examples:
https://github.com/Lightning-AI/lightning/blob/859a228a915893f14759dda6d35f6050fe6df382/src/pytorch_lightning/strategies/ddp.py#L170-L173
https://github.com/Lightning-AI/lightning/blob/859a228a915893f14759dda6d35f6050fe6df382/src/pytorch_lightning/strategies/ddp.py#L357-L363
https://github.com/Lightning-AI/lightning/blob/859a228a915893f14759dda6d35f6050fe6df382/src/pytorch_lightning/strategies/ipu.py#L146-L149
This requires us to pass in the entire Trainer instance just to check the stage. Unifying this with the base strategies from Lite is not possible, because of the dependency on Trainer.
### Pitch
Introduce `strategy.train()` and `strategy.eval()` with a similar mechanism as in `nn.Module.train/eval()`. This could be a method or a boolean attribute. The call to change this state would happen in the training loop/the trainer.
For now, this proposal would only be for the strategy definition in Trainer. This helps us standardize the interface for both strategies, while maintaining the flexibility of adding Trainer-specific logic outside Lite.
### Alternatives
Alternatively, the state could be passed in to the individual methods of the strategy.
### Additional context
_No response_
cc @borda @justusschock @carmocca
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 by reading the referenced strategy examples in src/pytorch_lightning/strategies/ddp.py and src/pytorch_lightning/strategies/ipu.py, then trace how the Trainer or training loop currently determines the stage. Done means strategies have a standardized training/evaluation state without requiring the entire Trainer just to inspect it.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100