Make installable without torch (microsoft#5001)

Fixes microsoft#4984
ROCm · May 9, 2024 · df1efca · df1efca
1 parent dc4a3af
commit df1efca
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -16,9 +16,6 @@ jobs:
         ref: "master"
     - id: setup-venv
       uses: ./.github/workflows/setup-venv
-    - name: Install torch
-      run: |
-        pip install torch
     - name: Get release version from tag
       run: |
         echo "RELEASE_VERSION=${GITHUB_REF#refs/*/v}" >> $GITHUB_ENV

diff --git a/op_builder/xpu/async_io.py b/op_builder/xpu/async_io.py
@@ -5,7 +5,6 @@
 
 import distutils.spawn
 import subprocess
-import torch
 
 from .builder import OpBuilder
 
@@ -33,6 +32,7 @@ def include_paths(self):
         return ['csrc/aio/py_lib', 'csrc/aio/common']
 
     def cxx_args(self):
+        import torch
         # -O0 for improved debugging, since performance is bound by I/O
         CPU_ARCH = self.cpu_arch()
         SIMD_WIDTH = self.simd_width()

diff --git a/op_builder/xpu/builder.py b/op_builder/xpu/builder.py
@@ -11,9 +11,9 @@
     # 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 # type: ignore
-    from op_builder.builder import OpBuilder, TORCH_MAJOR, TORCH_MINOR
+    from op_builder.builder import OpBuilder
 except ImportError:
-    from deepspeed.ops.op_builder.builder import OpBuilder, TORCH_MAJOR, TORCH_MINOR
+    from deepspeed.ops.op_builder.builder import OpBuilder
 
 
 class SYCLOpBuilder(OpBuilder):
@@ -35,6 +35,10 @@ def builder(self):
         return dpcpp_ext
 
     def version_dependent_macros(self):
+        try:
+            from op_builder.builder import TORCH_MAJOR, TORCH_MINOR
+        except ImportError:
+            from deepspeed.ops.op_builder.builder import TORCH_MAJOR, TORCH_MINOR
         # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456
         version_ge_1_1 = []
         if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):