XPU support for float8 dtensor test

Author: arlesniak

test/float8/test_dtensor.py

@@ -45,14 +45,20 @@
 
 torch.set_float32_matmul_precision("high")
 
+_GPU_DEVICE = (
+    str(torch.accelerator.current_accelerator())
+    if torch.accelerator.is_available()
+    else "no_gpu"
+)
+
 
 def setup_distributed():
     world_size = int(os.environ.get("WORLD_SIZE", -1))
-    device_mesh = init_device_mesh("cuda", (world_size,))
+    device_mesh = init_device_mesh(_GPU_DEVICE, (world_size,))
     # seed must be the same in all processes
     torch.manual_seed(1)
     local_rank = torch.distributed.get_rank()
-    torch.cuda.set_device(local_rank)
+    torch.get_device_module(_GPU_DEVICE).set_device(local_rank)
     return device_mesh
 
 
@@ -213,7 +219,7 @@ def _test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=32):
 
 def _test_distribute_fsdp_tensor_subclass(tp_mesh: DeviceMesh):
     torch.manual_seed(42)
-    model = Transformer(ModelArgs(dropout_p=0.0, weight_tying=False)).cuda()
+    model = Transformer(ModelArgs(dropout_p=0.0, weight_tying=False)).to(_GPU_DEVICE)
     convert_to_float8_training(
         model,
         config=Float8LinearConfig(

test/float8/test_dtensor.sh

@@ -8,8 +8,15 @@
 # terminate script on first error
 set -e
 
-if python -c 'import torch;print(torch.cuda.is_available())' | grep -q "False"; then
-    echo "Skipping test_dtensor.sh because no CUDA devices are available."
+if ! python - <<'PY'
+import sys
+import torch
+has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
+has_cuda = torch.cuda.is_available()
+sys.exit(0 if (has_xpu or has_cuda) else 1)
+PY
+then
+    echo "Skipping test_dtensor.sh because no XPU/CUDA devices are available."
     exit
 fi