Skip to content

auto detect available device when exporting #1743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Apr 25, 2024
Merged

auto detect available device when exporting #1743

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5f61b06
use cuda if exists for XPU format export
xin3he Apr 16, 2024
ee2578f
support cpu export
xin3he Apr 18, 2024
757e3c8
auto select availiable device when export
xin3he Apr 19, 2024
929ecf1
add log
xin3he Apr 19, 2024
2306c69
fix bug
xin3he Apr 19, 2024
2f98936
Merge branch 'master' into xpu_export
xin3he Apr 19, 2024
e567ab2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 19, 2024
75ca99a
Merge branch 'master' into xpu_export
xin3he Apr 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion neural_compressor/model/torch_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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":
Expand Down