Optimum Neuron Container

We provide pre-built Optimum Neuron containers for Amazon SageMaker. These containers come with all of the Hugging Face libraries and dependencies pre-installed, so you can start using them right away. We have containers for training and inference, and optimized text generation containers with TGI. The table is up to date and only includes the latest versions of each container. You can find older versions in the Deep Learning Container Release Notes

We recommend using the sagemaker Python SDK to retrieve the image URI for the container you want to use.

Available Optimum Neuron Containers

Type	Optimum Version	Image URI
Training	0.0.13	763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training-neuronx:1.13.1-neuronx-py310-sdk2.15.0-ubuntu20.04
Inference	0.0.13	763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference-neuronx:1.13.1-neuronx-py310-sdk2.15.0-ubuntu20.04
Text Generation Inference	0.0.20	763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:1.13.1-optimum0.0.20-neuronx-py310-ubuntu22.04

Please replace 763104351884 with the correct AWS account ID and region with the AWS region you are working in.