From 0a356010181651c1be0f045cdd01960dcbadd89b Mon Sep 17 00:00:00 2001
From: Mathieu Broillet <mathieu@broillet.ch>
Date: Wed, 28 Aug 2024 15:13:16 +0200
Subject: [PATCH] fix gpu targets and paths (workflow)

---
 .github/workflows/build-and_publish-prebuilts-whl.yml | 11 ++++-------
 prebuilts/Dockerfile                                  | 11 +++++++++--
 prebuilts/entrypoint.sh                               |  4 ++--
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/build-and_publish-prebuilts-whl.yml b/.github/workflows/build-and_publish-prebuilts-whl.yml
index d2db00e..aaeb400 100644
--- a/.github/workflows/build-and_publish-prebuilts-whl.yml
+++ b/.github/workflows/build-and_publish-prebuilts-whl.yml
@@ -36,12 +36,10 @@ jobs:
       - name: Build Docker image
         env:
           ROCM_VERSION: ${{ inputs.rocm_version }}
-          GPU_ARCH: ${{ inputs.gpu_arch }}
           TORCH_VERSION: ${{ inputs.torch_version }}
         run: |
           docker build \
             --build-arg ROCM_VERSION=${ROCM_VERSION} \
-            --build-arg GPU_ARCH=${GPU_ARCH} \
             --build-arg TORCH_VERSION=${TORCH_VERSION} \
             -t prebuilt-wheels ./prebuilts
 
@@ -49,12 +47,10 @@ jobs:
       - name: Run Docker container and generate wheels
         env:
           ROCM_VERSION: ${{ inputs.rocm_version }}
-          GPU_ARCH: ${{ inputs.gpu_arch }}
           TORCH_VERSION: ${{ inputs.torch_version }}
         run: |
           docker create --name prebuilt-container \
             -e ROCM_ARCH=${ROCM_ARCH} \
-            -e GPU_ARCH=${GPU_ARCH} \
             -e TORCH_VERSION=${TORCH_VERSION} \
             prebuilt-wheels
           docker start -a prebuilt-container
@@ -63,14 +59,15 @@ jobs:
       - name: Copy bitsandbytes wheel to host
         run: |
           mkdir -p /tmp/bitsandbytes/dist
-          docker cp prebuilt-container:/tmp/bitsandbytes/dist /tmp/bitsandbytes/dist
-      
+          docker cp prebuilt-container:/tmp/bitsandbytes/dist/ /tmp/bitsandbytes/dist/
+          ls -l /tmp/bitsandbytes/dist/
 
       # Step 6: Copy llama-cpp-python wheel artifact to host
       - name: Copy llama-cpp-python wheel to host
         run: |
           mkdir -p /tmp/llama-cpp-python/dist
-          docker cp prebuilt-container:/tmp/llama-cpp-python/dist /tmp/llama-cpp-python/dist
+          docker cp prebuilt-container:/tmp/llama-cpp-python/dist/ /tmp/llama-cpp-python/dist/
+          ls -l /tmp/llama-cpp-python/dist/
 
       # Step 7: Upload bitsandbytes wheel artifact
       - name: Upload bitsandbytes wheel
diff --git a/prebuilts/Dockerfile b/prebuilts/Dockerfile
index 6478fac..e980d81 100644
--- a/prebuilts/Dockerfile
+++ b/prebuilts/Dockerfile
@@ -1,8 +1,15 @@
-FROM rocm/dev-ubuntu-22.04:6.1.2
+FROM rocm/dev-ubuntu-22.04:${ROCM_VERSION}
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
-    PYTHONIOENCODING=UTF-8
+    PYTHONIOENCODING=UTF-8 \
+
+# For bitsandbytes
+ENV BNB_GPU_TARGETS="gfx803;gfx900;gfx906;gfx908;gfx90a;gfx1010;gfx1030;gfx1100;gfx1101;gfx1102"
+
+# For LLAMA
+ENV GPU_TARGETS="gfx803 gfx900 gfx906 gfx908 gfx90a gfx1010 gfx1030 gfx1100 gfx1101 gfx1102"
+
 
 WORKDIR /tmp
 
diff --git a/prebuilts/entrypoint.sh b/prebuilts/entrypoint.sh
index 12cc6e7..5903fb6 100755
--- a/prebuilts/entrypoint.sh
+++ b/prebuilts/entrypoint.sh
@@ -6,7 +6,7 @@ git clone --depth 1 -b multi-backend-refactor https://github.com/bitsandbytes-fo
 cd /tmp/bitsandbytes
 pip3 install -r requirements-dev.txt
 ## Build
-cmake -DCOMPUTE_BACKEND=hip -S . -DBNB_ROCM_ARCH=${ROCM_ARCH}
+cmake -DCOMPUTE_BACKEND=hip -S . -DBNB_ROCM_ARCH=${BNB_GPU_TARGETS}
 make
 python3.10 setup.py bdist_wheel --universal
 
@@ -15,4 +15,4 @@ python3.10 setup.py bdist_wheel --universal
 ## Clone repo and install python requirements
 git clone --recurse-submodules https://github.com/abetlen/llama-cpp-python.git /tmp/llama-cpp-python
 cd /tmp/llama-cpp-python
-CMAKE_ARGS="-DGGML_HIPBLAS=on" python3.10 -m build --wheel
\ No newline at end of file
+CMAKE_ARGS="-D GGML_HIPBLAS=on -D AMDGPU_TARGETS=${GPU_TARGETS}" FORCE_CMAKE=1 python3.10 -m build --wheel
\ No newline at end of file