auto detect available device when exporting (#1743)

xin3he · pre-commit-ci[bot] · web-flow · commit 7be355dee66c · 2024-04-25T15:59:07.000+08:00
Description
auto detect available packing device when exporting, priority: ["xpu", "cuda", "cpu"]

Expected Behavior &amp; Potential Risk
When XPU format is required and XPU device is not available, INC will select packing device from ["xpu", "cuda", "cpu"] automatically.

How has this PR been tested?
test on cuda device when using xpu argument.
---------

Signed-off-by: He, Xin3 &lt;xin3.he@intel.com&gt;
Signed-off-by: y &lt;xin3.he@intel.com&gt;
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/neural_compressor/model/torch_model.py b/neural_compressor/model/torch_model.py
@@ -498,6 +498,18 @@ def export_compressed_model(
             gptq_config = self.gptq_config if hasattr(self, "gptq_config") else {}
 
         autoround_config = self.autoround_config if hasattr(self, "autoround_config") else {}
+        # check available device, priority: ["xpu", "cuda", "cpu"]
+        availiable_device = []
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            availiable_device.append("xpu")
+        if torch.cuda.is_available():
+            availiable_device.append("cuda")
+        availiable_device.append("cpu")
+        orig_device = device
+        if device not in availiable_device and "cuda" not in device:  # cuda in cuda:0
+            logger.info(f"{device} is not detected in current environment, please check.")
+            device = availiable_device[0]
+            logger.info(f"The compression device has been changed to {device}.")
         if gptq_config:
             for k, v in weight_config.items():
                 logger.debug(f"Compressing {k} on device {device}")
@@ -558,7 +570,7 @@ def export_compressed_model(
                 new_module.pack(int_weight, gptq_scale, gptq_zp, m.bias, gptq_perm)
                 set_module(self.model, k, new_module)
         elif autoround_config:
-            if device == "xpu":
+            if orig_device == "xpu":
                 for k, v in weight_config.items():
                     logger.debug(f"Compressing {k} on device {device}")
                     if v["dtype"] == "fp32":
