Lightning-AI / Lightning-AI/pytorch-lightning
Remove Trainer reference from lightning module and datamodule
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
Motivation: Remove trainer reference from LightningModule and LightningDataModule.
Benefits:
- Better backwards compatibility: it should be easier to use the LightningModule of one version of Lightning with the Trainer from another version
- Better API clarity: users today who call Trainer functions like fit/validate/predict inside of the LightningModule will run into issues with state management. The Trainer assumes these top-level APIs are not called from within another trainer context.
Possible solution: think about introducing a TrainerContext / TrainerState object to pass state and deprecate references to the trainer. This would be mostly read-only data that the LightningModule could leverage for various settings like progress (epoch/step count), distributed setting (global rank, local rank, data parallel/ddp/deepspeed/etc) and more.
The ambition is the lightning module has a tightly controlled view of the trainer, while the trainer has full insight into the module.
cc @borda @tchaton @justusschock @awaelchli
Contributor guide
Assessment
This issue has not been assessed yet.