Use AWS Neuron sdk 2.18 (#547)

* chore(tgi): bump router version to 1.4.4 * chore: bump AWS Neuron SDK to 2.18 * chore(tgi): use AWS Neuron SDK 2.18 * ci: AWS Neuron SDK 2.18 on inf2 * docs: update TGI container
huggingface · Apr 8, 2024 · 09ddd67 · 09ddd67
1 parent e3bc576
commit 09ddd67
Show file tree

Hide file tree

Showing 11 changed files with 25 additions and 25 deletions.
diff --git a/.github/workflows/test_inf2.yml b/.github/workflows/test_inf2.yml
@@ -31,7 +31,7 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.0.0 aws-neuronx-oci-hook=2.2.45.0 aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b aws-neuronx-collectives=2.20.11.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/.github/workflows/test_inf2_export.yml b/.github/workflows/test_inf2_export.yml
@@ -31,7 +31,7 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.0.0 aws-neuronx-oci-hook=2.2.45.0 aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b aws-neuronx-collectives=2.20.11.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/.github/workflows/test_inf2_full_export.yml b/.github/workflows/test_inf2_full_export.yml
@@ -29,7 +29,7 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.0.0 aws-neuronx-oci-hook=2.2.45.0 aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b aws-neuronx-collectives=2.20.11.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/.github/workflows/test_inf2_inference.yml b/.github/workflows/test_inf2_inference.yml
@@ -31,7 +31,7 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.0.0 aws-neuronx-oci-hook=2.2.45.0 aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b aws-neuronx-collectives=2.20.11.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/.github/workflows/test_inf2_tgi.yml b/.github/workflows/test_inf2_tgi.yml
@@ -33,7 +33,7 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.0.0 aws-neuronx-oci-hook=2.2.45.0 aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b aws-neuronx-collectives=2.20.11.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/Makefile b/Makefile
@@ -40,7 +40,7 @@ PACKAGE_FILES = $(PACKAGE_PYTHON_FILES)  \
 $(PACKAGE_DIST) $(PACKAGE_WHEEL): $(PACKAGE_FILES)
 	python -m build
 
-TGI_VERSION ?= 1.4.1
+TGI_VERSION ?= 1.4.4
 
 neuronx-tgi: $(PACKAGE_DIST)
 	docker build --rm -f text-generation-inference/Dockerfile \

diff --git a/docs/source/containers.mdx b/docs/source/containers.mdx
@@ -10,20 +10,20 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o
 specific language governing permissions and limitations under the License.
 -->
 
-# Optimum Neuron Container 
+# Optimum Neuron Container
 
 We provide pre-built Optimum Neuron containers for Amazon SageMaker. These containers come with all of the Hugging Face libraries and dependencies pre-installed, so you can start using them right away.
 We have containers for training and inference, and optimized text generation containers with TGI. The table is up to date and only includes the latest versions of each container. You can find older versions in the [Deep Learning Container Release Notes](https://github.com/aws/deep-learning-containers/releases?q=hf-neuronx&expanded=true)
 
-We recommend using the `sagemaker` Python SDK to retrieve the image URI for the container you want to use. 
+We recommend using the `sagemaker` Python SDK to retrieve the image URI for the container you want to use.
 
 ## Available Optimum Neuron Containers
 
 | Type                       | Optimum Version | Image URI                                   |
 |-----------------------------|-----------------|---------------------------------------------|
 | Training  | 0.0.13           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training-neuronx:1.13.1-neuronx-py310-sdk2.15.0-ubuntu20.04`   |
 | Inference      | 0.0.13           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference-neuronx:1.13.1-neuronx-py310-sdk2.15.0-ubuntu20.04`      |
-| Text Generation Inference        | 0.0.17           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:1.13.1-optimum0.0.17-neuronx-py310-ubuntu22.04`        |
+| Text Generation Inference        | 0.0.20           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:1.13.1-optimum0.0.20-neuronx-py310-ubuntu22.04`        |
 
 
-Please replace `763104351884` with the correct [AWS account ID](https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/image_uri_config/huggingface-neuronx.json) and `region` with the AWS region you are working in. 
+Please replace `763104351884` with the correct [AWS account ID](https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/image_uri_config/huggingface-neuronx.json) and `region` with the AWS region you are working in.
diff --git a/optimum/neuron/version.py b/optimum/neuron/version.py
@@ -14,4 +14,4 @@
 
 __version__ = "0.0.21.dev0"
 
-__sdk_version__ = "2.17.0"
+__sdk_version__ = "2.18.0"
diff --git a/setup.py b/setup.py
@@ -47,20 +47,20 @@
     "quality": QUALITY_REQUIRES,
     "neuron": [
         "wheel",
-        "torch-neuron==1.13.1.2.9.17.0",
+        "torch-neuron==1.13.1.2.9.74.0",
         "torch==1.13.1.*",
-        "neuron-cc[tensorflow]==1.21.0.0",
+        "neuron-cc[tensorflow]==1.22.0.0",
         "protobuf",
         "torchvision",
     ],
     "neuronx": [
         "wheel",
-        "neuronx-cc==2.12.68.0",
-        "torch-neuronx==1.13.1.1.13.1",
-        "transformers-neuronx==0.9.474",
+        "neuronx-cc==2.13.66.0",
+        "torch-neuronx==1.13.1.1.14.0",
+        "transformers-neuronx==0.10.0.21",
         "torch==1.13.1.*",
         "torchvision==0.14.*",
-        "neuronx_distributed==0.6.0",
+        "neuronx_distributed==0.7.0",
     ],
     "diffusers": ["diffusers ~= 0.26.1", "peft"],
     "sentence-transformers": ["sentence-transformers >= 2.2.0"],

diff --git a/text-generation-inference/Dockerfile b/text-generation-inference/Dockerfile
@@ -91,19 +91,19 @@ RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEU
 # Install neuronx packages
 RUN apt-get update -y \
  && apt-get install -y --no-install-recommends \
-    aws-neuronx-dkms=2.15.9.0 \
-    aws-neuronx-collectives=2.20.11.0-c101c322e \
-    aws-neuronx-runtime-lib=2.20.11.0-b7d33e68b \
-    aws-neuronx-tools=2.17.0.0 \
+    aws-neuronx-dkms=2.16.7.0 \
+    aws-neuronx-collectives=2.20.22.0-c101c322e \
+    aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 \
+    aws-neuronx-tools=2.17.1.0 \
     && rm -rf /var/lib/apt/lists/* \
     && apt-get clean
 
 ENV PATH="/opt/bin/:/opt/aws/neuron/bin:${PATH}"
 
 RUN pip3 install \
-    neuronx-cc==2.12.68.0 \
-    torch-neuronx==1.13.1.1.13.1 \
-    transformers-neuronx==0.9.474 \
+    neuronx-cc==2.13.66.0 \
+    torch-neuronx==1.13.1.1.14.0 \
+    transformers-neuronx==0.10.0.21 \
     --extra-index-url=https://pip.repos.neuron.amazonaws.com
 
 # Install HuggingFace packages

diff --git a/text-generation-inference/server/Makefile b/text-generation-inference/server/Makefile
@@ -2,7 +2,7 @@
 pkg_name := text_generation_server
 BUILDDIR ?= $(CURDIR)/build
 VERSION ?= 0.0.1
-TGI_VERSION ?= 1.4.1
+TGI_VERSION ?= 1.4.4
 mkfile_path := $(abspath $(lastword $(MAKEFILE_LIST)))
 mkfile_dir := $(dir $(mkfile_path))
 pkg_dir := $(BUILDDIR)/$(pkg_name)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -14,4 +14,4 @@

		__version__ = "0.0.21.dev0"

		__sdk_version__ = "2.17.0"
		__sdk_version__ = "2.18.0"