Fix load_from_checkpoint to return model on correct device

[ryan597]

What does this PR do?

Fixes #17304

When using

model = BoringModel.load_from_checkpoint('path.ckpt', map_location='cuda')

the returned model is always on the CPU due to _load_from_checkpoint not using the map_location on the returned object.

https://github.com/Lightning-AI/lightning/blob/fd4697c62c059fc7b9946e84d91625ecb6efdbe5/src/lightning/pytorch/core/saving.py#L51-L92

This PR returns the created model with the correct mapped location.

Also create tests for checking that the model is loaded onto GPU after the checkpoint load. I have included tests for mapping to the CPU for completeness although these do not fail currently fail, only the GPU tests do.

Created tests failing on master

All tests checking for model to be on cuda will fail on current master branch.
All tests checking for model to be on cpu will pass on current branch, but included for completeness.

import pytest
import torch

import lightning.pytorch as pl
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.demos.boring_classes import BoringModel
from tests_pytorch.helpers.runif import RunIf


def create_boring_checkpoint(tmpdir, accelerator="gpu", testmodel=BoringModel):
    checkpoint_callback = ModelCheckpoint(dirpath=tmpdir, filename="checkpoint")
    model = testmodel()
    trainer = pl.Trainer(
        devices=1,
        accelerator=accelerator,
        max_epochs=1,
        enable_model_summary=False,
        enable_progress_bar=False,
        callbacks=[checkpoint_callback],
    )
    trainer.fit(model)


@pytest.mark.parametrize(
    "map_location", (None, "cpu", torch.device("cpu"), lambda storage, loc: storage, {"cpu": "cpu"})
)
def test_load_from_checkpoint_map_location_cpu(tmpdir, map_location, testmodel=BoringModel):
    create_boring_checkpoint(tmpdir, accelerator="cpu", testmodel=testmodel)
    model = testmodel.load_from_checkpoint(f"{tmpdir}/checkpoint.ckpt", map_location=map_location)
    assert model.device.type == "cpu"


@RunIf(min_cuda_gpus=1)
@pytest.mark.parametrize(
    "map_location", (None, "cuda", torch.device("cuda"), lambda storage, loc: storage.cuda(), {"cpu": "cuda"})
)
def test_load_from_checkpoint_map_location_gpu(tmpdir, map_location, testmodel=BoringModel):
    create_boring_checkpoint(tmpdir, accelerator="gpu", testmodel=testmodel)
    model = testmodel.load_from_checkpoint(f"{tmpdir}/checkpoint.ckpt", map_location=map_location)
    assert model.device.type == "cuda"


@RunIf(min_cuda_gpus=1)
@pytest.mark.parametrize("map_location", ("cpu", torch.device("cpu"), lambda storage, loc: storage, {"cuda": "cpu"}))
def test_load_from_checkpoint_map_location_gpu_to_cpu(tmpdir, map_location, testmodel=BoringModel):
    create_boring_checkpoint(tmpdir, accelerator="gpu", testmodel=testmodel)
    model = testmodel.load_from_checkpoint(f"{tmpdir}/checkpoint.ckpt", map_location=map_location)
    assert model.device.type == "cpu"


@RunIf(min_cuda_gpus=1)
@pytest.mark.parametrize(
    "map_location", ("cuda", torch.device("cuda"), lambda storage, loc: storage.cuda(), {"cpu": "cuda"})
)
def test_load_from_checkpoint_map_location_cpu_to_gpu(tmpdir, map_location, testmodel=BoringModel):
    create_boring_checkpoint(tmpdir, accelerator="cpu", testmodel=testmodel)
    model = testmodel.load_from_checkpoint(f"{tmpdir}/checkpoint.ckpt", map_location=map_location)
    assert model.device.type == "cuda"

What Can be Improved

When the object is created on the GPU without setting map_location (i.e. when the model checkpoint is from the GPU, shouldn't it automatically load onto the GPU as when you use torch.load("boring.pth")?).

import torch
import torch.nn as nn

class TestModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.fc = nn.Linear(32, 2)

    def forward(self, x):
        return self.fc(x)

testmodel = TestModel().to("cuda")

torch.save(testmodel, "torch.pth")
testmodel = torch.load("torch.pth")
print(next(testmodel.parameters()).device)  # Prints cuda:0

import lightning.pytorch as pl
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.demos.boring_classes import BoringModel

checkpoint_callback = ModelCheckpoint(dirpath='', filename="boring")

model = BoringModel()
trainer = pl.Trainer(
    max_epochs=1,
    enable_progress_bar=False,
    enable_model_summary=False,
    callbacks=[checkpoint_callback],
    accelerator="gpu",
    devices=1
)
trainer.fit(model)

model = BoringModel.load_from_checkpoint("boring.ckpt") 
print(model.device.type)  # Should print cuda but prints cpu

So for this the issue is when the object is created in _load_state() it is not created on or moved to the GPU, thus when map_location is set to lambda storage, loc: storage (as it is when map_location=None) it just returns the object on the CPU.

If I try to retrieve this information from the checkpoint['state_dict'], it requires some ugly work that I'm not sure will work in all cases as I have only tested with the BoringModel so far

...
  if issubclass(cls, pl.LightningModule):
      loaded_location = list(checkpoint["state_dict"].items())[0][1].device  # This could be done better?
      if map_location is None:
          map_location = loaded_location
  
      storage = _load_state(cls, checkpoint, strict=strict, **kwargs)
      restore_location = torch.serialization._get_restore_location(map_location)
      if isinstance(map_location, dict) and isinstance(storage, pl.LightningModule):
          return restore_location(storage, str(storage.device))
  
      return restore_location(storage, map_location)

This works (for the BoringModel anyways), but I'm sure it can be done much better if someone has better ideas.

Failing tests / Breaking changes

tests/tests_pytorch/strategies/test_fsdp.py contains tests which are failing. It loads using checkpoints with the changed function load_from_checkpoint, however in the assertion, it moves one of the tensors to the cpu resulting in the assertion error because the loaded checkpoint is now correctly on cuda.

https://github.com/Lightning-AI/lightning/blob/fd4697c62c059fc7b9946e84d91625ecb6efdbe5/tests/tests_pytorch/strategies/test_fsdp.py#L142-L152

In changing this assertion to keep both state_dicts on their devices all tests pass.

- assert torch.equal(ddp_param.float().cpu(), shard_param) 
+ assert torch.equal(ddp_param, shard_param) 

If there's other cases where users have similar conditions or rely on the current behavior this PR could break them, although it would only be unexpected for the case of map_location=None.

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified