diff --git a/optimum/intel/openvino/trainer.py b/optimum/intel/openvino/trainer.py index 51edb4e9bc..e16bf0f58d 100644 --- a/optimum/intel/openvino/trainer.py +++ b/optimum/intel/openvino/trainer.py @@ -81,6 +81,7 @@ from optimum.exporters.onnx import OnnxConfig from ..utils.constant import _TASK_ALIASES +from ..utils.import_utils import is_transformers_version from .configuration import OVConfig from .quantization import OVDataLoader from .training_args import OVTrainingArguments @@ -285,9 +286,16 @@ def _inner_training_loop( if args.gradient_checkpointing: self.model.gradient_checkpointing_enable() - if self.args.local_rank != -1: - if self.compression_controller is not None: - self.compression_controller.distributed() + if is_transformers_version("<", "4.29.0"): + is_distributed = self.args.local_rank != -1 + else: + from accelerate.utils import DistributedType + + is_distributed = self.args.distributed_state.distributed_type != DistributedType.NO + + if self.compression_controller is not None and is_distributed: + self.compression_controller.distributed() + model = self._wrap_model(self.model_wrapped) if is_sagemaker_mp_enabled() and resume_from_checkpoint is not None: diff --git a/setup.py b/setup.py index c3e0503527..b94283fb80 100644 --- a/setup.py +++ b/setup.py @@ -33,7 +33,7 @@ ], "openvino": ["openvino>=2023.0.0.dev20230217", "onnx", "onnxruntime"], "nncf": ["nncf>=2.4.0", "openvino-dev>=2023.0.0.dev20230217"], - "ipex": ["intel-extension-for-pytorch"], + "ipex": ["intel-extension-for-pytorch", "onnx"], "diffusers": ["diffusers"], "quality": QUALITY_REQUIRE, "tests": TESTS_REQUIRE,