add INC_PT_ONLY and INC_TF_ONLY

docs/source/installation_guide.md

@@ -82,6 +82,24 @@ pip install neural-compressor-tf
   [optional] pip install -r requirements_tf.txt # for TensorFlow framework extension API
   ```
 
+- [PyTorch only]
+    ```Shell
+    git clone https://github.com/intel/neural-compressor.git
+    cd neural-compressor
+    pip install -r requirements_pt.txt # for PyTorch framework extension API
+    export INC_PT_ONLY=1
+    pip install -e .
+    ```
+
+- [TensorFlow only]
+    ```Shell
+    git clone https://github.com/intel/neural-compressor.git
+    cd neural-compressor
+    pip install -r requirements_tf.txt # for TensorFlow framework extension API
+    export INC_TF_ONLY=1
+    pip install -e .
+    ```
+
 ### Install from AI Kit
 
 The Intel® Neural Compressor library is released as part of the [Intel® oneAPI AI Analytics Toolkit](https://software.intel.com/content/www/us/en/develop/tools/oneapi/ai-analytics-toolkit.html) (AI Kit). The AI Kit provides a consolidated package of Intel's latest deep learning and machine optimizations all in one place for ease of development. Along with Neural Compressor, the AI Kit includes Intel-optimized versions of deep learning frameworks (such as TensorFlow and PyTorch) and high-performing Python libraries to streamline end-to-end data science and AI workflows on Intel architectures.

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,17 @@ 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 = {}
-
-    if "pt" in sys.argv:
-        sys.argv.remove("pt")
+    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"
-
-    if "tf" in sys.argv:
-        sys.argv.remove("tf")
+    elif os.environ.get("INC_TF_ONLY", False):
         cfg_key = "neural_compressor_tf"
+    else:
+        cfg_key = "neural_compressor"
 
+    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():