LightningLite not updating DeviceDtypeModuleMixin correctly

[justusschock]

🐛 Bug

When using LightningLite and transferring the _LiteModule to cpu, attributes of DeviceDtypeModule are not updated.

To Reproduce

class SomeDummy(DeviceDtypeModuleMixin):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Linear(1,1)

class MyClass(LightningLite):
    def run(self):
        model = SomeDummy()
        model, optimiser = self.setup(model, torch.optim.Adam(model.parameters()))

        #  do some stuff
        
        # now clean up gpu memory for later stages
        model.cpu()
        assert str(model.module.device) == 'cpu'

MyClass(accelerator='gpu', devices=1).run()

Expected behavior

model.module.device should be Cpu

Additional context

Could probably be solved by using DeviceDtypeModuleMixin as base class for the _LiteModule since this is an issue with the to function only calling _apply on all child tensors instead of calling .to on every child module.

cc @carmocca @justusschock @awaelchli

[awaelchli]

That's an issue we have even without Lite, e.g., it could occur with torchmetrics too. This example shows it:

import torch

from pytorch_lightning.core.mixins import DeviceDtypeModuleMixin


class SomeDummy(DeviceDtypeModuleMixin):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Linear(1, 1)


class WrapperModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.module = SomeDummy()  # this could be a torchmetric


w = WrapperModule().cuda()
print(w.module.device)  # prints cpu !!! should be cuda:0
w.cpu()
print(w.module.device)  # prints cpu

As you said, the only solution is to add the DeviceDtypeModuleMixin to the base of the wrapper class because of how child modules get called.

[awaelchli]

I noticed, in torchmetrics they abandoned the DeviceDtypeMixin and override the _apply method:

https://github.com/PyTorchLightning/metrics/blob/93cb842f24d15804dd2e7677ca7fc6631b234773/torchmetrics/metric.py#L466-L490

[justusschock]

@awaelchli yes, but without Lite this is not our concern :D

From our offline chat (for completeness):
For torch metrics removing it is fine since they don't nest modules there. And overriding the _apply is exactly what the mixin does :)