This document is relevant for: Trn2, Trn3

Install PyTorch via Deep Learning Container#

Deploy PyTorch with Neuron support using pre-configured Docker images from AWS ECR.

⏱️ Estimated time: 10 minutes

Note

For a non-containerized setup, consider the DLAMI-based installation or manual installation instead.

What are Neuron DLCs?#

AWS Neuron Deep Learning Containers (DLCs) are pre-configured Docker images with the Neuron SDK and ML frameworks pre-installed. They provide Docker-based isolation, reproducibility, and portability across deployment platforms including EC2, EKS, ECS, and SageMaker.

Available PyTorch Neuron DLC images:

Container Type

Use Case

Links

PyTorch Inference vLLM (NeuronX)

LLM serving with vLLM

vLLM images

PyTorch Training (TorchNeuron)

Native PyTorch training with torch-neuronx 2.14 and TorchTitan

torch-neuronx images

Prerequisites#

Requirement

Details

Instance Type

Trn2 or Trn3

Docker

Docker Engine installed and running

AWS CLI

Configured with ECR access permissions

Neuron Driver

aws-neuronx-dkms installed on the host

Quick Start: vLLM Inference Container#

The fastest way to get started with LLM inference on Neuron:

# Pull the vLLM inference container
docker pull public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.29.0.1.2.0-neuronx-py314-sdk2.33.0-ubuntu26.04

# Run with Neuron device access
docker run -it --device=/dev/neuron0 \
  public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.29.0.1.2.0-neuronx-py314-sdk2.33.0-ubuntu26.04

For the latest image tags and a step-by-step walkthrough, see Quickstart: Configure and deploy a vLLM server using Neuron Deep Learning Container (DLC).

Note

The image tags above are examples. For the latest available images, see the Neuron DLC repository.

Note

Starting with Neuron 2.33.0, Neuron DLC images are published on Ubuntu 26.04 only. Image tags carry a -ubuntu26.04 suffix (for example, ...-neuronx-py314-sdk2.33.0-ubuntu26.04). If you need an Ubuntu 24.04 image, use a tag from Neuron 2.32.0 or earlier. Check the Neuron DLC repository for the exact published tags.

Customizing a DLC#

You can extend a Neuron DLC with additional packages by creating a custom Dockerfile:

FROM public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.29.0.1.2.0-neuronx-py314-sdk2.33.0-ubuntu26.04

# Install additional packages
RUN pip install transformers datasets

# Copy your application code
COPY app/ /app/

For more details, see Customize Neuron DLC.

Deployment Platforms#

Neuron DLCs can be deployed across multiple AWS services:

Amazon EC2

Deploy containers directly on EC2 instances with Neuron devices.

Amazon EKS

Run containers on managed Kubernetes with the Neuron device plugin.

Amazon ECS

Deploy containers using Amazon Elastic Container Service.

Next Steps#

This document is relevant for: Trn2, Trn3