Fix issue between our abstract accelerator and colossalai's version o…

…f op_builder (microsoft#2963) Co-authored-by: Logan Adams <[email protected]>
EleutherAI · Mar 8, 2023 · 58a4a4d · 58a4a4d
1 parent 6379def
commit 58a4a4d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git 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
@@ -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