Add check_neural_compressor_min_version for 4 bit behavior (#1500)

Signed-off-by: Xin <[email protected]>
Signed-off-by: xinhe3 <[email protected]>
Co-authored-by: xinhe3 <[email protected]>

Author: 2 people

examples/text-generation/utils.py

@@ -38,6 +38,7 @@
 )
 from optimum.habana.utils import (
     check_habana_frameworks_version,
+    check_neural_compressor_min_version,
     check_optimum_habana_min_version,
     get_habana_frameworks_version,
     set_seed,
@@ -267,9 +268,8 @@ def setup_model(args, model_dtype, model_kwargs, logger):
             original_model=org_model,
             **model_kwargs,
         )
-        # TODO: This will be removed in v1.19 Synapse release
-        # the loaded model should have the same dtype as original_model
-        model = model.to(model_kwargs["torch_dtype"])
+        if not check_neural_compressor_min_version("3.2"):
+            model = model.to(model_kwargs["torch_dtype"])
     else:
         if args.assistant_model is not None:
             assistant_model = AutoModelForCausalLM.from_pretrained(

optimum/habana/utils.py

@@ -384,6 +384,15 @@ def check_habana_frameworks_version(req_version):
     )
 
 
+def check_neural_compressor_min_version(req_version):
+    """
+    Checks if the installed version of `neural_compressor` is larger than or equal to `req_version`.
+    """
+    import neural_compressor
+
+    return version.Version(neural_compressor.__version__) >= version.Version(req_version)
+
+
 def get_device_name():
     """
     Returns the name of the current device: Gaudi or Gaudi2.