From 16d6a06c212fbc8908d443b16a354fdfd6337629 Mon Sep 17 00:00:00 2001 From: Sihan Chen <39623753+Spycsh@users.noreply.github.com> Date: Wed, 22 Nov 2023 14:37:03 +0800 Subject: [PATCH] add quantization support on MacOS (#1407) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- neural_compressor/benchmark.py | 5 ++--- neural_compressor/utils/utility.py | 5 ++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/neural_compressor/benchmark.py b/neural_compressor/benchmark.py index f14d444737b..22b8cbc152f 100644 --- a/neural_compressor/benchmark.py +++ b/neural_compressor/benchmark.py @@ -510,11 +510,10 @@ def fit(model, conf, b_dataloader=None, b_func=None): if b_dataloader is not None: check_dataloader(b_dataloader) - - assert sys.platform in ["linux", "win32"], "only support platform windows and linux..." + assert sys.platform in ["linux", "win32", "darwin"], "platform not supported..." # disable multi-instance for running benchmark on GPU device set_all_env_var(conf) - if conf.device == "gpu": + if conf.device == "gpu" or sys.platform == "darwin": set_env_var("NC_ENV_CONF", True, overwrite_existing=True) if conf.diagnosis and os.environ.get("NC_ENV_CONF", None) in [None, "False"]: diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index b46446cd5dc..1777d52f6f5 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -249,7 +249,10 @@ def __init__(self): b"\xB8\x07\x00\x00\x00" b"\x0f\xa2" b"\xC3", # mov eax, 7 # cpuid # ret ) self._bf16 = bool(eax & (1 << 5)) - self._sockets = self.get_number_of_sockets() + if "arch" in info and "ARM" in info["arch"]: # pragma: no cover + self._sockets = 1 + else: + self._sockets = self.get_number_of_sockets() self._cores = psutil.cpu_count(logical=False) self._cores_per_socket = int(self._cores / self._sockets)