From d2f5c9d70cc6377fd5b333003f2c961933e3e117 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Wed, 20 Sep 2023 17:30:50 +0200 Subject: [PATCH] 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 --- optimum/neuron/modeling_decoder.py | 2 +- setup.py | 8 ++++---- text-generation-inference/Dockerfile | 14 +++++++------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py index e1cd5beb5..74fc67b87 100644 --- 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 index 61f31592b..b34b7645f 100644 --- 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 index 047185f41..07fef24f8 100644 --- 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