[checkpoint] Resolve 2 different checkpoint loading paths across fit vs validate/test/predict

[ananthsub]

Proposed refactoring or deprecation

Consolidate checkpoint loading code across fit and validate/test/predict

Motivation

We have 2 different code paths for checkpoint loading

1. Trainer.fit and the constructor argument resume_from_checkpoint
2. ckpt_path passed to Trainer.validate/test/predict

Offering multiple code paths here risks divergence. Lightning must ensure a consistent experience for checkpoint loading across these different entry points.

Background

These are the paths today.

Trainer.fit:

• Trainer constructor is initialized with resume_from_checkpoint: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L143
• Which initializes the checkpoint connector with resume_from_checkpoint: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L383
• Which sets this field in the connector: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L38
• In _run we call Trainer._restore_modules_and_callbacks: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L960-L962
• Which explicitly checks if we're fitting, and if so, first loads the checkpoint, before restoring the model

Trainer.validate/test/predict:

• ckpt_path is an argument to the function
• There's a common helper to validate the checkpoint path: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L1200
• A property is set on the trainer for the corresponding path: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L671-L673
• We use an umbrella property to work across validate/test/predict: https://github.com/PyTorchLightning/pytorch-lightning/blob/6e124e7207f6459cb43f540cfb5a1c6cc9b00f7a/pytorch_lightning/trainer/properties.py#L626-L633
• There are checks for loading based on if this property is set inside Trainer._run
  -- https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L950-L951
  -- https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L1006-L1007

Notably:

• trainer.fit calls restore_model whereas trainer.validate/test/predict calls restore_model_weights in the checkpoint connector. https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L116-L150
• Load checkpoint weights is called before the accelerator setup while restore model is called after: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L950-L951r: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L960-L962

Pitch

1. Unify in CheckpointConnector
   CheckpointConnector.restore_model is a superset of CheckpointConnector.restore_model_weights which suggests we don't need both.

2. Unify in trainer.py

The overall sequence in _run can be shared:

• Load & validate the checkpoint dict through the training type plugin
• Load the datamodule state
• Load the LightningModule state through the training type plugin
• Load the callback states
• Load the loop states
• If fitting, load the precision, optimizer, and LR scheduler states

We have 5 different properties exposed for the checkpoint path to resume from (excluding HPC stuff):

• trainer.resume_from_checkpoint:
• trainer.validated_ckpt_path: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L671-L674
• trainer.tested_ckpt_path: https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L760-L762
• trainer.predicted_ckpt_path : https://github.com/PyTorchLightning/pytorch-lightning/blob/a079d7fccc0a9be25b40296f2a348c4b4f40c8cf/pytorch_lightning/trainer/trainer.py#L843-L845
• An umbrella property across these 3: trainer._ckpt_path: https://github.com/PyTorchLightning/pytorch-lightning/blob/6e124e7207f6459cb43f540cfb5a1c6cc9b00f7a/pytorch_lightning/trainer/properties.py#L626-L633
  -- This is also inconsistent: the attributes initialized in validate/test/predict are public, while _ckpt_path is private. Why?

It's unclear what the lifecycle of these properties should be. Do successive calls to validate/test/predict end up relying on this?

Proposal:

• Pass the checkpoint path to resume from as an argument to Trainer._run to avoid our dependency on these properties
• Deprecate these properties

3. Unify in the Trainer API
   Proposal: deprecate resume_from_checkpoint from the Trainer constructor, and add a new argument ckpt_path to Trainer.fit . This provides API consistency with validate/test/predict

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning

• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[ananthsub]

paraphrasing @awaelchli

Remove the trainer argument will greatly impact usability and configurability through CLI. To that you will counter argue that Trainer should not be a configuration system. Still, sad to see that we are just taking it away.

but how does the CLI work today for passing ckpt_path to validate, test, or predict ? why is this any different from that?

Related to that I was hoping we could take this one up some day: #5339

I commented there with what could be viewed as an extension to the trainer argument discussion here

[carmocca]

but how does the CLI work today for passing ckpt_path to validate, test, or predict ? why is this any different from that?

In the CLI you specify the fit checkpoint with

python script.py fit --trainer.resume_from_checkpoint=...

and the test checkpoint with:

python script.py test ckpt_path=...

This PR would make it consistent.

---

I personally think there is a lot of value in doing this - if this wasn't pursued then I'd suggest renaming resume_from_checkpoint to fit_ckpt_path to address the confusion