feat: add dtype option to LightningModule.load_from_checkpoint

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
 
 -

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 <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
+            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)

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}")

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"),
     [