Use AWS Neuron SDK 2.14 (#236)

* chore: use neuron SDK 2.14 * chore(tgi): use neuron SDK 2.14 * fix(generate): use new neuronx compiler flag
huggingface · Sep 20, 2023 · d2f5c9d · d2f5c9d
1 parent 74d0905
commit d2f5c9d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
@@ -206,7 +206,7 @@ def _from_pretrained(
 
         # Compile the Neuron model (if present compiled artifacts will be reloaded instead of compiled)
         neuron_cc_flags = os.environ.get("NEURON_CC_FLAGS", "")
-        os.environ["NEURON_CC_FLAGS"] = neuron_cc_flags + " --model-type=transformer-inference"
+        os.environ["NEURON_CC_FLAGS"] = neuron_cc_flags + " --model-type=transformer"
         neuronx_model.to_neuron()
         os.environ["NEURON_CC_FLAGS"] = neuron_cc_flags
 

diff --git a/setup.py b/setup.py
@@ -53,12 +53,12 @@
     ],
     "neuronx": [
         "wheel",
-        "neuronx-cc>=2.9",
-        "torch-neuronx>=1.13.1.1.10.1",
-        "transformers-neuronx>=0.6.106",
+        "neuronx-cc>=2.10",
+        "torch-neuronx>=1.13.1.1.11.0",
+        "transformers-neuronx>=0.7.84",
         "torch==1.13.1.*",
         "torchvision==0.14.*",
-        "neuronx_distributed>=0.3.0",
+        "neuronx_distributed>=0.4.0",
     ],
     "diffusers": ["diffusers"],
 }

diff --git a/text-generation-inference/Dockerfile b/text-generation-inference/Dockerfile
@@ -89,19 +89,19 @@ RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEU
 # Install neuronx 2.12.2 packages
 RUN apt-get update -y \
  && apt-get install -y --no-install-recommends \
-    aws-neuronx-dkms=2.12.18.0 \
-    aws-neuronx-collectives=2.16.16.0-e59c7bb3e \
-    aws-neuronx-runtime-lib=2.16.14.0-61fdc395f \
-    aws-neuronx-tools=2.13.4.0 \
+    aws-neuronx-dkms=2.13.4.0 \
+    aws-neuronx-collectives=2.17.9.0-fb6d14044 \
+    aws-neuronx-runtime-lib=2.17.7.0-df62e3f70 \
+    aws-neuronx-tools=2.14.6.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.9.0.40 \
-    torch-neuronx==1.13.1.1.10.1 \
-    transformers-neuronx==0.6.106 \
+    neuronx-cc==2.10.0.34 \
+    torch-neuronx==1.13.1.1.11.0 \
+    transformers-neuronx==0.7.84 \
     --extra-index-url=https://pip.repos.neuron.amazonaws.com
 
 # Install HuggingFace packages