add xpu device

Author: jenniew

install/install_requirements.sh

@@ -35,14 +35,9 @@ fi
 # newer version of torch nightly installed later in this script.
 #
 
-#(
-#  set -x
-#  $PIP_EXECUTABLE install -r install/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cu121
-#)
-
 (
   set -x
-  $PIP_EXECUTABLE install -r install/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/xpu
+  $PIP_EXECUTABLE install -r install/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cu121
 )
 
 # Since torchchat often uses main-branch features of pytorch, only the nightly
@@ -52,7 +47,12 @@ fi
 # NOTE: If a newly-fetched version of the executorch repo changes the value of
 # PYTORCH_NIGHTLY_VERSION, you should re-run this script to install the necessary
 # package versions.
-PYTORCH_NIGHTLY_VERSION=dev20241001
+if [[ -x "$(command -v xpu-smi)" ]];
+then
+  PYTORCH_NIGHTLY_VERSION=dev20241001
+else
+  PYTORCH_NIGHTLY_VERSION=dev20241002
+fi
 
 # Nightly version for torchvision
 VISION_NIGHTLY_VERSION=dev20241002
@@ -69,22 +69,34 @@ TUNE_NIGHTLY_VERSION=dev20241010
 # The pip repository that hosts nightly torch packages. cpu by default.
 # If cuda is available, based on presence of nvidia-smi, install the pytorch nightly
 # with cuda for faster execution on cuda GPUs.
-#if [[ -x "$(command -v nvidia-smi)" ]];
-#then
-#  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cu121"
-#elif [[ -x "$(command -v rocminfo)" ]];
-#then
-#  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/rocm6.2"
-#else
-#  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cpu"
-#fi
-TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/xpu"
+if [[ -x "$(command -v nvidia-smi)" ]];
+then
+  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cu121"
+elif [[ -x "$(command -v rocminfo)" ]];
+then
+  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/rocm6.2"
+elif [[ -x "$(command -v xpu-smi)" ]];
+then
+  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/xpu"
+else
+  TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cpu"
+fi
+
 # pip packages needed by exir.
-REQUIREMENTS_TO_INSTALL=(
-  torch=="2.6.0.${PYTORCH_NIGHTLY_VERSION}"
-  torchvision=="0.20.0.${VISION_NIGHTLY_VERSION}"
-  #torchtune=="0.3.0.${TUNE_NIGHTLY_VERSION}"
-)
+if [[ -x "$(command -v xpu-smi)" ]];
+then
+  REQUIREMENTS_TO_INSTALL=(
+    torch=="2.6.0.${PYTORCH_NIGHTLY_VERSION}"
+    torchvision=="0.20.0.${VISION_NIGHTLY_VERSION}"
+    torchtune=="0.3.1"
+  )
+else
+  REQUIREMENTS_TO_INSTALL=(
+    torch=="2.6.0.${PYTORCH_NIGHTLY_VERSION}"
+    torchvision=="0.20.0.${VISION_NIGHTLY_VERSION}"
+    torchtune=="0.4.0.${TUNE_NIGHTLY_VERSION}"
+  )
+fi
 
 # Install the requirements. --extra-index-url tells pip to look for package
 # versions on the provided URL if they aren't available on the default URL.
@@ -94,12 +106,6 @@ REQUIREMENTS_TO_INSTALL=(
     "${REQUIREMENTS_TO_INSTALL[@]}"
 )
 
-(
-  set -x
-  $PIP_EXECUTABLE install --extra-index-url "${TORCH_NIGHTLY_URL}" \
-    torchtune=="0.3.1"
-)
-
 (
   set -x
   $PIP_EXECUTABLE install torchao=="0.5.0"

torchchat/cli/builder.py

@@ -68,7 +68,13 @@ class BuilderArgs:
 
     def __post_init__(self):
         if self.device is None:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            # self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            if torch.cuda.is_available():
+                self.device = "cuda"
+            elif hasattr(torch, 'xpu') and torch.xpu.is_available():
+                self.device = "xpu"
+            else:
+                self.device = "cpu"
 
         if not (
             (self.checkpoint_path and self.checkpoint_path.is_file())

torchchat/cli/cli.py

@@ -168,7 +168,7 @@ def _add_model_config_args(parser, verb: str) -> None:
         type=str,
         default=default_device,
         choices=["fast", "cpu", "cuda", "mps", "xpu"],
-        help="Hardware device to use. Options: cpu, cuda, mps",
+        help="Hardware device to use. Options: cpu, cuda, mps, xpu",
     )
 
 

torchchat/utils/build_utils.py

@@ -203,6 +203,8 @@ def find_multiple(n: int, k: int) -> int:
 def device_sync(device="cpu"):
     if "cuda" in device:
         torch.cuda.synchronize(device)
+    elif "xpu" in device:
+        torch.xpu.synchronize(device)
     elif ("cpu" in device) or ("mps" in device):
         pass
     else:
@@ -251,7 +253,8 @@ def get_device_str(device) -> str:
         device = (
             "cuda"
             if torch.cuda.is_available()
-            else "mps" if is_mps_available() else "cpu"
+            else "mps" if is_mps_available()
+            else "xpu" if torch.xpu.is_available()  else "cpu"
         )
         return device
     else:
@@ -263,7 +266,8 @@ def get_device(device) -> str:
         device = (
             "cuda"
             if torch.cuda.is_available()
-            else "mps" if is_mps_available() else "cpu"
+            else "mps" if is_mps_available()
+            else "xpu" if torch.xpu.is_available()  else "cpu"
         )
     return torch.device(device)
 

torchchat/utils/device_info.py

@@ -14,7 +14,7 @@ def get_device_info(device: str) -> str:
     """Returns a human-readable description of the hardware based on a torch.device.type
 
     Args:
-        device: A torch.device.type string: one of {"cpu", "cuda"}.
+        device: A torch.device.type string: one of {"cpu", "cuda", "xpu"}.
     Returns:
         str: A human-readable description of the hardware or an empty string if the device type is unhandled.
 
@@ -37,4 +37,20 @@ def get_device_info(device: str) -> str:
             )
     if device == "cuda":
         return torch.cuda.get_device_name(0)
+    if device == "xpu":
+        # return (
+        #     check_output(
+        #         ["sycl-ls | grep gpu"], shell=True
+        #     )
+        #     .decode("utf-8")
+        #     .split("\n")[0]
+        #     )
+        return (
+            check_output(
+                ["xpu-smi discovery |grep 'Device Name:'"], shell=True
+            )
+            .decode("utf-8")
+            .split("\n")[0]
+            .split("Device Name:")[1]
+            )
     return ""