$CUDA_HOME_PATH

Author: Borda

.azure/gpu-tests-fabric.yml

@@ -86,8 +86,23 @@ jobs:
         condition: endsWith(variables['Agent.JobName'], 'future')
         displayName: "extend env. vars 4 future"
 
+      - bash: |
+          set -e
+          # Try to find nvcc and set CUDA_HOME so that deepspeed can be installed
+          NVCC_PATH=$(find /usr/local -name "nvcc" | head -n 1)
+          if [[ -z "$NVCC_PATH" ]]; then
+            echo "nvcc not found, skipping setting CUDA_HOME"
+          else
+            CUDA_HOME_PATH=$(dirname $(dirname "$NVCC_PATH"))
+            echo "Found nvcc at: ${NVCC_PATH}"
+            echo "Setting CUDA_HOME to: ${CUDA_HOME_PATH}"
+            echo "##vso[task.setvariable variable=CUDA_HOME]${CUDA_HOME_PATH}"
+          fi
+        displayName: "Set CUDA_HOME for DeepSpeed"
+
       - bash: |
           echo $(DEVICES)
+          echo $CUDA_HOME_PATH
           echo $CUDA_VISIBLE_DEVICES
           echo $CUDA_VERSION_MM
           echo $TORCH_URL

.azure/gpu-tests-pytorch.yml

@@ -90,8 +90,23 @@ jobs:
         condition: endsWith(variables['Agent.JobName'], 'future')
         displayName: "extend env. vars 4 future"
 
+      - bash: |
+          set -e
+          # Try to find nvcc and set CUDA_HOME so that deepspeed can be installed
+          NVCC_PATH=$(find /usr/local -name "nvcc" | head -n 1)
+          if [[ -z "$NVCC_PATH" ]]; then
+            echo "nvcc not found, skipping setting CUDA_HOME"
+          else
+            CUDA_HOME_PATH=$(dirname $(dirname "$NVCC_PATH"))
+            echo "Found nvcc at: ${NVCC_PATH}"
+            echo "Setting CUDA_HOME to: ${CUDA_HOME_PATH}"
+            echo "##vso[task.setvariable variable=CUDA_HOME]${CUDA_HOME_PATH}"
+          fi
+        displayName: "Set CUDA_HOME for DeepSpeed"
+
       - bash: |
           echo $(DEVICES)
+          echo $CUDA_HOME_PATH
           echo $CUDA_VISIBLE_DEVICES
           echo $CUDA_VERSION_MM
           echo $TORCH_URL
@@ -212,4 +227,4 @@ jobs:
         workingDirectory: examples
         env:
           PL_USE_MOCKED_MNIST: "1"
-        displayName: "Testing: PyTorch examples"
+        displayName: "Testing: PyTorch examples"