enable fsdp tests

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

Author: jiqing-feng

tests/fsdp_state_dict_save.py

@@ -20,6 +20,14 @@
 import bitsandbytes as bnb
 
 
+def _get_device_and_backend():
+    """Auto-detect accelerator device and distributed backend."""
+    device_type = str(torch.accelerator.current_accelerator())
+    backend_map = {"cuda": "nccl", "xpu": "ccl"}
+    backend = backend_map.get(device_type, "gloo")
+    return device_type, backend
+
+
 class SimpleQLoRAModel(nn.Module):
     """Minimal model with a frozen 4-bit base layer and a trainable adapter."""
 
@@ -33,15 +41,16 @@ def forward(self, x):
 
 
 def main():
-    dist.init_process_group(backend="nccl")
+    device_type, backend = _get_device_and_backend()
+    dist.init_process_group(backend=backend)
     rank = dist.get_rank()
-    torch.cuda.set_device(rank)
+    torch.accelerator.set_device_index(rank)
 
     errors = []
 
     for quant_type in ("nf4", "fp4"):
         model = SimpleQLoRAModel(quant_type=quant_type)
-        model = model.to("cuda")
+        model = model.to(device_type)
 
         # Freeze quantized base weights (as in real QLoRA)
         for p in model.base.parameters():

tests/test_linear4bit.py

@@ -569,11 +569,7 @@ def test_params4bit_quant_state_attr_access(device, quant_type, compress_statist
     assert w.bnb_quantized is True
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="FSDP requires CUDA")
-@pytest.mark.skipif(
-    not getattr(torch.distributed, "is_nccl_available", lambda: False)(),
-    reason="FSDP test requires NCCL backend",
-)
+@pytest.mark.skipif(not torch.accelerator.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).