From 58a4a4d4c19bda86d489ac171fa10f3ddb27c9d6 Mon Sep 17 00:00:00 2001 From: Jeff Rasley Date: Wed, 8 Mar 2023 09:55:41 -0800 Subject: [PATCH] Fix issue between our abstract accelerator and colossalai's version of op_builder (#2963) Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> --- accelerator/cuda_accelerator.py | 5 +++-- op_builder/__init__.py | 3 +++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py index cdf4cd447764..945ba42a3a9c 100644 --- a/accelerator/cuda_accelerator.py +++ b/accelerator/cuda_accelerator.py @@ -224,8 +224,9 @@ def on_accelerator(self, tensor): def op_builder_dir(self): try: - # during installation time op_builder is visible, otherwise return deepspeed.ops.op_builder - import op_builder # noqa: F401 + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 return "op_builder" except ImportError: return "deepspeed.ops.op_builder" diff --git a/op_builder/__init__.py b/op_builder/__init__.py index 51546be6ccaf..b41f074d9ba0 100644 --- a/op_builder/__init__.py +++ b/op_builder/__init__.py @@ -8,6 +8,9 @@ from .builder import get_default_compute_capabilities, OpBuilder +# Do not remove, required for abstract accelerator to detect if we have a deepspeed or 3p op_builder +__deepspeed__ = True + # List of all available op builders from deepspeed op_builder try: import deepspeed.ops.op_builder # noqa: F401