Add context manager to save tensors on CPU

Fix #57100

ghstack-source-id: 9285a37caf37f90fd65b309235d7027504ce4ca5
Pull Request resolved: #61928

test/test_autograd.py

@@ -5900,6 +5900,30 @@ def test_default_saved_variable_hooks_double_backward(self):
         finally:
             torch.autograd.graph.reset_saved_tensors_default_hooks()
 
+    def test_graph_save_on_cpu(self):
+        a = torch.randn(5, requires_grad=True)
+        with torch.autograd.graph.save_on_cpu():
+            y = a * a
+        self.assertEqual(a, y.grad_fn._saved_self)
+        self.assertEqual(a, y.grad_fn._saved_other)
+        y.sum().backward()
+        self.assertEqual(2 * a, a.grad)
+
+    @unittest.skipIf(not TEST_CUDA, "test requires CUDA")
+    def test_graph_save_on_cpu_cuda(self):
+        a = torch.randn(5, requires_grad=True, device="cuda")
+        with torch.autograd.graph.save_on_cpu():
+            y = a * a
+        self.assertTrue(y.is_cuda)
+        before = CudaMemoryLeakCheck.get_cuda_memory_usage()
+        self.assertTrue(y.grad_fn._saved_self.is_cuda)
+        after = CudaMemoryLeakCheck.get_cuda_memory_usage()
+        self.assertGreater(after, before)
+        self.assertEqual(a, y.grad_fn._saved_self)
+        self.assertEqual(a, y.grad_fn._saved_other)
+        y.sum().backward()
+        self.assertEqual(2 * a, a.grad)
+
 
 def index_perm_variable(shape, max_indices):
     if not isinstance(shape, tuple):

torch/autograd/saved_variable_default_hooks.py

@@ -1,7 +1,30 @@
 import torch
 
+from typing import Any
+
 def set_saved_tensors_default_hooks(pack_hook, unpack_hook):
     torch._C._autograd._register_default_hooks(pack_hook, unpack_hook)
 
 def reset_saved_tensors_default_hooks():
     torch._C._autograd._reset_default_hooks()
+
+class save_on_cpu(object):
+    r"""Context-manager under which tensors saved by the forward pass will be
+    stored on cpu, then retrieved for backward
+    """
+    def __init__(self):
+        def pack_hook(tensor):
+            return (tensor.device, tensor.cpu())
+
+        def unpack_hook(packed):
+            device, tensor = packed
+            return tensor.to(device)
+
+        self.pack_hook = pack_hook
+        self.unpack_hook = unpack_hook
+
+    def __enter__(self) -> None:
+        torch._C._autograd._register_default_hooks(self.pack_hook, self.unpack_hook)
+
+    def __exit__(self, *args: Any) -> None:
+        torch._C._autograd._reset_default_hooks()