fix torch version accelerator

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

Author: jiqing-feng

tests/fsdp_state_dict_save.py

@@ -20,9 +20,29 @@
 import bitsandbytes as bnb
 
 
+def _current_accelerator_type():
+    if hasattr(torch, "accelerator") and torch.accelerator.is_available():
+        return str(torch.accelerator.current_accelerator())
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        return "xpu"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def _set_device_index(index: int, device_type: str):
+    if hasattr(torch, "accelerator"):
+        torch.accelerator.set_device_index(index)
+        return
+    if device_type == "cuda":
+        torch.cuda.set_device(index)
+    elif device_type == "xpu" and hasattr(torch, "xpu") and hasattr(torch.xpu, "set_device"):
+        torch.xpu.set_device(index)
+
+
 def _get_device_and_backend():
     """Auto-detect accelerator device and distributed backend."""
-    device_type = str(torch.accelerator.current_accelerator())
+    device_type = _current_accelerator_type()
     backend_map = {"cuda": "nccl", "xpu": "xccl"}
     backend = backend_map.get(device_type, "gloo")
     return device_type, backend
@@ -44,7 +64,7 @@ def main():
     device_type, backend = _get_device_and_backend()
     dist.init_process_group(backend=backend)
     rank = dist.get_rank()
-    torch.accelerator.set_device_index(rank)
+    _set_device_index(rank, device_type)
 
     errors = []
 

tests/test_linear4bit.py

@@ -569,7 +569,14 @@ def test_params4bit_quant_state_attr_access(device, quant_type, compress_statist
     assert w.bnb_quantized is True
 
 
-@pytest.mark.skipif(not torch.accelerator.is_available(), reason="FSDP requires an accelerator device")
+@pytest.mark.skipif(
+    not (
+        (hasattr(torch, "accelerator") and torch.accelerator.is_available())
+        or torch.cuda.is_available()
+        or (hasattr(torch, "xpu") and torch.xpu.is_available())
+    ),
+    reason="FSDP requires an accelerator device",
+)
 def test_fsdp_state_dict_save_4bit():
     """Integration test: FSDP get_model_state_dict with cpu_offload on a 4-bit model (#1405).