ROCm 6.1 package

mlc-ai · Aug 10, 2024 · 4a8ff9a · 4a8ff9a
1 parent f4fed53
commit 4a8ff9a
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 31 deletions.
diff --git a/.github/workflows/wheel_manylinux_nightly.yaml b/.github/workflows/wheel_manylinux_nightly.yaml
@@ -13,27 +13,7 @@ on:
     - cron: '0 6 * * *' # 6 AM UTC
 
 jobs:
-  Build:
-    strategy:
-      matrix:
-        pkg_kind: ['nightly']
-        # matrix of build configs
-        config:
-          - gpu: 'none'
-            image: 'mlcaidev/package-cpu:7a83f97'
-          - gpu: 'cuda-12.1'
-            image: 'mlcaidev/package-cu121:7a83f97'
-          - gpu: 'cuda-12.2'
-            image: 'mlcaidev/package-cu122:7a83f97'
-          # - gpu: 'rocm-5.6'
-          #   image: 'mlcaidev/package-rocm56:7a83f97'
-          # - gpu: 'rocm-5.7'
-          #   image: 'mlcaidev/package-rocm57:7a83f97'
-        exclude:
-          - pkg_kind: ${{ github.event_name == 'schedule' && 'stable' || '' }}
-
-    runs-on: ubuntu-latest
-
+  Build_template:
     steps:
     - name: Reclaim disk space
       run: |
@@ -85,3 +65,37 @@ jobs:
           ${{ matrix.pkg_kind == 'nightly' && 'mlc-llm/python/repaired_wheels/*.whl' || '' }}
         tag_name: v0.9.dev0
         prerelease: true
+
+  Build_CUDA:
+    <<: *Build_template
+    strategy:
+      matrix:
+        pkg_kind: ['nightly']
+        # matrix of build configs
+        config:
+          - gpu: 'none'
+            image: 'mlcaidev/package-cpu:16b1781'
+          - gpu: 'cuda-12.1'
+            image: 'mlcaidev/package-cu121:16b1781'
+          - gpu: 'cuda-12.2'
+            image: 'mlcaidev/package-cu122:16b1781'
+        exclude:
+          - pkg_kind: ${{ github.event_name == 'schedule' && 'stable' || '' }}
+
+    runs-on: ubuntu-latest
+
+  Build_ROCm:
+    <<: *Build_template
+    strategy:
+      matrix:
+        pkg_kind: ['nightly']
+        # matrix of build configs
+        config:
+          - gpu: 'rocm-6.1'
+            image: 'mlcaidev/package-rocm61:16b1781'
+          - gpu: 'rocm-6.2'
+            image: 'mlcaidev/package-rocm62:16b1781'
+        exclude:
+          - pkg_kind: ${{ github.event_name == 'schedule' && 'stable' || '' }}
+
+    runs-on: [self-hosted, Linux, X64]
diff --git a/scripts/build_mlc_ai_wheel_manylinux.sh b/scripts/build_mlc_ai_wheel_manylinux.sh
@@ -6,7 +6,7 @@ source /opt/rh/gcc-toolset-11/enable # GCC-11 is the hightest GCC version compat
 function usage() {
 	echo "Usage: $0 [--gpu GPU-VERSION]"
 	echo
-	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 rocm-5.6 rocm-5.7}"
+	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 rocm-6.1 rocm-6.2}"
 	echo -e "\tSpecify the GPU version (CUDA/ROCm) in the TVM (default: none)."
 }
 
@@ -44,7 +44,7 @@ function audit_mlc_ai_wheel() {
 TVM_PYTHON_DIR="/workspace/tvm/python"
 PYTHON_VERSIONS_CPU=("3.7" "3.8" "3.9" "3.10" "3.11" "3.12")
 PYTHON_VERSIONS_GPU=("3.7" "3.8" "3.9" "3.10" "3.11" "3.12")
-GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-5.6" "rocm-5.7")
+GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-6.1" "rocm-6.2")
 GPU="none"
 
 while [[ $# -gt 0 ]]; do
@@ -71,7 +71,7 @@ done
 if ! in_array "${GPU}" "${GPU_OPTIONS[*]}"; then
 	echo "Invalid GPU option: ${GPU}"
 	echo
-	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-5.6" "rocm-5.7"}'
+	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-6.1" "rocm-6.2"}'
 	exit -1
 fi
 
@@ -105,9 +105,9 @@ elif [[ ${GPU} == cuda* ]]; then
 fi
 
 if [[ ${GPU} == rocm* ]]; then
-	echo set\(USE_LLVM \"/opt/rocm/llvm/bin/llvm-config --ignore-libllvm --link-static\"\) >>config.cmake
+	echo set\(USE_LLVM \"llvm-config --ignore-libllvm --link-static\"\) >>config.cmake
 	echo set\(USE_ROCM ON\) >>config.cmake
-	echo set\(USE_RCCL /opt/rocm/rccl/ \) >>config.cmake
+	echo set\(USE_RCCL /opt/rocm/ \) >>config.cmake
 elif [[ ${GPU} == cuda* ]]; then
 	echo set\(USE_LLVM \"llvm-config --ignore-libllvm --link-static\"\) >>config.cmake
 	echo set\(USE_CUDA ON\) >>config.cmake

diff --git a/scripts/build_mlc_llm_wheel_manylinux.sh b/scripts/build_mlc_llm_wheel_manylinux.sh
@@ -6,7 +6,7 @@ source /opt/rh/gcc-toolset-11/enable # GCC-11 is the hightest GCC version compat
 function usage() {
 	echo "Usage: $0 [--gpu GPU-VERSION]"
 	echo
-	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 rocm-5.6 rocm-5.7}"
+	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 rocm-6.1 rocm-6.2}"
 	echo -e "\tSpecify the GPU version (CUDA/ROCm) in the MLC-LLM (default: none)."
 }
 
@@ -44,7 +44,7 @@ function audit_mlc_llm_wheel() {
 MLC_LLM_PYTHON_DIR="/workspace/mlc-llm/python"
 PYTHON_VERSIONS_CPU=("3.7" "3.8" "3.9" "3.10" "3.11" "3.12")
 PYTHON_VERSIONS_GPU=("3.7" "3.8" "3.9" "3.10" "3.11" "3.12")
-GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-5.6" "rocm-5.7")
+GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-6.1" "rocm-6.2")
 GPU="none"
 
 while [[ $# -gt 0 ]]; do
@@ -71,7 +71,7 @@ done
 if ! in_array "${GPU}" "${GPU_OPTIONS[*]}"; then
 	echo "Invalid GPU option: ${GPU}"
 	echo
-	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-5.6" "rocm-5.7"}'
+	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "rocm-6.1" "rocm-6.2"}'
 	exit -1
 fi
 

diff --git a/scripts/sync_package.py b/scripts/sync_package.py
@@ -167,8 +167,8 @@ def main():
             "cuda-11.8",
             "cuda-12.1",
             "cuda-12.2",
-            "rocm-5.6",
-            "rocm-5.7",
+            "rocm-6.1",
+            "rocm-6.2",
         ],
         help="GPU (CUDA/ROCm) version to be linked to the resultant binaries,"
         "or none, to disable CUDA/ROCm. Defaults to none.",