dev: added cuda12.4 build support

vectorch-ai · Jun 30, 2024 · 8070796 · 8070796
1 parent 3c514fc
commit 8070796
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 1 deletion.
diff --git a/.github/workflows/publish_devel_image.yml b/.github/workflows/publish_devel_image.yml
@@ -27,6 +27,22 @@ jobs:
       - name: Create cache directory
         run: mkdir -p $CI_CACHE_DIR/.buildx-cache
 
+      - name: Build devel image for cuda 12.4
+        uses: docker/build-push-action@v5
+        with:
+          context: ./docker
+          file: ./docker/Dockerfile.devel
+          push: true
+          cache-from: type=local,src=$CI_CACHE_DIR/.buildx-cache
+          cache-to: type=local,dest=$CI_CACHE_DIR/.buildx-cache
+          build-args: |
+            UBUNTU_VERSION=22.04
+            CUDA_VERSION=12.4
+            GCC_VERSION=12
+          tags: |
+            vectorchai/scalellm_devel:cuda12.4-ubuntu22.04
+            vectorchai/scalellm_devel:cuda12.4
+
       - name: Build devel image for cuda 12.1
         uses: docker/build-push-action@v5
         with:

diff --git a/.github/workflows/publish_manylinux_image.yml b/.github/workflows/publish_manylinux_image.yml
@@ -27,6 +27,19 @@ jobs:
       - name: Create cache directory
         run: mkdir -p $CI_CACHE_DIR/.buildx-cache
 
+      - name: Build base for cuda 12.4
+        uses: docker/build-push-action@v5
+        with:
+          context: ./docker
+          file: ./docker/Dockerfile.manylinux
+          push: true
+          cache-from: type=local,src=$CI_CACHE_DIR/.buildx-cache
+          cache-to: type=local,dest=$CI_CACHE_DIR/.buildx-cache
+          build-args: |
+            CUDA_VERSION=12.4
+          tags: |
+            vectorchai/scalellm_manylinux:cuda12.4
+
       - name: Build base for cuda 12.1
         uses: docker/build-push-action@v5
         with:

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -189,7 +189,14 @@ if (DEFINED ENV{LIBTORCH_ROOT})
   message(STATUS "Using libtorch at $ENV{LIBTORCH_ROOT}")
 else()
   include(FetchContent)
-  if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.1)
+  if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.4)
+    # download nightly libtorch with cuda 12.4 from pytorch.org
+    if (USE_CXX11_ABI)
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/nightly/cu124/libtorch-cxx11-abi-shared-with-deps-latest.zip")
+    else()
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/nightly/cu124/libtorch-shared-with-deps-latest.zip")
+    endif()
+  elseif(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.1)
     # download libtorch 2.3 with cuda 12.1 from pytorch.org
     if (USE_CXX11_ABI)
       set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.3.0%2Bcu121.zip")