add INC_PT_ONLY and INC_TF_ONLY (#2202)

* add INC_PT_ONLY and INC_TF_ONLY
* compatible with previous install method

---------

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

Author: xin3he

neural_compressor/__init__.py

@@ -17,16 +17,18 @@
 """Intel® Neural Compressor: An open-source Python library supporting popular model compression techniques."""
 from .version import __version__
 
-# we need to set a global 'NA' backend, or Model can't be used
-from .config import (
-    DistillationConfig,
-    PostTrainingQuantConfig,
-    WeightPruningConfig,
-    QuantizationAwareTrainingConfig,
-    MixedPrecisionConfig,
-)
-from .contrib import *
-from .model import *
-from .metric import *
-from .utils import options
-from .utils.utility import set_random_seed, set_tensorboard, set_workspace, set_resume_from
+import os
+
+if not (os.environ.get("INC_PT_ONLY", False) or os.environ.get("INC_TF_ONLY", False)):
+    from .config import (
+        DistillationConfig,
+        PostTrainingQuantConfig,
+        WeightPruningConfig,
+        QuantizationAwareTrainingConfig,
+        MixedPrecisionConfig,
+    )
+    from .contrib import *
+    from .model import *
+    from .metric import *
+    from .utils import options
+    from .utils.utility import set_random_seed, set_tensorboard, set_workspace, set_resume_from

setup.py

@@ -95,22 +95,25 @@ def get_build_version():
 
 
 if __name__ == "__main__":
-    cfg_key = "neural_compressor"
-
-    # Temporary implementation of fp8 tensor saving and loading
-    # Will remove after Habana torch applies below patch:
-    # https://github.com/pytorch/pytorch/pull/114662
-    ext_modules = []
-    cmdclass = {}
-
+    # for setuptools>=80.0.0, `INC_PT_ONLY=1 pip install -e .`
+    if os.environ.get("INC_PT_ONLY", False) and os.environ.get("INC_TF_ONLY", False):
+        raise ValueError("Both INC_PT_ONLY and INC_TF_ONLY are set. Please set only one.")
+    if os.environ.get("INC_PT_ONLY", False):
+        cfg_key = "neural_compressor_pt"
+    elif os.environ.get("INC_TF_ONLY", False):
+        cfg_key = "neural_compressor_tf"
+    else:
+        cfg_key = "neural_compressor"
+    # for setuptools < 80.0.0, `python setup.py develop pt`
     if "pt" in sys.argv:
         sys.argv.remove("pt")
         cfg_key = "neural_compressor_pt"
-
     if "tf" in sys.argv:
         sys.argv.remove("tf")
         cfg_key = "neural_compressor_tf"
 
+    ext_modules = []
+    cmdclass = {}
     project_name = PKG_INSTALL_CFG[cfg_key].get("project_name")
     include_packages = PKG_INSTALL_CFG[cfg_key].get("include_packages") or {}
     package_data = PKG_INSTALL_CFG[cfg_key].get("package_data") or {}

test/3x/torch/algorithms/fp8_quant/conftest.py

@@ -1,12 +1,15 @@
 # Called once at the beginning of the test session
 def pytest_sessionstart():
+    import os
     import habana_frameworks.torch.core as htcore
     import torch
 
     htcore.hpu_set_env()
 
     # Use reproducible results
     torch.use_deterministic_algorithms(True)
+    # Ensure that only 3x PyTorch part of INC is imported
+    os.environ.setdefault("INC_PT_ONLY", "1")
 
     # Fix the seed - just in case
     torch.manual_seed(0)

test/3x/torch/quantization/fp8_quant/conftest.py

@@ -2,6 +2,8 @@
 # Ensure that the HPU is in lazy mode and weight sharing is disabled
 os.environ.setdefault("PT_HPU_LAZY_MODE", "1")
 os.environ.setdefault("PT_HPU_WEIGHT_SHARING", "0")
+# Ensure that only 3x PyTorch part of INC is imported
+os.environ.setdefault("INC_PT_ONLY", "1")
 
 
 def pytest_sessionstart():