Repository: deepjavalibrary/djl
Pull request #2037 (2022-09-27 10:31:43)

Add support for Apples' Metal Performance Shaders (MPS) in pytorch
## Description ##
Pytorch 1.12 supports Apple’s Metal Performance Shaders (MPS) for accelerated training/inference. A simple test on HuggingFace models shows 2x-3x improvements on the inference latency over the CPU cores on an Apple M1Pro with 8 CPU and 14 GPU cores:
```
--------------------------------------------------------------------------------
Model:              roberta-hf
Device:             cpu()
Samples:            34
Model Input Length: 384
Metrics:
	Total: 140.206 ± 5.880      msec
	Inference: 138.794 ± 5.713      msec
	Preprocess: 0.441 ± 0.553      msec
	Postprocess: 0.824 ± 0.567      msec
--------------------------------------------------------------------------------
Model:              roberta-hf
Device:             mps(-1)
Samples:            34
Model Input Length: 384
Metrics:
	Total: 65.500 ± 9.166      msec
	Inference: 41.235 ± 6.198      msec
	Preprocess: 1.059 ± 0.338      msec
	Postprocess: 23.206 ± 6.794      msec

--------------------------------------------------------------------------------
```
This improvement has been requested by me in the DJL help slack channel, as well as more recently in #2018 . 

Brief description of what this PR is about

- Adds a new "MPS" device type and maps to the corresponding pytorch device type.

Some caveats:
- The "MPS" mode only works if  `torch::jit::load(path, device, map);`  is called with `device= torch::nullopt` as in torch  the model gets deserialized in 'legacy' mode, which only support CPU and GPU devices. The model get desiralized on the CPU and then can be converted to MPS. For this, in traced torch model's directory  in the file "serving.properties" the `option.mapLocation=false` should be set, which triggers this implementation already present in DJL native module.
Alternatively, the the following check can be added the pytorch-native module: https://github.com/deepjavalibrary/djl/blob/f097502b9fe7607203320c4aecee14c759e29d5e/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc#L49-L52

``` 
if (jmap_location && (device.is_cpu() || device.is_cuda())) {
```

-  An env variable s of hould be set`PYTORCH_ENABLE_MPS_FALLBACK=1` as a work around a pytorch implementation deficiency of 'aten::index.Tensor' for MPS, see https://github.com/pytorch/pytorch/issues/77764





Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index fea755e68..ef942fe17 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -197,6 +197,45 @@ jobs:
           aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*"
 
+  build-pytorch-jni-arm64-macos:
+    if: github.repository == 'deepjavalibrary/djl'
+    runs-on: [ self-hosted, ARM64, macOS ]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up JDK 11
+        uses: actions/setup-java@v3
+        with:
+          java-version: 11
+          distribution: corretto
+          architecture: aarch64
+      - uses: actions/cache@v2
+        with:
+          path: ~/.gradle/caches
+          key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
+          restore-keys: |
+            ${{ runner.os }}-gradle-
+      - name: Release JNI prep
+        shell: bash
+        run: |
+          PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
+          export PYTORCH_VERSION=${PYTORCH_VERSION:-$(cat gradle.properties | awk -F '=' '/pytorch_version/ {print $2}')}
+          echo $PYTORCH_VERSION
+          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION
+          ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: us-east-2
+      - name: Copy files to S3 with the AWS CLI
+        shell: bash
+        run: |
+          PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
+          PYTORCH_VERSION=${PYTORCH_VERSION:-$(cat gradle.properties | awk -F '=' '/pytorch_version/ {print $2}')}
+          aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*"
+
   create-aarch64-runner:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, scheduler ]
diff --git a/.github/workflows/native_s3_huggingface.yml b/.github/workflows/native_s3_huggingface.yml
index d20d98396..ce7f08a10 100644
--- a/.github/workflows/native_s3_huggingface.yml
+++ b/.github/workflows/native_s3_huggingface.yml
@@ -119,6 +119,42 @@ jobs:
           aws s3 sync extensions/tokenizers/jnilib s3://djl-ai/publish/tokenizers/${TOKENIZERS_VERSION}/jnilib/
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/tokenizers/${TOKENIZERS_VERSION}/jnilib/*"
 
+  build-tokenizers-jni-arm64-osx:
+    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
+    runs-on: [ self-hosted, ARM64, macOS ]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up JDK 11
+        uses: actions/setup-java@v3
+        with:
+          java-version: 11
+          distribution: corretto
+          architecture: aarch64
+      - uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+      - uses: actions/cache@v2
+        with:
+          path: ~/.gradle/caches
+          key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
+          restore-keys: |
+            ${{ runner.os }}-gradle-
+      - name: Release JNI prep
+        run: |
+          ./gradlew :extensions:tokenizers:compileJNI
+          ./gradlew -Pjni :extensions:tokenizers:test
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: us-east-2
+      - name: Copy files to S3 with the AWS CLI
+        run: |
+          TOKENIZERS_VERSION="$(cat gradle.properties | awk -F '=' '/tokenizers_version/ {print $2}')"
+          aws s3 sync extensions/tokenizers/jnilib s3://djl-ai/publish/tokenizers/${TOKENIZERS_VERSION}/jnilib/
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/tokenizers/${TOKENIZERS_VERSION}/jnilib/*"
+
   create-aarch64-runner:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, scheduler ]
diff --git a/android/README.md b/android/README.md
index 0786d5ce7..fd99b283b 100644
--- a/android/README.md
+++ b/android/README.md
@@ -18,8 +18,8 @@ In gradle, you can include the snapshot repository and add the 4 modules in your
 dependencies {
     implementation "ai.djl:api:0.19.0-SNAPSHOT"
     implementation "ai.djl.android:core:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.pytorch:pytorch-engine:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.android:pytorch-native:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.android:onnxruntime:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.pytorch:pytorch-engine:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.android:pytorch-native:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.android:onnxruntime:0.19.0-SNAPSHOT"
 }
 ```
diff --git a/docs/development/dependency_management.md b/docs/development/dependency_management.md
index 9b4cda14c..267995268 100644
--- a/docs/development/dependency_management.md
+++ b/docs/development/dependency_management.md
@@ -6,7 +6,7 @@ This document outlines how DJL manages its dependencies and how developers can f
 
 We try our best to minimize the external dependencies. The core DJL library only has 4 external dependencies:
 
-- [sl4fj-api](https://mvnrepository.com/artifact/ai.djl/api)
+- [sl4fj-api](https://mvnrepository.com/artifact/org.slf4j/slf4j-api)
 - [Gson](https://mvnrepository.com/artifact/com.google.code.gson/gson)
 - [JNA](https://mvnrepository.com/artifact/net.java.dev.jna/jna)
 - [apache-commons-compress](https://mvnrepository.com/artifact/org.apache.commons/commons-compress)
@@ -58,17 +58,20 @@ See [How to use DJL's BOM](../../bom/README.md#how-to-use-djls-bom) for detail.
 | [ai.djl.pytorch](https://search.maven.org/search?q=g:ai.djl.pytorch)             | [pytorch-engine](../../engines/pytorch/pytorch-engine/README.md#installation)                                              | PyTorch engine adapter                                                                           |
 |                                                                                  | [pytorch-model-zoo](../../engines/pytorch/pytorch-model-zoo/README.md#installation)                                        | Contains state of the art PyTorch torch script models                                            |
 |                                                                                  | [pytorch-native-cpu(osx-x86_64)](../../engines/pytorch/pytorch-engine/README.md#macos)                                     | Contains PyTorch native library for macOS                                                        |
+|                                                                                  | [pytorch-native-cpu(osx-aarch64)](../../engines/pytorch/pytorch-engine/README.md#macos-m1)                                 | Contains PyTorch native library for macOS M1                                                     |
 |                                                                                  | [pytorch-native-cpu(win-x86_64)](../../engines/pytorch/pytorch-engine/README.md#windows-cpu)                               | Contains PyTorch native library for Windows                                                      |
 |                                                                                  | [pytorch-native-cpu(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#linux-cpu)                               | Contains PyTorch native library for Linux                                                        |
-|                                                                                  | [pytorch-native-cu113(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#linux-gpu)                             | Contains PyTorch native library for Linux with CUDA 11.3                                         |
-|                                                                                  | [pytorch-native-cu113(win-x86_64)](../../engines/pytorch/pytorch-engine/README.md#windows-gpu)                             | Contains PyTorch native library for Windows with CUDA 11.3                                       |
+|                                                                                  | [pytorch-native-cu116(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#linux-gpu)                             | Contains PyTorch native library for Linux with CUDA 11.6                                         |
+|                                                                                  | [pytorch-native-cu116(win-x86_64)](../../engines/pytorch/pytorch-engine/README.md#windows-gpu)                             | Contains PyTorch native library for Windows with CUDA 11.6                                       |
 |                                                                                  | [pytorch-native-cu102(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#linux-gpu)                             | Contains PyTorch native library for Linux with CUDA 10.2                                         |
 |                                                                                  | [pytorch-native-cu102(win-x86_64)](../../engines/pytorch/pytorch-engine/README.md#windows-gpu)                             | Contains PyTorch native library for Windows with CUDA 10.2                                       |
-|                                                                                  | [pytorch-native-cu113-precxx11(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#for-pre-cxx11-build)          | Contains PyTorch native library for Linux with CUDA 11.3                                         |
+|                                                                                  | [pytorch-native-cu116-precxx11(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#for-pre-cxx11-build)          | Contains PyTorch native library for Linux with CUDA 11.6                                         |
 |                                                                                  | [pytorch-native-cpu-precxx11(linux-x86_64)](../../engines/pytorch/pytorch-engine/README.md#for-pre-cxx11-build)            | Contains PyTorch native library for centOS 7 and Ubuntu 14.04                                    |
 |                                                                                  | [pytorch-native-cpu-precxx11(linux-aarch64)](../../engines/pytorch/pytorch-engine/README.md#for-aarch64-build)             | Contains PyTorch native library for Linux ARM                                                    |
 |                                                                                  | [pytorch-jni](../../engines/pytorch/pytorch-engine/README.md)                                                              | Contains PyTorch JNI native library                                                              |
 |                                                                                  | pytorch-native-auto (deprecated)                                                                                           | No longer needed since DJL 0.15.0                                                                |
+|                                                                                  | pytorch-native-cu113(linux-x86_64) (deprecated)                                                                            | Contains PyTorch native library for Linux with CUDA 11.3                                         |
+|                                                                                  | pytorch-native-cu113(win-x86_64) (deprecated)                                                                              | Contains PyTorch native library for Windows with CUDA 11.3                                       |
 |                                                                                  | pytorch-native-cu111(linux-x86_64) (deprecated)                                                                            | Contains PyTorch native library for Linux with CUDA 11.1                                         |
 |                                                                                  | pytorch-native-cu111(win-x86_64) (deprecated)                                                                              | Contains PyTorch native library for Windows with CUDA 11.1                                       |
 |                                                                                  | pytorch-native-cu101(linux-x86_64) (deprecated)                                                                            | Contains PyTorch native library <= 1.7.1 for Linux with CUDA 10.1                                |
@@ -115,6 +118,8 @@ See [How to use DJL's BOM](../../bom/README.md#how-to-use-djls-bom) for detail.
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.tensorrt](https://search.maven.org/search?q=g:ai.djl.tensorrt)           | [tensorrt](../../engines/tensorrt/README.md#installation)                                                                  | TensorRT engine adapter                                                                          |
 |                                                                                  |                                                                                                                            |                                                                                                  |
+| [ai.djl.java](https://search.maven.org/search?q=g:ai.djl.java)                   | [java](https://github.com/deepjavalibrary/djl-serving/tree/master/engines/java/README.md#installation)                     | Java engine adapter                                                                              |
+|                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.python](https://search.maven.org/search?q=g:ai.djl.python)               | [python](https://github.com/deepjavalibrary/djl-serving/tree/master/engines/python/README.md#installation)                 | Python engine adapter                                                                            |
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.ml.lightgbm](https://search.maven.org/search?q=g:ai.djl.ml.lightgbm)     | [lightgbm](../../engines/ml/lightgbm/README.md#installation)                                                               | LightGBM engine adapter                                                                          |
@@ -126,7 +131,7 @@ See [How to use DJL's BOM](../../bom/README.md#how-to-use-djls-bom) for detail.
 |                                                                                  | [pytorch-native](../../android/pytorch-native/README.md)                                                                   | Contains DJL PyTorch Android native package                                                      |
 |                                                                                  | [onnxruntime](../../engines/onnxruntime/onnxruntime-android/README.md#installation)                                        | Contains DJL Onnxruntime Android package                                                         |
 |                                                                                  |                                                                                                                            |                                                                                                  |
-| [ai.djl.audio](https://search.maven.org/search?q=g:ai.djl.audio)                 | [audio](../../extensions/audio/README.md#installation)                                                                     | Contains classes that make it easy for DJL to access AWS services                                |
+| [ai.djl.audio](https://search.maven.org/search?q=g:ai.djl.audio)                 | [audio](../../extensions/audio/README.md#installation)                                                                     | DJL audio support extension                                                                      |
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.aws](https://search.maven.org/search?q=g:ai.djl.aws)                     | [aws-ai](../../extensions/aws-ai/README.md#installation)                                                                   | Contains classes that make it easy for DJL to access AWS services                                |
 |                                                                                  |                                                                                                                            |                                                                                                  |
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index 31d71b6df..632cf4e5f 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -99,6 +99,28 @@ For macOS, you can use the following library:
 </dependency>
 ```
 
+### macOS M1
+For macOS M1, you can use the following library:
+
+- ai.djl.pytorch:pytorch-jni:1.12.1-0.19.0
+- ai.djl.pytorch:pytorch-native-cpu:1.12.1:osx-aarch64
+
+```xml
+<dependency>
+    <groupId>ai.djl.pytorch</groupId>
+    <artifactId>pytorch-native-cpu</artifactId>
+    <classifier>osx-aarch64</classifier>
+    <version>1.12.1</version>
+    <scope>runtime</scope>
+</dependency>
+<dependency>
+    <groupId>ai.djl.pytorch</groupId>
+    <artifactId>pytorch-jni</artifactId>
+    <version>1.12.1-0.19.0</version>
+    <scope>runtime</scope>
+</dependency>
+```
+
 ### Linux
 For the Linux platform, you can choose between CPU, GPU. If you have NVIDIA [CUDA](https://en.wikipedia.org/wiki/CUDA)
 installed on your GPU machine, you can use one of the following library:
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
index 78b620345..a7396dd8a 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
@@ -32,8 +32,10 @@ public final class PtDeviceType {
             return 0;
         } else if (Device.Type.GPU.equals(deviceType)) {
             return 1;
+        } else if ("mps".equals(deviceType)) {
+            return 13;
         } else {
-            throw new IllegalArgumentException("Unsupported device: " + device.toString());
+            throw new IllegalArgumentException("Unsupported device: " + device);
         }
     }
 
@@ -49,6 +51,8 @@ public final class PtDeviceType {
                 return Device.Type.CPU;
             case 1:
                 return Device.Type.GPU;
+            case 13:
+                return "mps";
             default:
                 throw new IllegalArgumentException("Unsupported deviceType: " + deviceType);
         }


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
new file mode 100644
index 000000000..8b4e2326f
--- /dev/null
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.pytorch.integration;
+
+import ai.djl.Device;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+
+import org.testng.Assert;
+import org.testng.SkipException;
+import org.testng.annotations.Test;
+
+public class MpsTest {
+
+    @Test
+    public void testMps() {
+        if (!"aarch64".equals(System.getProperty("os.arch"))
+                || !System.getProperty("os.name").startsWith("Mac")) {
+            throw new SkipException("MPS test requires M1 macOS.");
+        }
+
+        Device device = Device.of("mps", -1);
+        try (NDManager manager = NDManager.newBaseManager(device)) {
+            NDArray array = manager.zeros(new Shape(1, 2));
+            Assert.assertEquals(array.getDevice().getDeviceType(), "mps");
+        }
+    }
+}