Disable non blocking to device with MPS

src/pytorch_lightning/utilities/apply_func.py

@@ -15,6 +15,7 @@
 
 import dataclasses
 import operator
+import platform
 from abc import ABC
 from collections import defaultdict, OrderedDict
 from collections.abc import Mapping, Sequence
@@ -27,7 +28,7 @@
 from torch import Tensor
 
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.imports import _compare_version, _TORCHTEXT_LEGACY
+from pytorch_lightning.utilities.imports import _compare_version, _TORCH_GREATER_EQUAL_1_12, _TORCHTEXT_LEGACY
 from pytorch_lightning.utilities.warnings import rank_zero_deprecation
 
 if _TORCHTEXT_LEGACY:
@@ -40,6 +41,10 @@
 
 
 _CPU_DEVICES = ("cpu", torch.device("cpu"))
+if _TORCH_GREATER_EQUAL_1_12 and torch.backends.mps.is_available() and platform.processor() in ("arm", "arm64"):
+    _MPS_DEVICES = ("mps", torch.device("mps"))
+else:
+    _MPS_DEVICES = ("mps",)
 
 
 def to_dtype_tensor(
@@ -343,7 +348,9 @@ def batch_to(data: Any) -> Any:
 
         kwargs = {}
         # Don't issue non-blocking transfers to CPU
-        if isinstance(data, Tensor) and device not in _CPU_DEVICES:
+        # Don't issue non-blocking transfers to MPS due to race condition bug:
+        # https://github.com/pytorch/pytorch/issues/83015
+        if isinstance(data, Tensor) and device not in _CPU_DEVICES and device not in _MPS_DEVICES:
             kwargs["non_blocking"] = True
         data_output = data.to(device, **kwargs)
         if data_output is not None:

tests/tests_pytorch/accelerators/test_mps.py

@@ -162,3 +162,14 @@ def to(self, *args, **kwargs):
 
     assert batch.text.type() == "torch.mps.LongTensor"
     assert batch.label.type() == "torch.mps.LongTensor"
+
+
+@RunIf(mps=True)
+def test_data_is_not_changed_after_move_to_mps_device():
+    trainer = Trainer(accelerator="mps", devices=1)
+    x = torch.zeros([10, 10])
+    device = torch.device("mps")
+
+    for _ in range(1000):
+        x_mps = trainer.strategy.batch_to_device(x.clone(), device)
+        torch.testing.assert_close(x_mps, x)