[INF2][DLC] Update Neuron to 2.10 (#681)

* add a few inf2 build flags * update
deepjavalibrary · May 2, 2023 · a5607d1 · a5607d1
1 parent d68cedb
commit a5607d1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 12 deletions.
diff --git a/serving/docker/pytorch-inf2.Dockerfile b/serving/docker/pytorch-inf2.Dockerfile
@@ -13,7 +13,8 @@ FROM ubuntu:20.04
 ARG djl_version=0.23.0~SNAPSHOT
 ARG torch_version=1.13.1
 ARG python_version=3.8
-ARG torch_neuronx_version=1.13.0.1.4.0
+ARG torch_neuronx_version=1.13.1.1.7.0
+ARG transformers_neuronx_version=0.3.32
 ARG transformers_version=4.28.1
 ARG accelerate_version=0.18.0
 EXPOSE 8080
@@ -52,9 +53,10 @@ RUN mkdir -p /opt/djl/bin && cp scripts/telemetry.sh /opt/djl/bin && \
     echo "${djl_version} inf2" > /opt/djl/bin/telemetry && \
     scripts/install_python.sh ${python_version} && \
     scripts/install_djl_serving.sh $djl_version && \
-    scripts/install_inferentia2.sh $torch_neuronx_version && \
-    pip install transformers==${transformers_version} accelerate==${accelerate_version} && \
-    pip install git+https://github.com/aws-neuron/transformers-neuronx.git@v2.8.0 && \
+    scripts/install_inferentia2.sh && \
+    pip install transformers==${transformers_version} accelerate==${accelerate_version} \
+    neuronx-cc==2.6.* torch_neuronx==${torch_neuronx_version} transformers-neuronx==${transformers_neuronx_version} \
+    --extra-index-url=https://pip.repos.neuron.amazonaws.com && \
     scripts/install_s5cmd.sh x64 && \
     scripts/patch_oss_dlc.sh python && \
     useradd -m -d /home/djl djl && \

diff --git a/serving/docker/scripts/install_inferentia2.sh b/serving/docker/scripts/install_inferentia2.sh
@@ -1,5 +1,4 @@
 #!/usr/bin/env bash
-TORCH_NEURONX_VERSION=$1
 
 set -ex
 
@@ -15,12 +14,9 @@ echo "deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main" >/etc
 curl -L https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
 # https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/releasecontent.html#inf2-packages
-apt-get update -y && apt-get install -y linux-headers-$(uname -r) && apt-get install -y aws-neuronx-dkms=2.8.4.* \
-    aws-neuronx-collectives=2.12.27.* \
-    aws-neuronx-runtime-lib=2.12.16.* \
-    aws-neuronx-tools=2.9.5.*
+apt-get update -y && apt-get install -y linux-headers-$(uname -r) && apt-get install -y aws-neuronx-dkms=2.9.* \
+    aws-neuronx-collectives=2.13.* \
+    aws-neuronx-runtime-lib=2.13.* \
+    aws-neuronx-tools=2.10.*
 
 export PATH=/opt/aws/neuron/bin:$PATH
-
-python3 -m pip install numpy awscli
-python3 -m pip install neuronx-cc==2.4.* torch_neuronx==${TORCH_NEURONX_VERSION} --extra-index-url=https://pip.repos.neuron.amazonaws.com