- slim build down targeting colab env

- enable pytorch nightly installation in ci

Author: deftlion

.github/workflows/publish.yml

@@ -4,6 +4,7 @@
 name: Create Release
 
 on:
+  workflow_dispatch:
   push:
     tags:
       - v*
@@ -48,9 +49,9 @@ jobs:
       fail-fast: false
       matrix:
           os: ['ubuntu-20.04']
-          python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
-          pytorch-version: ['2.4.0']  # Must be the most recent version that meets requirements-cuda.txt.
-          cuda-version: ['11.8', '12.1']
+          python-version: ['3.10']
+          pytorch-version: ['nightly']  # Must be the most recent version that meets requirements-cuda.txt.
+          cuda-version: ['12.1']
 
     steps:
       - name: Checkout

.github/workflows/scripts/build.sh

@@ -14,7 +14,6 @@ $python_executable -m pip install -r requirements-cuda.txt
 # Limit the number of parallel jobs to avoid OOM
 export MAX_JOBS=1
 # Make sure release wheels are built for the following architectures
-export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
-export VLLM_FA_CMAKE_GPU_ARCHES="80-real;90-real"
+export TORCH_CUDA_ARCH_LIST="7.5"
 # Build
 $python_executable setup.py bdist_wheel --dist-dir=dist

.github/workflows/scripts/pytorch-install.sh

@@ -6,8 +6,12 @@ cuda_version=$3
 
 # Install torch
 $python_executable -m pip install numpy pyyaml scipy ipython mkl mkl-include ninja cython typing pandas typing-extensions dataclasses setuptools && conda clean -ya
-$python_executable -m pip install torch==${pytorch_version}+cu${cuda_version//./} --extra-index-url https://download.pytorch.org/whl/cu${cuda_version//./}
 
+if [[ $pytorch_version == "nightly" ]]; then
+  $python_executable -m pip install torch --pre --upgrade --index-url https://download.pytorch.org/whl/nightly/cu121
+else
+  $python_executable -m pip install torch==${pytorch_version}+cu${cuda_version//./} --extra-index-url https://download.pytorch.org/whl/cu${cuda_version//./}
+fi
 # Print version information
 $python_executable --version
 $python_executable -c "import torch; print('PyTorch:', torch.__version__)"