[docker] Bump up neuron to 2.21.1

serving/docker/pytorch-inf2.Dockerfile

@@ -19,15 +19,15 @@ ARG torch_version=2.1.2
 ARG torchvision_version=0.16.2
 
 # Neuron SDK components
-ARG neuronsdk_version=2.21.0
+ARG neuronsdk_version=2.21.1
 ARG torch_neuronx_version=2.1.2.2.4.0
-ARG transformers_neuronx_version=0.13.322
-ARG neuronx_distributed_version=0.10.0
-ARG neuronx_distributed_inference_version=0.1.0
-ARG neuronx_cc_version=2.16.345.0
-ARG neuronx_cc_stubs_version=2.16.345.0
+ARG transformers_neuronx_version=0.13.380
+ARG neuronx_distributed_version=0.10.1
+ARG neuronx_distributed_inference_version=0.1.1
+ARG neuronx_cc_version=2.16.372.0
+ARG neuronx_cc_stubs_version=2.16.372.0
 ARG torch_xla_version=2.1.6
-ARG libneuronxla_version=2.1.681.0
+ARG libneuronxla_version=2.1.714.0
 
 ARG transformers_version=4.45.2
 ARG tokenizers_version=0.20.3

serving/docker/scripts/install_inferentia2.sh

@@ -16,8 +16,8 @@ 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 aws-neuronx-collectives=2.23.133.0* \
-  aws-neuronx-runtime-lib=2.23.110.0* \
+apt-get update -y && apt-get install -y aws-neuronx-collectives=2.23.135.0* \
+  aws-neuronx-runtime-lib=2.23.112.0* \
   aws-neuronx-tools=2.20.204.0
 
 # TODO: Remove this hack after aws-neuronx-dkms install no longer throws an error, this bypasses the `set -ex`