Bug description

Usually, I would expect Fabric.load(ckpt_path, state) to operate in-place, i.e. for it to call the load_state_dict methods of the loadable objects in state. This does indeed happen when state is a flat dictionary. However, when state is nested (e.g. because one handles multiple models with corresponding optimizers), this changes all of a sudden. Specifically, the modules in the state dictionary will reference different objects after the call to load. I have provided a minimal reproducible example below.

To me this is definitely unexpected behavior and will lead to problems when we are using the references to the old modules when resuming training (e.g. imagine in the example below we use model during training and not state["generator"]["model"]). Currently one needs to explicitly update the old references, which is cumbersome and error-prone, especially when users are not aware of this behavior.

What version are you seeing the problem on?

v2.3

How to reproduce the bug

from lightning import Fabric
from torch import nn

if __name__ == "__main__":
    fabric = Fabric(accelerator="cpu")
    model = nn.Linear(2, 2)
    model = fabric.setup(model)
    state_flat = {"model": model}
    state_nested = {"generator": {"model": model, "optim": None}}
    fabric.save("flat.pt", state_flat)
    fabric.save("nested.pt", state_nested)

    fabric.load("flat.pt", state_flat, strict=True)
    assert model is state_flat["model"]     # This is fine

    fabric.load("nested.pt", state_nested, strict=True)
    assert model is state_nested["generator"]["model"]      # This will fail

Error messages and logs

No response

Environment

More info

No response