[DOC] Document checkpoint resuming on lightning 2.6 and use pytorch_forecasting Tuner in examples

README.md

@@ -86,7 +86,7 @@ from lightning.pytorch.loggers import TensorBoardLogger
 from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor
 # import dataset, network to train and metric to optimize
 from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer, QuantileLoss
-from lightning.pytorch.tuner import Tuner
+from pytorch_forecasting.tuning import Tuner
 
 # load data: this is pandas dataframe with at least a column for
 # * the target (what you want to predict)

docs/source/faq.rst

@@ -78,6 +78,12 @@ Training models
   the learning rate finder. To tune hyperparameters, the `optuna package <https://optuna.org/>`_
   is a great place to start with.
 
+* **Why does resuming training from a checkpoint with ckpt_path fail with an UnpicklingError?**
+
+  From ``lightning`` 2.6, pass ``weights_only=False`` together with ``ckpt_path`` for
+  checkpoints that you trust. See
+  :ref:`resuming from a checkpoint <models_resume_checkpoint>`.
+
 
 Interpreting models
 -------------------

docs/source/getting-started.rst

@@ -63,8 +63,8 @@ Example
 
     import lightning.pytorch as pl
     from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor
-    from lightning.pytorch.tuner import Tuner
     from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer
+    from pytorch_forecasting.tuning import Tuner
 
     # load data
     data = ...

docs/source/models.rst

@@ -45,8 +45,8 @@ One example of using :py:class:`~data.timeseries.TimeSeriesDataSet` and models i
 
     import lightning.pytorch as pl
     from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor
-    from lightning.pytorch.tuner import Tuner
     from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer
+    from pytorch_forecasting.tuning import Tuner
 
     # load data
     data = ...
@@ -118,6 +118,40 @@ One example of using :py:class:`~data.timeseries.TimeSeriesDataSet` and models i
         tft, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader,
     )
 
+.. _models_resume_checkpoint:
+
+Resuming from a checkpoint
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To resume training, pass the checkpoint as ``ckpt_path`` to ``trainer.fit``. From
+``lightning`` 2.6, also pass ``weights_only=False``:
+
+.. code-block:: python
+
+    trainer.fit(
+        tft,
+        train_dataloaders=train_dataloader,
+        val_dataloaders=val_dataloader,
+        ckpt_path="path/to/checkpoint.ckpt",
+        weights_only=False,
+    )
+
+The same applies to ``trainer.validate`` and ``trainer.test``. ``weights_only`` controls
+how the checkpoint is loaded; it does not select a checkpoint or resume training by
+itself. A v1 PyTorch Forecasting checkpoint contains more than tensors: the fitted encoders and
+normalizers of the dataset as well as the loss and metrics are pickled into it.
+``lightning`` 2.6 passes ``weights_only`` on to ``torch.load``, which defaults to
+``weights_only=True`` from ``torch`` 2.6 and then rejects these objects. Only load
+checkpoints from trusted sources this way, as unpickling can execute arbitrary code.
+``lightning`` 2.4 and 2.5 already load full checkpoints and do not have the keyword.
+
+``load_from_checkpoint`` of v1 models defaults to ``weights_only=False``, and ``lr_find``
+of ``pytorch_forecasting.tuning.Tuner`` loads with it. With ``lightning`` 2.6 and
+``torch`` 2.6, Lightning's ``Tuner.lr_find`` and the ``LearningRateFinder`` callback fail,
+as they restore their temporary checkpoint without ``weights_only``; passing it to
+``trainer.fit`` does not reach the callback. Use ``pytorch_forecasting.tuning.Tuner``
+as in the example above.
+
 
 Selecting an architecture
 --------------------------

docs/source/tutorials/stallion.ipynb

@@ -2062,7 +2062,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "PyTorch Lightning automatically checkpoints training and thus, we can easily retrieve the best model and load it.\n"
+    "PyTorch Lightning automatically checkpoints training and thus, we can easily retrieve the best model and load it.\n",
+    "\n",
+    "To resume training from a checkpoint instead, pass it as `ckpt_path` to `trainer.fit`. From `lightning` 2.6 with `torch` 2.6, the checkpoint is loaded with `weights_only=True` by default, which rejects the encoders, normalizers and metrics that PyTorch Forecasting pickles into it. For checkpoints that you trust, also pass `weights_only=False`, e.g. `trainer.fit(tft, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader, ckpt_path=best_model_path, weights_only=False)`. The same applies to `trainer.validate` and `trainer.test`. On `lightning` 2.6 and later, `load_from_checkpoint` below already loads with `weights_only=False`.\n"
    ]
   },
   {

tests/test_models/test_temporal_fusion_transformer.py

@@ -308,6 +308,65 @@ def test_pickle(model):
     pickle.loads(pkl)  # noqa: S301
 
 
+@pytest.mark.skipif(
+    _check_soft_dependencies("lightning<2.6", severity="none"),
+    reason="Trainer methods accept weights_only only from lightning 2.6",
+)
+def test_trainer_ckpt_path_weights_only(dataloaders_with_covariates, tmp_path):
+    train_dataloader = dataloaders_with_covariates["train"]
+    val_dataloader = dataloaders_with_covariates["val"]
+
+    def make_model():
+        return TemporalFusionTransformer.from_dataset(
+            train_dataloader.dataset,
+            hidden_size=4,
+            attention_head_size=1,
+            hidden_continuous_size=2,
+        )
+
+    trainer_kwargs = dict(
+        accelerator="cpu",
+        limit_train_batches=2,
+        limit_val_batches=1,
+        limit_test_batches=1,
+        logger=False,
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+    )
+    trainer = pl.Trainer(max_epochs=1, **trainer_kwargs)
+    trainer.fit(make_model(), train_dataloaders=train_dataloader)
+    ckpt_path = tmp_path / "model.ckpt"
+    trainer.save_checkpoint(ckpt_path)
+
+    epochs_run = []
+
+    class CountEpochs(pl.Callback):
+        def on_train_epoch_start(self, trainer, pl_module):
+            epochs_run.append(trainer.current_epoch)
+
+    trainer = pl.Trainer(max_epochs=2, callbacks=[CountEpochs()], **trainer_kwargs)
+    trainer.fit(
+        make_model(),
+        train_dataloaders=train_dataloader,
+        ckpt_path=ckpt_path,
+        weights_only=False,
+    )
+    assert trainer.current_epoch == 2
+    assert epochs_run == [1]
+    trainer.validate(
+        make_model(),
+        dataloaders=val_dataloader,
+        ckpt_path=ckpt_path,
+        weights_only=False,
+    )
+    trainer.test(
+        make_model(),
+        dataloaders=val_dataloader,
+        ckpt_path=ckpt_path,
+        weights_only=False,
+    )
+
+
 @pytest.mark.parametrize(
     "kwargs", [dict(mode="dataframe"), dict(mode="series"), dict(mode="raw")]
 )