Avoid errors in DeepSpeed build on ROCm

rraminen · Aug 28, 2024 · 2e1d1ce · 2e1d1ce
1 parent 1041c8a
commit 2e1d1ce
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/op_builder/gds.py b/op_builder/gds.py
@@ -43,8 +43,13 @@ def is_compatible(self, verbose=False):
                 self.warning("Please install torch if trying to pre-compile GDS")
             return False
 
-        CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
-        CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
+        if not self.is_rocm_pytorch():
+            CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
+            CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
+        else:
+            CUDA_HOME = torch.utils.cpp_extension.ROCM_HOME
+            CUDA_LIB64 = os.path.join(CUDA_HOME, "lib")
+
         gds_compatible = self.has_function(funcname="cuFileDriverOpen",
                                            libraries=("cufile", ),
                                            library_dirs=(