Merge branch 'master' into cifar

microsoft · Aug 22, 2024 · 33aae5c · 33aae5c
2 parents 1a924a9 + b81b197
commit 33aae5c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml
@@ -21,7 +21,7 @@ on:
       - "deepspeed/runtime/zero/parameter_offload.py"
       - "deepspeed/runtime/pipe/engine.py"
       - "deepspeed/runtime/utils.py"
-      - "opbuilder/xpu/**"
+      - "op_builder/xpu/**"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}

diff --git a/op_builder/evoformer_attn.py b/op_builder/evoformer_attn.py
@@ -52,11 +52,27 @@ def is_compatible(self, verbose=False):
             if verbose:
                 self.warning("Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH")
             return False
-        with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f:
-            if '3.1.0' not in f.read():
+        if os.path.exists(f'{self.cutlass_path}/CHANGELOG.md'):
+            with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f:
+                if '3.1.0' not in f.read():
+                    if verbose:
+                        self.warning("Please use CUTLASS version >= 3.1.0")
+                    return False
+        else:
+            # pip install nvidia-cutlass package
+            try:
+                import cutlass
+            except ImportError:
+                if verbose:
+                    self.warning("Please pip install nvidia-cutlass if trying to pre-compile kernels")
+                return False
+            cutlass_major, cutlass_minor = cutlass.__version__.split('.')[:2]
+            cutlass_compatible = (int(cutlass_major) >= 3 and int(cutlass_minor) >= 1)
+            if not cutlass_compatible:
                 if verbose:
                     self.warning("Please use CUTLASS version >= 3.1.0")
                 return False
+
         cuda_okay = True
         if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
             sys_cuda_major, _ = installed_cuda_version()

diff --git 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.")