Set _DeviceDtypeModuleMixin _device from torch's default device function (#21164)

* fix: set `_DeviceDtypeModuleMixin` _device from torch's default device function.
* add torch 2.2.2 or below compatibility
* add import for `_TORCH_GREATER_EQUAL_2_3
* fix: restore torch default dtype once test_submodules_context_device_and_dtype is finished.

---------

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Bhimraj Yadav <[email protected]>
Co-authored-by: Nicki Skafte Detlefsen <[email protected]>

Author: 4 people

src/lightning/fabric/CHANGELOG.md

@@ -26,6 +26,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     [#21057](https://github.com/Lightning-AI/pytorch-lightning/pull/21057), [#21093](https://github.com/Lightning-AI/pytorch-lightning/pull/21093))
 
 
+- Set `_DeviceDtypeModuleMixin._device` from torch's default device function ([#21164](https://github.com/Lightning-AI/pytorch-lightning/pull/21164))
+
+
 ### Fixed
 
 - Fixed with adding a missing device id for pytorch 2.8 ([#21105](https://github.com/Lightning-AI/pytorch-lightning/pull/21105))

src/lightning/fabric/utilities/device_dtype_mixin.py

@@ -18,14 +18,18 @@
 from torch.nn import Module
 from typing_extensions import Self, override
 
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_3
+
 
 class _DeviceDtypeModuleMixin(Module):
     __jit_unused_properties__: list[str] = ["device", "dtype"]
 
     def __init__(self) -> None:
         super().__init__()
         self._dtype: Union[str, torch.dtype] = torch.get_default_dtype()
-        self._device = torch.device("cpu")
+        # Workarounds from the original pytorch issue:
+        # https://github.com/pytorch/pytorch/issues/115333#issuecomment-1848449687
+        self._device = torch.get_default_device() if _TORCH_GREATER_EQUAL_2_3 else torch.empty(0).device
 
     @property
     def dtype(self) -> Union[str, torch.dtype]:

tests/tests_fabric/utilities/test_device_dtype_mixin.py

@@ -2,6 +2,7 @@
 import torch
 from torch import nn as nn
 
+from lightning.fabric.plugins.precision.utils import _DtypeContextManager
 from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
 from tests_fabric.helpers.runif import RunIf
 
@@ -50,6 +51,30 @@ def test_submodules_device_and_dtype(dst_device_str, dst_type):
     assert model.dtype == model.module.module.dtype == dst_type
 
 
+@pytest.mark.parametrize(
+    "dst_device_str",
+    [
+        "cpu",
+        pytest.param("cuda:0", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("mps:0", marks=RunIf(mps=True)),
+    ],
+)
+@pytest.mark.parametrize(
+    "dst_type",
+    [
+        torch.float,
+        pytest.param(torch.half, marks=RunIf(mps=False)),
+        pytest.param(torch.double, marks=RunIf(mps=False)),
+    ],
+)
+def test_submodules_context_device_and_dtype(dst_device_str, dst_type):
+    dst_device = torch.device(dst_device_str)
+    with _DtypeContextManager(dst_type), dst_device:
+        model = TopModule()
+    assert model.device == dst_device
+    assert model.dtype == dst_type
+
+
 @pytest.mark.parametrize(
     "device",
     [

tests/tests_pytorch/trainer/test_trainer.py

@@ -2107,6 +2107,22 @@ def test_init_module_context(monkeypatch):
     strategy.tensor_init_context.reset_mock()
 
 
+@pytest.mark.parametrize(
+    ("target_device", "accelerator", "devices"),
+    [
+        ("cpu", "cpu", "auto"),
+        pytest.param("cuda:0", "gpu", [0], marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("cuda:1", "gpu", [1], marks=RunIf(min_cuda_gpus=2)),
+    ],
+)
+def test_init_module_device_type(target_device, accelerator, devices):
+    """Test that the strategy returns the context manager for initializing the module."""
+    trainer = Trainer(accelerator=accelerator, devices=devices)
+    with trainer.init_module():
+        model = BoringModel()
+        assert model.device == torch.device(target_device)
+
+
 def test_expand_home_trainer():
     """Test that the dirpath gets expanded if it contains `~`."""
     home_root = Path.home()