Option to install upstream PyTorch from nightly wheels

.github/actions/setup-pytorch/action.yml

@@ -16,6 +16,9 @@ inputs:
   ref:
     description: Branch, tag, commit id
     default: ""
+  mode:
+    description: Source or wheels
+    default: source
 runs:
   using: "composite"
   steps:
@@ -71,7 +74,7 @@ runs:
     - name: Generate PyTorch cache key
       shell: bash
       run: |
-        PYTORCH_CACHE_KEY=$(echo $PYTHON_VERSION $PYTORCH_COMMIT_ID ${{ hashFiles('scripts/patch-pytorch.sh') }} | sha256sum - | cut -d\  -f1)
+        PYTORCH_CACHE_KEY=$(echo $PYTHON_VERSION $PYTORCH_COMMIT_ID ${{ hashFiles('scripts/patch-pytorch.sh') }} ${{ inputs.mode }} | sha256sum - | cut -d\  -f1)
         echo "PYTORCH_CACHE_KEY=$PYTORCH_CACHE_KEY" | tee -a "$GITHUB_ENV"
 
     - name: Load PyTorch from a cache
@@ -90,11 +93,12 @@ runs:
       with:
         repository: ${{ env.PYTORCH_REPO }}
         ref: ${{ env.PYTORCH_COMMIT_ID }}
-        submodules: recursive
+        # To build PyTorch from source we need all submodules, they are not required for benchmarks
+        submodules: ${{ inputs.mode == 'source' && 'recursive' || 'false' }}
         path: pytorch
 
     - name: Apply additional PR patches
-      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.repository == 'pytorch/pytorch' }}
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.repository == 'pytorch/pytorch' && inputs.mode == 'source' }}
       shell: bash
       run: |
         cd pytorch
@@ -108,7 +112,7 @@ runs:
         pip install 'numpy<2.0.0'
 
     - name: Build PyTorch
-      if: ${{ steps.pytorch-cache.outputs.status == 'miss' }}
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.mode == 'source' }}
       shell: bash
       run: |
         source ${{ inputs.oneapi }}/setvars.sh
@@ -117,11 +121,24 @@ runs:
         pip install -r requirements.txt
         python setup.py bdist_wheel
 
-    - name: Install PyTorch
+    - name: Install PyTorch (built from source)
+      if: ${{ inputs.mode == 'source' }}
       shell: bash
       run: |
         source ${{ inputs.oneapi }}/setvars.sh
         pip install pytorch/dist/*.whl
+
+    - name: Install PyTorch (from wheels)
+      if: ${{ inputs.mode == 'wheels' }}
+      shell: bash
+      run: |
+        source ${{ inputs.oneapi }}/setvars.sh
+        pip install torch --index-url https://download.pytorch.org/whl/nightly/xpu
+
+    - name: Get PyTorch version
+      shell: bash
+      run: |
+        source ${{ inputs.oneapi }}/setvars.sh
         PYTORCH_VERSION="$(python -c 'import torch;print(torch.__version__)')"
         echo "PYTORCH_VERSION=$PYTORCH_VERSION" | tee -a "$GITHUB_ENV"
 

.github/workflows/build-test-gpu.yml

@@ -12,6 +12,13 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: choice
+        options:
+          - source
+          - wheels
+        default: source
       upload_test_reports:
         description: Upload test reports
         type: boolean
@@ -46,6 +53,7 @@ jobs:
       device: ${{ inputs.runner_label }}
       runner_label: ${{ inputs.runner_label }}
       pytorch_ref: ${{ inputs.pytorch_ref }}
+      pytorch_mode: ${{ inputs.pytorch_mode || 'source' }}
       python_version: ${{ matrix.python }}
       upload_test_reports: ${{ inputs.upload_test_reports }}
       ignore_errors: ${{ inputs.ignore_errors }}

.github/workflows/build-test-reusable.yml

@@ -20,6 +20,10 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: string
+        default: "source"
       python_version:
         description: Python version
         type: string
@@ -96,6 +100,7 @@ jobs:
         with:
           repository: pytorch/pytorch
           ref: ${{ inputs.pytorch_ref }}
+          mode: ${{ inputs.pytorch_mode }}
 
       - name: Install pass_rate dependencies
         run: |

.github/workflows/build-test.yml

@@ -12,6 +12,13 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: choice
+        options:
+          - source
+          - wheels
+        default: source
       upload_test_reports:
         description: Upload test reports
         type: boolean
@@ -120,6 +127,7 @@ jobs:
       driver_version: ${{ matrix.driver }}
       runner_label: ${{ inputs.runner_label }}
       pytorch_ref: ${{ inputs.pytorch_ref }}
+      pytorch_mode: ${{ inputs.pytorch_mode || 'source' }}
       python_version: ${{ matrix.python }}
       upload_test_reports: ${{ inputs.upload_test_reports || false }}
       ignore_errors: ${{ inputs.ignore_errors || false }}