Lightning-AI / Lightning-AI/pytorch-lightning

Introduce training and eval modes for the strategy

Open
#16,226 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature strategy
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.