Sync states for xpu fsdp (#2005)

* sync states for xpu fsdp * Update src/accelerate/utils/dataclasses.py Co-authored-by: Zach Mueller <[email protected]> --------- Co-authored-by: Zach Mueller <[email protected]>
huggingface · Oct 2, 2023 · 420743a · 420743a
1 parent 206ab49
commit 420743a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
@@ -32,6 +32,7 @@
 
 from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH, FSDP_STATE_DICT_TYPE
 from .environment import str_to_bool
+from .imports import is_xpu_available
 from .versions import compare_versions
 
 
@@ -920,7 +921,8 @@ def __post_init__(self):
         self.activation_checkpointing = str_to_bool(os.environ.get(prefix + "ACTIVATION_CHECKPOINTING", "False")) == 1
 
         if self.sync_module_states:
-            self.param_init_fn = lambda x: x.to_empty(device=torch.cuda.current_device(), recurse=False)
+            device = torch.cuda.current_device() if not is_xpu_available() else torch.xpu.current_device()
+            self.param_init_fn = lambda x: x.to_empty(device=device, recurse=False)
 
     @staticmethod
     def get_module_class_from_name(module, name):