[pytorch] Updates PyTorch engine to 2.1.1 (#2864)

deepjavalibrary · Nov 21, 2023 · 5a9191e · 5a9191e
1 parent b5e5100
commit 5a9191e
Show file tree

Hide file tree

Showing 9 changed files with 56 additions and 45 deletions.
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
@@ -136,7 +136,7 @@ jobs:
           echo $PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
-          ./gradlew :engines:pytorch:pytorch-native:cleanJNI clean
+          ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           rm -rf ~/.djl.ai
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials

diff --git a/.github/workflows/nightly_publish.yml b/.github/workflows/nightly_publish.yml
@@ -122,11 +122,11 @@ jobs:
           ./gradlew :integration:test "-Dai.djl.default_engine=OnnxRuntime"
           ./gradlew :integration:clean
 
-  test-cuda-118:
+  test-cuda-121:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, gpu ]
     container:
-      image: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu18.04
+      image: nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu20.04
       options: --gpus all --runtime=nvidia
     timeout-minutes: 30
     needs: create-runners
@@ -163,7 +163,7 @@ jobs:
   publish:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: ubuntu-latest
-    needs: [ build, test-pytorch, test-tensorflow, test-aarch64, test-cuda-118 ]
+    needs: [ build, test-pytorch, test-tensorflow, test-aarch64, test-cuda-121 ]
     steps:
       - uses: actions/checkout@v3
       - name: Set up JDK 11
@@ -181,10 +181,9 @@ jobs:
       - name: Publish to snapshot repository
         if: ${{ github.event.inputs.mode == '' || github.event.inputs.mode == 'snapshot' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -Psnapshot
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.13.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.0.1 -Psnapshot
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.1 -Psnapshot
           ./gradlew clean engines:ml:xgboost:publish -Pgpu -Psnapshot
           ./gradlew clean publish -Psnapshot
           cd bom
@@ -197,10 +196,9 @@ jobs:
       - name: Publish to staging repository
         if: ${{ github.event.inputs.mode == 'staging' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -P${{ github.event.inputs.mode }}
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.13.1 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.0.1 -P${{ github.event.inputs.mode }}
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.1 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:ml:xgboost:publish -Pgpu -P${{ github.event.inputs.mode }}
           ./gradlew clean publish -P${{ github.event.inputs.mode }}
           cd bom
@@ -246,7 +244,7 @@ jobs:
   stop-runners:
     if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
     runs-on: [ self-hosted, scheduler ]
-    needs: [ create-runners, test-aarch64, test-cuda-118 ]
+    needs: [ create-runners, test-aarch64, test-cuda-121 ]
     steps:
       - name: Stop all instances
         run: |

diff --git a/bom/build.gradle b/bom/build.gradle
@@ -109,21 +109,21 @@ publishing {
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "linux-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "win-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-cu112mkl", "linux-x86_64", "${mxnet_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116", "linux-x86_64", "1.12.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116", "win-x86_64", "1.12.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116-precxx11", "linux-x86_64", "1.12.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "win-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "linux-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "win-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117-precxx11", "linux-x86_64", "1.13.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "win-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118-precxx11", "linux-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "linux-x86_64", "2.0.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "win-x86_64", "2.0.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118-precxx11", "linux-x86_64", "2.0.1")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "osx-x86_64", "${tensorflow_version}")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "linux-x86_64", "${tensorflow_version}")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "win-x86_64", "${tensorflow_version}")

diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
@@ -46,6 +46,8 @@ The following table illustrates which pytorch version that DJL supports:
 
 | PyTorch engine version | PyTorch native library version            |
 |------------------------|-------------------------------------------|
+| pytorch-engine:0.26.0  | 1.13.1, 2.0.1, **2.1.1**                  |
+| pytorch-engine:0.25.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
 | pytorch-engine:0.24.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
 | pytorch-engine:0.23.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
 | pytorch-engine:0.22.1  | 1.11.0, 1.12.1, **1.13.1**, 2.0.0         |

diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
@@ -24,7 +24,13 @@ processResources {
                 "osx-x86_64/cpu/libdjl_torch.dylib",
                 "win-x86_64/cpu/djl_torch.dll"
         ]
-        if (ptVersion.startsWith("2.0.")) {
+        if (ptVersion.startsWith("2.1.")) {
+            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
+            files.add("linux-x86_64/cu121/libdjl_torch.so")
+            files.add("linux-x86_64/cu121-precxx11/libdjl_torch.so")
+            files.add("win-x86_64/cu121/djl_torch.dll")
+            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
+        } else if (ptVersion.startsWith("2.0.")) {
             files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
             files.add("linux-x86_64/cu118/libdjl_torch.so")
             files.add("linux-x86_64/cu118-precxx11/libdjl_torch.so")

diff --git a/engines/pytorch/pytorch-native/CMakeLists.txt b/engines/pytorch/pytorch-native/CMakeLists.txt
@@ -64,7 +64,7 @@ add_library(djl_torch SHARED ${SOURCE_FILES})
 if(NOT BUILD_ANDROID)
     target_link_libraries(djl_torch "${TORCH_LIBRARIES}")
     target_include_directories(djl_torch PUBLIC build/include ${JNI_INCLUDE_DIRS} ${UTILS_INCLUDE_DIR})
-    set_property(TARGET djl_torch PROPERTY CXX_STANDARD 14)
+    set_property(TARGET djl_torch PROPERTY CXX_STANDARD 17)
     # We have to kill the default rpath and use current dir
     set(CMAKE_SKIP_RPATH TRUE)
     if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")

diff --git a/engines/pytorch/pytorch-native/build.gradle b/engines/pytorch/pytorch-native/build.gradle
@@ -24,6 +24,8 @@ if (project.hasProperty("cu11")) {
         FLAVOR = "cu117"
     } else if (VERSION.startsWith("2.0.")) {
         FLAVOR = "cu118"
+    } else if (VERSION.startsWith("2.1.")) {
+        FLAVOR = "cu121"
     } else {
         throw new GradleException("Unsupported PyTorch version: ${VERSION}")
     }
@@ -88,15 +90,17 @@ def prepareNativeLib(String binaryRoot, String ver) {
 
     def officialPytorchUrl = "https://download.pytorch.org/libtorch"
     def aarch64PytorchUrl = "https://djl-ai.s3.amazonaws.com/publish/pytorch"
-    String cu11
+    String cuda
     if (ver.startsWith("1.11.")) {
-        cu11 = "cu113"
+        cuda = "cu113"
     } else if (ver.startsWith("1.12.")) {
-        cu11 = "cu116"
+        cuda = "cu116"
     } else if (ver.startsWith("1.13.")) {
-        cu11 = "cu117"
+        cuda = "cu117"
     } else if (ver.startsWith("2.0.")) {
-        cu11 = "cu118"
+        cuda = "cu118"
+    } else if (ver.startsWith("2.1.")) {
+        cuda = "cu121"
     } else {
         throw new GradleException("Unsupported PyTorch version: ${ver}")
     }
@@ -105,10 +109,10 @@ def prepareNativeLib(String binaryRoot, String ver) {
             "cpu/libtorch-cxx11-abi-shared-with-deps-${ver}%2Bcpu.zip"        : "cpu/linux-x86_64",
             "cpu/libtorch-macos-${ver}.zip"                                   : "cpu/osx-x86_64",
             "cpu/libtorch-win-shared-with-deps-${ver}%2Bcpu.zip"              : "cpu/win-x86_64",
-            "${cu11}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cu11}.zip": "${cu11}/linux-x86_64",
-            "${cu11}/libtorch-win-shared-with-deps-${ver}%2B${cu11}.zip"      : "${cu11}/win-x86_64",
+            "${cuda}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cuda}.zip": "${cuda}/linux-x86_64",
+            "${cuda}/libtorch-win-shared-with-deps-${ver}%2B${cuda}.zip"      : "${cuda}/win-x86_64",
             "cpu/libtorch-shared-with-deps-${ver}%2Bcpu.zip"                  : "cpu-precxx11/linux-x86_64",
-            "${cu11}/libtorch-shared-with-deps-${ver}%2B${cu11}.zip"          : "${cu11}-precxx11/linux-x86_64"
+            "${cuda}/libtorch-shared-with-deps-${ver}%2B${cuda}.zip"          : "${cuda}-precxx11/linux-x86_64"
     ]
 
     def aarch64Files = [
@@ -138,17 +142,12 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
             from zipTree(file)
             into outputDir
         }
-        // CPU dependencies
-        copy {
-            from("${outputDir}/libtorch/lib/") {
-                include "libc10.*", "c10.dll", "libiomp5*.*", "libarm_compute*.*", "libgomp*.*", "libnvfuser_codegen.so", "libtorch.*", "libtorch_cpu.*", "torch.dll", "torch_cpu.dll", "fbgemm.dll", "asmjit.dll", "uv.dll", "nvfuser_codegen.dll"
-            }
-            into("${outputDir}/native/lib")
-        }
-        // GPU dependencies
+        delete "${outputDir}/libtorch/lib/*.lib"
+        delete "${outputDir}/libtorch/lib/*.a"
+
         copy {
             from("${outputDir}/libtorch/lib/") {
-                include "libtorch_cuda*.so", "torch_cuda*.dll", "libc10_cuda.so", "c10_cuda.dll", "libcaffe2_nvrtc.so", "libnvrtc*.so.*", "libcudart*.*", "*nvToolsExt*.*", "cudnn*.dll", "caffe2_nvrtc.dll", "nvrtc64*.dll", "uv.dll", "libcublas*", "zlibwapi.dll"
+                include "libarm_compute*", "libc10_cuda.so", "libc10.*", "libcaffe2_nvrtc.so", "libcu*", "libgfortran-*", "libgomp*", "libiomp*", "libnv*", "libopenblasp-*", "libtorch_cpu.*", "libtorch_cuda*.so", "libtorch.*", "asmjit.dll", "c10_cuda.dll", "c10.dll", "caffe2_nvrtc.dll", "cu*.dll", "fbgemm.dll", "nv*.dll", "torch_cpu.dll", "torch_cuda*.dll", "torch.dll", "uv.dll", "zlibwapi.dll"
             }
             into("${outputDir}/native/lib")
         }
@@ -287,9 +286,9 @@ tasks.register('uploadS3') {
                 "${BINARY_ROOT}/cpu/win-x86_64/native/lib/",
                 "${BINARY_ROOT}/cpu-precxx11/linux-aarch64/native/lib/",
                 "${BINARY_ROOT}/cpu-precxx11/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118/win-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118-precxx11/linux-x86_64/native/lib/"
+                "${BINARY_ROOT}/cu121/linux-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu121/win-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu121-precxx11/linux-x86_64/native/lib/"
         ]
         uploadDirs.each { item ->
             fileTree(item).files.name.each {

diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
@@ -23,7 +23,7 @@ ARCH=$4
 
 if [[ ! -d "libtorch" ]]; then
   if [[ $PLATFORM == 'linux' ]]; then
-    if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu113|cu116|cu117|cu118)$ ]]; then
+    if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu113|cu116|cu117|cu118|cu121)$ ]]; then
       echo "$FLAVOR is not supported."
       exit 1
     fi
@@ -62,6 +62,12 @@ mkdir classes
 javac -sourcepath ../../pytorch-engine/src/main/java/ ../../pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java -h include -d classes
 cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=${PT_VERSION} -DUSE_CUDA=$USE_CUDA ..
 cmake --build . --config Release -- -j "${NUM_PROC}"
+if [[ "$FLAVOR" = cu* ]]; then
+  # avoid link with libcudart.so.11.0
+  sed -i -r "s/\/usr\/local\/cuda(.{5})?\/lib64\/lib(cudart|nvrtc).so//g" CMakeFiles/djl_torch.dir/link.txt
+  rm libdjl_torch.so
+  . CMakeFiles/djl_torch.dir/link.txt
+fi
 
 if [[ $PLATFORM == 'darwin' ]]; then
   install_name_tool -add_rpath @loader_path libdjl_torch.dylib

diff --git a/gradle.properties b/gradle.properties
@@ -13,7 +13,7 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true
 
 djl_version=0.26.0
 mxnet_version=1.9.1
-pytorch_version=2.0.1
+pytorch_version=2.1.1
 tensorflow_version=2.10.1
 tflite_version=2.6.2
 trt_version=8.4.1