diff --git a/neural_compressor/__init__.py b/neural_compressor/__init__.py index 5ee86bf561a..47401052d2f 100644 --- a/neural_compressor/__init__.py +++ b/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 diff --git a/setup.py b/setup.py index d0be2291d6f..77097e18176 100644 --- a/setup.py +++ b/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 {} diff --git a/test/3x/torch/algorithms/fp8_quant/conftest.py b/test/3x/torch/algorithms/fp8_quant/conftest.py index 3497af8b3f4..8143a41e6d8 100644 --- a/test/3x/torch/algorithms/fp8_quant/conftest.py +++ b/test/3x/torch/algorithms/fp8_quant/conftest.py @@ -1,5 +1,6 @@ # Called once at the beginning of the test session def pytest_sessionstart(): + import os import habana_frameworks.torch.core as htcore import torch @@ -7,6 +8,8 @@ def pytest_sessionstart(): # 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) diff --git a/test/3x/torch/quantization/fp8_quant/conftest.py b/test/3x/torch/quantization/fp8_quant/conftest.py index 966b4edb296..1072052eeb9 100644 --- a/test/3x/torch/quantization/fp8_quant/conftest.py +++ b/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():