`self.device` is not correct during `setup`

[carmocca]

🐛 Bug

To Reproduce

import os

from pytorch_lightning import Trainer
from pytorch_lightning.demos.boring_classes import BoringModel


class BugReportModel(BoringModel):
    def setup(self, stage=None):
        # current workaround
        assert self.trainer.strategy.root_device.type == "cuda"
        # error! it is on cpu
        assert self.device.type == "cuda"


def run():
    model = BugReportModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        fast_dev_run=True,
        accelerator="gpu",
        devices=1,
    )
    trainer.fit(model)


if __name__ == "__main__":
    run()

Expected behavior

lightning_module.device and strategy.device point to the same device

Additional context

Reported in Slack by Joshua Chia

cc @tchaton @rohitgr7 @justusschock @kaushikb11 @awaelchli @akihironitta

[awaelchli]

We cannot expect this to be true at any point in time:

lightning_module.device and strategy.device point to the same device

The value of strategy.device is available as soon as the trainer instantiation finishes. Between then and trainer.fit, it is clear that the lightning_module.device cannot possibly have the updated value already. However, we can expect that after a certain point in the hook call timeline, these two become identical. When is that point?

We know that it certainly must be after configure_sharded_model(), because only then do we know that all layers are created completely. configure_sharded_model runs after the setup() hook, so that excluses setup() from being able to meet the requirement. In essence, the sequence is this:

setup()
configure_sharded_model()
move_model_to_device()

Could one just change the order to

configure_sharded_model()
move_model_to_device()
setup()

??

I'm not so sure if a breaking change like this is worth fixing this "bug".

This argumentation assumes that the value of lightning_module.device represents the true current placement of the module layers. One could in theory set lightning_module.device equal to strategy.device without actually moving the module layers, just for the sake of having them equal early on in the setup phase of fit, but that would be misleading IMO.

[Daniel451]

Is there any update on this?

@carmocca what type of "hotfix" are you using so far? I tried it with the good old lit_model = lit_model.to(device) so far but this is quite counter-intuitive.

@awaelchli having setup() and the end would obviously fix this problem but isn't it the case that dataloaders need distribution information when the sharding happens? So if you would configure something w.r.t. dataloaders in setup() this would have to happen before configure_sharded_model(). Might be totally wrong tho, maybe @Borda has an idea?