diff --git a/neural_compressor/torch/algorithms/weight_only/modules.py b/neural_compressor/torch/algorithms/weight_only/modules.py index c33310b61ea..9d0367c1853 100644 --- a/neural_compressor/torch/algorithms/weight_only/modules.py +++ b/neural_compressor/torch/algorithms/weight_only/modules.py @@ -26,10 +26,9 @@ from torch.nn import functional as F from neural_compressor.torch.utils import ( - Version, accelerator, can_pack_with_numba, - get_hpex_version, + is_hpex_support_g_idx, logger, ) @@ -731,7 +730,7 @@ def __init__( ) else: self.g_idx = None - self.support_g_idx = True if get_hpex_version() >= Version("1.23.0") else False + self.support_g_idx = is_hpex_support_g_idx() self.half_indim = self.in_features // 2 diff --git a/neural_compressor/torch/utils/environ.py b/neural_compressor/torch/utils/environ.py index 4cf2b32bf25..c8b8f596e61 100644 --- a/neural_compressor/torch/utils/environ.py +++ b/neural_compressor/torch/utils/environ.py @@ -17,6 +17,7 @@ import importlib import os import sys +from functools import lru_cache import torch from packaging.version import Version @@ -79,19 +80,20 @@ def is_hpu_available(): return get_accelerator().name() == "hpu" -def get_hpex_version(): - """Return ipex version if ipex exists.""" +@lru_cache(None) +def is_hpex_support_g_idx(): + """Check if HPEX supports group_index in the schema of hpu::convert_from_int4.""" if is_hpex_available(): try: import habana_frameworks.torch + import torch - hpex_version = habana_frameworks.torch.__version__ - except ValueError as e: # pragma: no cover - assert False, "Got an unknown version of habana_frameworks.torch: {}".format(e) - version = Version(hpex_version) - return version + schema = torch._C._get_schema("hpu::convert_from_int4", "") + return "group_index" in str(schema) + except: # pragma: no cover + return False else: - return None + return False ## check optimum