[MPS] Fix memory leak in copy_from_mps_ (#114197)

malfet · pytorchmergebot · commit b5dd37f23efe · 2023-11-21T14:52:55.000Z
By always calling `[destBuffer release]` before leaving the scope in which it was allocated. Leak was introduced by #84928 Add regression test. Before the change: ``` % python ../test/test_mps.py -v -k test_copy_cast_no_leak --repeat 10 test_copy_cast_no_leak (__main__.TestMemoryLeak) ... FAIL ====================================================================== FAIL: test_copy_cast_no_leak (__main__.TestMemoryLeak) ---------------------------------------------------------------------- Traceback (most recent call last): File "/Users/nshulga/git/pytorch/pytorch/torch/testing/_internal/common_utils.py", line 2554, in wrapper method(*args, **kwargs) File "/Users/nshulga/git/pytorch/pytorch/build/../test/test_mps.py", line 1064, in test_copy_cast_no_leak self.assertTrue(driver_before == driver_after, f"Detected {driver_after-driver_before} bytes leak of GPU memory") AssertionError: False is not true : Detected 65536 bytes leak of GPU memory To execute this test, run the following from the base repo dir: python test/test_mps.py -k test_copy_cast_no_leak This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0 ---------------------------------------------------------------------- Ran 1 test in 1.102s FAILED (failures=1) ``` After: ``` % python ../test/test_mps.py -k test_copy_cast_no_leak --repeat 10 . ---------------------------------------------------------------------- Ran 1 test in 0.819s OK . ---------------------------------------------------------------------- Ran 1 test in 0.001s OK . ---------------------------------------------------------------------- Ran 1 test in 0.002s OK ... ``` Fixes #114096 Pull Request resolved: #114197 Approved by: https://github.com/kit1980
diff --git a/aten/src/ATen/native/mps/operations/Copy.mm b/aten/src/ATen/native/mps/operations/Copy.mm
@@ -140,8 +140,8 @@ static void copy_cast_mps(at::Tensor& dst,
 
       stream->copy_and_sync(
           tmpBuffer, destBuffer, size_to_copy, storage_byte_offset, destOffset, non_blocking, profile_id);
-      [destBuffer release];
     }
+    [destBuffer release];
   }
   if (!dst.is_same(dst_)) {
     dst_.copy_(dst, non_blocking);
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -1053,6 +1053,16 @@ def leak_gpu0():
         with self.assertRaisesRegex(RuntimeError, r"MPS driver API confirmed .+"):
             leak_gpu0()
 
+    def test_copy_cast_no_leak(self):
+        a = torch.randn(128, 128, device='mps', dtype=torch.float16)
+        torch.mps.empty_cache()
+        driver_before = torch.mps.driver_allocated_memory()
+        a = a.to(device='cpu', dtype=torch.float32)
+        a = a.to(device='mps', dtype=torch.float16)
+        torch.mps.empty_cache()
+        driver_after = torch.mps.driver_allocated_memory()
+        self.assertTrue(driver_before == driver_after, f"Detected {driver_after-driver_before} bytes leak of GPU memory")
+
 
 class TestPixelShuffle(TestCaseMPS):
     def test_pixel_shuffle_unshuffle(self):
