update install

Author: songhappy

install/install_requirements.sh

@@ -53,16 +53,16 @@ echo "Using pip executable: $PIP_EXECUTABLE"
 # package versions.
 if [[ -x "$(command -v xpu-smi)" ]];
 then
-  PYTORCH_NIGHTLY_VERSION=dev20241217
+  PYTORCH_NIGHTLY_VERSION=dev20250110
 else
   PYTORCH_NIGHTLY_VERSION=dev20241218
 fi
 
 # Nightly version for torchvision
-VISION_NIGHTLY_VERSION=dev20241218
+VISION_NIGHTLY_VERSION=dev20250111
 
 # Nightly version for torchtune
-TUNE_NIGHTLY_VERSION=dev20241218
+TUNE_NIGHTLY_VERSION=dev20250105
 
 # Uninstall triton, as nightly will depend on pytorch-triton, which is one and the same
 (
@@ -90,9 +90,9 @@ fi
 if [[ -x "$(command -v xpu-smi)" ]];
 then
   REQUIREMENTS_TO_INSTALL=(
-    torch=="2.6.0.${PYTORCH_NIGHTLY_VERSION}"
+    torch=="2.7.0.${PYTORCH_NIGHTLY_VERSION}"
     torchvision=="0.22.0.${VISION_NIGHTLY_VERSION}"
-    torchtune=="0.4.0"
+   # torchtune=="0.4.0"
   )
 else
   REQUIREMENTS_TO_INSTALL=(
@@ -122,10 +122,28 @@ fi
 
 # For torchao need to install from github since nightly build doesn't have macos build.
 # TODO: Remove this and install nightly build, once it supports macos
-(
-  set -x
-  $PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@2f97b0955953fa1a46594a27f0df2bc48d93e79d
-)
+if [[ -x "$(command -v xpu-smi)" ]];
+then
+  # install torchao nightly for xpu
+  (
+    set -x
+    $PIP_EXECUTABLE install --extra-index-url "${TORCH_NIGHTLY_URL}" torchao=="0.8.0.dev20250110"
+  )
+else
+  (
+    set -x
+    $PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@2f97b0955953fa1a46594a27f0df2bc48d93e79d
+  )
+fi
+
+# install torchtune from source for xpu
+if [[ -x "$(command -v xpu-smi)" ]];
+then
+  (
+    set -x
+    $PIP_EXECUTABLE install git+https://github.com/pytorch/torchtune
+  )
+fi
 
 if [[ -x "$(command -v nvidia-smi)" ]]; then
   (
@@ -134,7 +152,6 @@ if [[ -x "$(command -v nvidia-smi)" ]]; then
   )
 fi
 
-
 (
   set -x
   $PIP_EXECUTABLE install evaluate=="0.4.3" lm-eval=="0.4.2" psutil=="6.0.0"

torchchat/generate.py

@@ -1168,7 +1168,6 @@ def callback(x, *, done_generating=False):
             else:
                 torch.profiler._utils._init_for_cuda_graphs()
                 prof = torch.profiler.profile()
-            print("prof is: ", prof)
             t0 = time.perf_counter()
             num_tokens_generated = 0
             with prof:

torchchat/utils/quantize.py

@@ -111,7 +111,7 @@ def quantize_model(
             raise RuntimeError(f"unknown quantizer {quantizer} specified")
         else:
             # Use tensor subclass API for int4 weight only.
-            if device == "cuda" and quantizer == "linear:int4":
+            if (device == "cuda" or device == "xpu") and quantizer == "linear:int4":
                 quantize_(model, int4_weight_only(q_kwargs["groupsize"]))
                 if not support_tensor_subclass:
                     unwrap_tensor_subclass(model)