From b81b197ec333331a92189ea17b35737104715a78 Mon Sep 17 00:00:00 2001 From: Yizhou Wang Date: Fri, 23 Aug 2024 01:24:38 +0800 Subject: [PATCH] [XPU] API align with new intel pytorch extension release (#6395) with intel_extension_for_pytorch new release, we changed deepspeed kernel API. Need to align with upstream DeepSpeed xpu op builder. Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> --- op_builder/xpu/inference.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/op_builder/xpu/inference.py b/op_builder/xpu/inference.py index 9114dcc2c315..a9ac4f84c2ca 100644 --- a/op_builder/xpu/inference.py +++ b/op_builder/xpu/inference.py @@ -30,7 +30,10 @@ def cxx_args(self): def load(self): try: - import intel_extension_for_pytorch.deepspeed - return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + import intel_extension_for_pytorch + if hasattr(intel_extension_for_pytorch, "deepspeed"): + return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + else: + return intel_extension_for_pytorch.xpu.deepspeed except ImportError: raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.")