diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 8775c149fb365..b3d3c02ce64f2 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added `log_key_prefix` parameter to `LearningRateMonitor` callback for prefixing logged metric names ([#21612](https://github.com/Lightning-AI/pytorch-lightning/issues/21612)) +- Added a `dtype` argument to `LightningModule.load_from_checkpoint` to cast the loaded model's floating-point parameters at load time ([#20833](https://github.com/Lightning-AI/pytorch-lightning/issues/20833)) + ### Changed - diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py index 289f9aa3da7cc..6f4ed257bd0df 100644 --- a/src/lightning/pytorch/core/module.py +++ b/src/lightning/pytorch/core/module.py @@ -1722,6 +1722,7 @@ def load_from_checkpoint( hparams_file: Optional[_PATH] = None, strict: Optional[bool] = None, weights_only: Optional[bool] = None, + dtype: Optional[torch.dtype] = None, **kwargs: Any, ) -> Self: r"""Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint it stores the arguments @@ -1760,6 +1761,11 @@ def load_from_checkpoint( ``weights_only=False``. If loading checkpoint from an untrusted source, we recommend using ``weights_only=True``. For more information, please refer to the `PyTorch Developer Notes on Serialization Semantics `_. + dtype: If provided, cast the loaded model's floating-point parameters and buffers to this + :class:`torch.dtype` (e.g. ``torch.float16``), equivalent to calling ``.to(dtype)`` on the + returned model. Non-floating-point tensors (such as integer buffers) are left unchanged. + Defaults to ``None``, which keeps the dtype stored in the checkpoint. Note that downcasting + (e.g. from ``float32`` to ``float16``) reduces numerical precision. \**kwargs: Any extra keyword args needed to init the model. Can also be used to override saved hyperparameter values. @@ -1816,6 +1822,7 @@ def load_from_checkpoint( hparams_file, strict, weights_only, + dtype=dtype, **kwargs, ) return cast(Self, loaded) diff --git a/src/lightning/pytorch/core/saving.py b/src/lightning/pytorch/core/saving.py index 391e9dd5d0f25..9623bf274b080 100644 --- a/src/lightning/pytorch/core/saving.py +++ b/src/lightning/pytorch/core/saving.py @@ -57,6 +57,7 @@ def _load_from_checkpoint( hparams_file: Optional[_PATH] = None, strict: Optional[bool] = None, weights_only: Optional[bool] = None, + dtype: Optional["torch.dtype"] = None, **kwargs: Any, ) -> Union["pl.LightningModule", "pl.LightningDataModule"]: map_location = map_location or _default_map_location @@ -98,6 +99,10 @@ def _load_from_checkpoint( device = next((t for t in state_dict.values() if isinstance(t, torch.Tensor)), torch.tensor(0)).device assert isinstance(model, pl.LightningModule) + if dtype is not None: + # cast floating-point parameters/buffers to the requested dtype; + # `torch.nn.Module.to` leaves non-floating-point tensors untouched + return model.to(device=device, dtype=dtype) return model.to(device) raise NotImplementedError(f"Unsupported {cls}") diff --git a/tests/tests_pytorch/core/test_saving.py b/tests/tests_pytorch/core/test_saving.py index 8e1e9584a7c68..060649feda123 100644 --- a/tests/tests_pytorch/core/test_saving.py +++ b/tests/tests_pytorch/core/test_saving.py @@ -122,6 +122,20 @@ def test_load_from_checkpoint_warn_on_empty_state_dict(tmp_path): assert model.device.type == "cpu" +@pytest.mark.parametrize("dtype", [torch.float16, torch.float64]) +def test_load_from_checkpoint_dtype(dtype, tmp_path): + """Test that `load_from_checkpoint` casts floating-point parameters to the requested dtype.""" + create_boring_checkpoint(tmp_path, BoringModel(), accelerator="cpu") + + # default: dtype is preserved as stored in the checkpoint (float32) + model = BoringModel.load_from_checkpoint(f"{tmp_path}/checkpoint.ckpt") + assert all(p.dtype == torch.float32 for p in model.parameters()) + + # explicit dtype casts floating-point parameters + model = BoringModel.load_from_checkpoint(f"{tmp_path}/checkpoint.ckpt", dtype=dtype) + assert all(p.dtype == dtype for p in model.parameters()) + + @pytest.mark.parametrize( ("strict", "strict_loading", "expected"), [