fix hpu detect issue

Signed-off-by: xinhe3 <xinhe3@habana.ai>

Author: xinhe3

auto_round/autoround.py

@@ -82,7 +82,7 @@
     is_debug_mode,
     is_mx_fp,
     is_nv_fp,
-    is_optimum_habana_available,
+    is_hpex_available,
     is_standard_fp,
     is_static_wfp8afp8,
     is_wfp8afp8,
@@ -375,7 +375,7 @@ def __init__(
         self._check_configs()
         torch.set_printoptions(precision=3, sci_mode=True)
 
-        if is_optimum_habana_available():
+        if is_hpex_available():
             logger.info("optimum Habana is available, import htcore explicitly.")
             import habana_frameworks.torch.core as htcore  # pylint: disable=E0401
             import habana_frameworks.torch.hpu as hthpu  # pylint: disable=E0401]
@@ -3238,7 +3238,7 @@ def _scale_loss_and_backward(self, scaler: Any, loss: torch.Tensor) -> torch.Ten
         """
         scale_loss = loss * 1000
         scale_loss.backward()
-        if is_optimum_habana_available():
+        if is_hpex_available():
             htcore.mark_step()
         return scale_loss
 
@@ -3255,7 +3255,7 @@ def _step(self, scaler: Any, optimizer: Any, lr_schedule: Any):
         """
         optimizer.step()
         # for hpu
-        if is_optimum_habana_available():
+        if is_hpex_available():
             htcore.mark_step()
         optimizer.zero_grad()
         lr_schedule.step()
@@ -3433,7 +3433,7 @@ def _scale_loss_and_backward(self, scaler, loss):
             loss = scaler.scale(loss)
 
         loss.backward()
-        if is_optimum_habana_available():
+        if is_hpex_available():
             htcore.mark_step()
         return loss
 
@@ -3447,5 +3447,5 @@ def _step(self, scaler, optimizer, lr_schedule):
             optimizer.step()
             optimizer.zero_grad()
             lr_schedule.step()
-        if is_optimum_habana_available():
+        if is_hpex_available():
             htcore.mark_step()

auto_round/utils.py

@@ -115,7 +115,7 @@ def __init__(self, module_name):
         """Init LazyImport object.
 
         Args:
-           module_name (string): The name of module imported later
+            module_name (string): The name of module imported later
         """
         self.module_name = module_name
         self.module = None
@@ -143,13 +143,31 @@ def __call__(self, *args, **kwargs):
 auto_gptq = LazyImport("auto_gptq")
 htcore = LazyImport("habana_frameworks.torch.core")
 
+################ Check available sys.module to decide behavior #################
+def is_package_available(package_name):
+    """Check if the package exists in the environment without importing.
+
+    Args:
+        package_name (str): package name
+    """
+    from importlib.util import find_spec
+
+    package_spec = find_spec(package_name)
+    return package_spec is not None
+
+
+## check hpex
+if is_package_available("habana_frameworks"):
+    _hpex_available = True
+    import habana_frameworks.torch.hpex  # pylint: disable=E0401
+else:
+    _hpex_available = False
+
 
 @torch._dynamo.disable()
 @lru_cache(None)
-def is_optimum_habana_available():
-    from transformers.utils.import_utils import is_optimum_available
-
-    return is_optimum_available() and importlib.util.find_spec("optimum.habana") is not None
+def is_hpex_available():
+    return _hpex_available
 
 
 def get_module(module, key):
@@ -552,7 +570,7 @@ def is_valid_digit(s):
         if torch.cuda.is_available():
             device = torch.device("cuda")
             # logger.info("Using GPU device")
-        elif is_optimum_habana_available():  # pragma: no cover
+        elif is_hpex_available():  # pragma: no cover
             device = torch.device("hpu")
             # logger.info("Using HPU device")
         elif torch.xpu.is_available():  # pragma: no cover