fix load_state_dict for npu (#3211)

* fix load_state_dict for npu * update
huggingface · Dec 11, 2024 · b92fb47 · b92fb47
1 parent 3e62fbb
commit b92fb47
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
@@ -1670,9 +1670,11 @@ def load_state_dict(checkpoint_file, device_map=None):
             if len(set(device_map.values())) == 1:
                 device = list(device_map.values())[0]
                 target_device = device
-                if is_xpu_available():
-                    if isinstance(device, int):
+                if isinstance(device, int):
+                    if is_xpu_available():
                         target_device = f"xpu:{device}"
+                    elif is_npu_available():
+                        target_device = f"npu:{device}"
 
                 return safe_load_file(checkpoint_file, device=target_device)
 
@@ -1704,9 +1706,11 @@ def load_state_dict(checkpoint_file, device_map=None):
                 progress_bar = None
             for device in devices:
                 target_device = device
-                if is_xpu_available():
-                    if isinstance(device, int):
+                if isinstance(device, int):
+                    if is_xpu_available():
                         target_device = f"xpu:{device}"
+                    elif is_npu_available():
+                        target_device = f"npu:{device}"
 
                 with safe_open(checkpoint_file, framework="pt", device=target_device) as f:
                     for key in device_weights[device]: