This document is relevant for: Trn2, Trn3

Update PyTorch in a Deep Learning Container

Update your DLC-based PyTorch Neuron deployment to the latest release.

Update the container image

DLC images are versioned and tagged with the Neuron SDK version. To update, pull the latest image tag from ECR:

# vLLM Inference
docker pull public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:<new_image_tag>

Replace <new_image_tag> with the tag for the desired SDK version (for example, 0.29.0.1.2.0-neuronx-py314-sdk2.33.0-ubuntu26.04).

Important

Starting with Neuron 2.33.0, Neuron DLC images are published on Ubuntu 26.04 only, and their tags end in -ubuntu26.04. If you update from a 2.32.0 or earlier image, the new image uses a different base OS and Python version (Python 3.14), so rebuild any custom images that you derived from the earlier image. To stay on an Ubuntu 24.04 image, use a tag from Neuron 2.32.0 or earlier. The host OS does not need to match the container OS.

Check available tags at the ECR Public Gallery:

• PyTorch vLLM Inference

For the full list of available images and tags, see Neuron Deep Learning Containers.

Update Neuron driver on the host

The Neuron driver runs on the host, not inside the container. Update it separately when moving to a new Neuron SDK release.

Ubuntu 26.04

sudo apt-get update
sudo apt-get install -y aws-neuronx-dkms

Ubuntu 24.04

sudo apt-get update
sudo apt-get install -y aws-neuronx-dkms

Amazon Linux 2023

sudo dnf install -y aws-neuronx-dkms

Verify the update

Launch the new container and verify:

docker run -it \
  --device=/dev/neuron0 \
  --cap-add SYS_ADMIN \
  --cap-add IPC_LOCK \
  public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:<new_image_tag> \
  bash

Inside the container:

python3 -c "import torch; import torch_neuronx; print(f'PyTorch {torch.__version__}, torch-neuronx {torch_neuronx.__version__}')"
neuron-ls

⚠️ Troubleshooting: Version mismatch between host driver and container

If you see runtime errors after updating the container image but not the host driver:

1. Check the host driver version: modinfo neuron on the host

2. Update the host driver to match the SDK version in the container

3. Reboot if the driver update requires it: sudo reboot

This document is relevant for: Trn2, Trn3