Device agnostic for DCP (#158337)

Chao1Han · guangyey · pytorchmergebot · commit 02ca965560e7 · 2025-07-25T05:24:09.000Z
Enable device-agnostic implementation of DCP-related functionality, allowing the new DCP features to be supported on XPU as well. use_cuda_non_blocking_copy to use_non_blocking_copy because non-blocking copy is supported by most GPUs and is not exclusive to CUDA devices. Test plan: test cases have not yet been updated to be fully device agnostic; this will be addressed in future work. Pull Request resolved: #158337 Approved by: https://github.com/guangyey, https://github.com/EikanWang, https://github.com/Saiteja64 Co-authored-by: Yu, Guangye <106960996+guangyey@users.noreply.github.com>
diff --git a/test/distributed/checkpoint/_experimental/test_builder.py b/test/distributed/checkpoint/_experimental/test_builder.py
@@ -123,8 +123,8 @@ def test_make_async_checkpointer(self) -> None:
         # Create async checkpointer using factory function with default parameters
         config: CheckpointerConfig = CheckpointerConfig()
         config.staging_config = CheckpointStagerConfig(
-            use_cuda_non_blocking_copy=torch.cuda.is_available(),
-            use_pinned_memory=torch.cuda.is_available(),
+            use_non_blocking_copy=torch.accelerator.is_available(),
+            use_pinned_memory=torch.accelerator.is_available(),
         )
         checkpointer = make_async_checkpointer(config=config, rank_info=self.rank_info)
 
diff --git a/test/distributed/checkpoint/_experimental/test_staging.py b/test/distributed/checkpoint/_experimental/test_staging.py
@@ -74,7 +74,7 @@ def test_cuda_non_blocking_without_cuda(self) -> None:
         if torch.cuda.is_available():
             self.skipTest("CUDA is available, cannot test CUDA unavailable scenario")
 
-        options = CheckpointStagerConfig(use_cuda_non_blocking_copy=True)
+        options = CheckpointStagerConfig(use_non_blocking_copy=True)
         with self.assertRaises(AssertionError):
             DefaultStager(options)
 
@@ -86,41 +86,41 @@ def test_different_option_combinations(self) -> None:
                 use_pinned_memory=False,
                 use_shared_memory=False,
                 use_async_staging=False,
-                use_cuda_non_blocking_copy=False,
+                use_non_blocking_copy=False,
             ),
             # Only pinned memory
             CheckpointStagerConfig(
                 use_pinned_memory=True,
                 use_shared_memory=False,
                 use_async_staging=False,
-                use_cuda_non_blocking_copy=False,
+                use_non_blocking_copy=False,
             ),
             # Only shared memory
             CheckpointStagerConfig(
                 use_pinned_memory=False,
                 use_shared_memory=True,
                 use_async_staging=False,
-                use_cuda_non_blocking_copy=False,
+                use_non_blocking_copy=False,
             ),
         ]
 
         if torch.cuda.is_available():
             # Only async staging
             test_cases.append(
                 CheckpointStagerConfig(
-                    use_pinned_memory=torch.cuda.is_available(),
+                    use_pinned_memory=torch.accelerator.is_available(),
                     use_shared_memory=False,
                     use_async_staging=True,
-                    use_cuda_non_blocking_copy=False,
+                    use_non_blocking_copy=False,
                 )
             )
             # Only CUDA non-blocking copy
             test_cases.append(
                 CheckpointStagerConfig(
-                    use_pinned_memory=torch.cuda.is_available(),
+                    use_pinned_memory=torch.accelerator.is_available(),
                     use_shared_memory=False,
                     use_async_staging=False,
-                    use_cuda_non_blocking_copy=torch.cuda.is_available(),
+                    use_non_blocking_copy=torch.accelerator.is_available(),
                 )
             )
 
@@ -129,7 +129,7 @@ def test_different_option_combinations(self) -> None:
                 stager = DefaultStager(options)
 
                 # Test staging works with these options
-                if options.use_async_staging and torch.cuda.is_available():
+                if options.use_async_staging and torch.accelerator.is_available():
                     result = stager.stage(self.state_dict)
                     self.assertIsInstance(result, Future)
                     staged_dict = result.result()
@@ -183,9 +183,9 @@ def test_multiple_staging_operations(self) -> None:
         """Test multiple staging operations with the same stager."""
         options = CheckpointStagerConfig(
             use_async_staging=False,
-            use_pinned_memory=torch.cuda.is_available(),
+            use_pinned_memory=torch.accelerator.is_available(),
             use_shared_memory=False,
-            use_cuda_non_blocking_copy=torch.cuda.is_available(),
+            use_non_blocking_copy=torch.accelerator.is_available(),
         )
         stager = DefaultStager(options)
 
diff --git a/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py b/test/distributed/checkpoint/e2e/test_e2e_save_and_load.py
@@ -279,7 +279,7 @@ def _run_e2e_test(
                     use_async_staging=zoc,
                     use_shared_memory=use_shared_memory,
                     use_pinned_memory=zoc,
-                    use_cuda_non_blocking_copy=zoc,
+                    use_non_blocking_copy=zoc,
                 )
                 stager = DefaultStager(staging_options)
             async_save_response_or_future = saver.async_save(
diff --git a/torch/distributed/checkpoint/_experimental/staging.py b/torch/distributed/checkpoint/_experimental/staging.py
@@ -77,7 +77,7 @@ class CheckpointStagerConfig:
         use_async_staging (bool): Enable asynchronous staging using a
             background thread pool. Allows overlapping computation with
             staging operations. Requires CUDA. Default: True
-        use_cuda_non_blocking_copy (bool): Use non-blocking CUDA memory
+        use_non_blocking_copy (bool): Use non-blocking device memory
             copies with stream synchronization. Improves performance by
             allowing CPU work to continue during GPU transfers. Default: True
 
@@ -88,7 +88,7 @@ class CheckpointStagerConfig:
     use_pinned_memory: bool = True
     use_shared_memory: bool = True
     use_async_staging: bool = True
-    use_cuda_non_blocking_copy: bool = True
+    use_non_blocking_copy: bool = True
 
 
 class DefaultStager(CheckpointStager):
@@ -148,15 +148,17 @@ def __init__(
 
         if self._config.use_async_staging:
             self._staging_executor = ThreadPoolExecutor(max_workers=1)
-            if torch.cuda.is_available():
+            if torch.accelerator.is_available():
                 # Note: stream needs to be initialized on the main thread after default cuda
                 # stream is setup/used to avoid the risk of accidentally reusing the main
                 # compute stream or in other cases kernels actually launching from the
                 # main thread.
-                self._staging_stream = torch.cuda.Stream()
+                self._staging_stream = torch.Stream()
 
-        if self._config.use_cuda_non_blocking_copy:
-            assert torch.cuda.is_available(), "Non-blocking copy requires CUDA"
+        if self._config.use_non_blocking_copy:
+            assert torch.accelerator.is_available(), (
+                "Non-blocking copy requires that the current accelerator is available."
+            )
 
     def stage(
         self,
@@ -177,16 +179,16 @@ def stage(
 
     def _stage(self, state_dict: STATE_DICT, **kwargs: Any) -> STATE_DICT:
         state_dict = self._state_dict_stager.stage(
-            state_dict, non_blocking=self._config.use_cuda_non_blocking_copy, **kwargs
+            state_dict, non_blocking=self._config.use_non_blocking_copy, **kwargs
         )
 
-        if self._config.use_cuda_non_blocking_copy:
+        if self._config.use_non_blocking_copy:
             assert self._staging_stream or not self._config.use_async_staging, (
-                "Non-blocking cuda copy in a background thread for async staging needs staging_stream to be initialized."
+                "Non-blocking copy in a background thread for async staging needs staging_stream to be initialized."
             )
 
             # waits for the enqued copy operations to finish.
-            self._staging_stream.synchronize() if self._staging_stream else torch.cuda.synchronize()
+            self._staging_stream.synchronize() if self._staging_stream else torch.accelerator.synchronize()
 
         return state_dict
 
diff --git a/torch/distributed/checkpoint/staging.py b/torch/distributed/checkpoint/staging.py
@@ -110,7 +110,7 @@ class StagingOptions:
         use_async_staging (bool): Enable asynchronous staging using a
             background thread pool. Allows overlapping computation with
             staging operations. Requires CUDA. Default: True
-        use_cuda_non_blocking_copy (bool): Use non-blocking CUDA memory
+        use_non_blocking_copy (bool): Use non-blocking device memory
             copies with stream synchronization. Improves performance by
             allowing CPU work to continue during GPU transfers. Default: True
 
@@ -121,7 +121,7 @@ class StagingOptions:
     use_pinned_memory: bool = True
     use_shared_memory: bool = True
     use_async_staging: bool = True
-    use_cuda_non_blocking_copy: bool = True
+    use_non_blocking_copy: bool = True
 
 
 class DefaultStager(AsyncStager):
@@ -177,15 +177,17 @@ def __init__(
         self._staging_stream = None
         if self._config.use_async_staging:
             self._staging_executor = ThreadPoolExecutor(max_workers=1)
-            if torch.cuda.is_available():
+            if torch.accelerator.is_available():
                 # Note: stream needs to be initialized on the main thread after default cuda
                 # stream is setup/used to avoid the risk of accidentally reusing the main
                 # compute stream or in other cases kernels actually launching from the
                 # main thread.
-                self._staging_stream = torch.cuda.Stream()
+                self._staging_stream = torch.Stream()
 
-        if self._config.use_cuda_non_blocking_copy:
-            assert torch.cuda.is_available(), "Non-blocking copy requires CUDA"
+        if self._config.use_non_blocking_copy:
+            assert torch.accelerator.is_available(), (
+                "Non-blocking copy requires that the current accelerator is available."
+            )
 
         self._staging_future: Optional[Future[STATE_DICT_TYPE]] = None
 
@@ -216,20 +218,20 @@ def stage(
             return self._stage(state_dict, **kwargs)
 
     def _stage(self, state_dict: STATE_DICT_TYPE, **kwargs: Any) -> STATE_DICT_TYPE:
-        if self._config.use_cuda_non_blocking_copy:
+        if self._config.use_non_blocking_copy:
             assert self._staging_stream or not self._config.use_async_staging, (
-                "Non-blocking cuda copy in a background thread for async staging needs staging_stream to be initialized."
+                "Non-blocking copy in a background thread for async staging needs staging_stream to be initialized."
             )
             with (
                 self._staging_stream
                 if self._staging_stream is not None
                 else nullcontext()
             ):
                 state_dict = self._state_dict_stager.stage(
-                    state_dict, non_blocking=self._config.use_cuda_non_blocking_copy
+                    state_dict, non_blocking=self._config.use_non_blocking_copy
                 )
             # waits for the enqued copy operations to finish.
-            self._staging_stream.synchronize() if self._staging_stream else torch.cuda.synchronize()
+            self._staging_stream.synchronize() if self._staging_stream else torch.accelerator.synchronize()
         else:
             state_dict = self._state_dict_stager.stage(state_dict, non_blocking=False)
         return state_dict
