Lightning-AI / Lightning-AI/pytorch-lightning

Remove Trainer reference from lightning module and datamodule

Open
#7,315 2 comments 0 reactions 1 assignee Claimed by @ananthsub View on GitHub
design discussion feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.