Device agnostic for DCP

Author: Chao1Han

test/distributed/checkpoint/_experimental/test_builder.py

@@ -123,7 +123,7 @@ 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_non_blocking_copy=torch.cuda.is_available(),
             use_pinned_memory=torch.cuda.is_available(),
         )
         checkpointer = make_async_checkpointer(config=config, rank_info=self.rank_info)

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,21 +86,21 @@ 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,
             ),
         ]
 
@@ -111,7 +111,7 @@ def test_different_option_combinations(self) -> None:
                     use_pinned_memory=torch.cuda.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
@@ -120,7 +120,7 @@ def test_different_option_combinations(self) -> None:
                     use_pinned_memory=torch.cuda.is_available(),
                     use_shared_memory=False,
                     use_async_staging=False,
-                    use_cuda_non_blocking_copy=torch.cuda.is_available(),
+                    use_non_blocking_copy=torch.cuda.is_available(),
                 )
             )
 
@@ -185,7 +185,7 @@ def test_multiple_staging_operations(self) -> None:
             use_async_staging=False,
             use_pinned_memory=torch.cuda.is_available(),
             use_shared_memory=False,
-            use_cuda_non_blocking_copy=torch.cuda.is_available(),
+            use_non_blocking_copy=torch.cuda.is_available(),
         )
         stager = DefaultStager(options)
 

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(

torch/distributed/checkpoint/_experimental/staging.py

@@ -82,7 +82,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 CUDA memory
             copies with stream synchronization. Improves performance by
             allowing CPU work to continue during GPU transfers. Default: True
 
@@ -93,7 +93,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):
@@ -153,15 +153,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 CUDA/XPU"
+            )
 
     def stage(
         self,
@@ -182,16 +184,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
 

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 CUDA 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 CUDA/XPU"
+            )
 
         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