Available DLCs on AWS

Below you can find a listing of our latest Deep Learning Containers (DLCs) available on AWS.

For each supported combination of use-case (training, inference), accelerator type (CPU, GPU, Neuron), and framework (PyTorch, TGI, TEI) containers are created.

Neuron DLCs for training and inference on AWS Trainium and AWS Inferentia instances can be found in the Optimum Neuron documentation.

If you want to keep track of all our available DLCs, you can also check the AWS Deep Learning Containers releases page.

Training

For training, the DLCs are available for PyTorch via Transformers. They include GPUs and AWS AI chips support, with libraries such as TRL, Sentence Transformers, or Diffusers.

You can also keep track of the latest PyTorch Training DLC releases here.

Inference

PyTorch Inference

For inference, there is a general-purpose PyTorch inference DLC, for serving models trained with any of those frameworks mentioned before on CPU, GPU, and AWS AI chips.

vLLM

In case you want to serve text generation models with vLLM, there are specific DLCs available for GPU and AWS AI chips.

SGLang

There is also a specific DLC for serving models with SGLang on GPU.

Text Embeddings Inference

Finally, there is the Text Embeddings Inference (TEI) DLC for high-performance serving of embedding models on CPU and GPU.

FAQ

How to choose the right inference container for my use case?

Note: See here for the list of supported task in the inference toolkit.

Note: Browse through the Hub to see if your model is tagged “text-generation-inference” or “text-embeddings-inference”.

How to find the URI of my container?

The URI is built with an AWS account ID and an AWS region. Those two values need to be replaced depending on your use case.

Let’s say you want to use the training DLC for GPUs:

• dlc-aws-account-id: The AWS account ID of the account that owns the ECR repository. You can find them in the here
• region: The AWS region where you want to use it.

How to find the URI of my container but simpler?

The image_uris.retrieve helper from sagemaker.core is not always up to date but it is much simpler than reconstructing the image URI yourself. The processor (CPU/GPU/Neuron) is inferred from the instance_type.

These docs and examples use the SageMaker Python SDK v3, which introduces a new framework-agnostic API built around ModelBuilder (inference) and ModelTrainer (training), replacing the v2 HuggingFaceModel and HuggingFace classes. Install it with pip install "sagemaker>=3.0.0". In v2 these URIs were retrieved with get_huggingface_llm_image_uri, which has been removed.

from sagemaker.core import image_uris

region = "us-east-1"  # the AWS region where you want to use the container

# TGI (LLMs) on GPU
print(f"TGI GPU: {image_uris.retrieve(framework='huggingface-llm', region=region, image_scope='inference', instance_type='ml.g5.2xlarge')}")
# TEI (embeddings) on GPU
print(f"TEI GPU: {image_uris.retrieve(framework='huggingface-tei', region=region, image_scope='inference', instance_type='ml.g5.2xlarge')}")
# TEI (embeddings) on CPU
print(f"TEI CPU: {image_uris.retrieve(framework='huggingface-tei-cpu', region=region, image_scope='inference', instance_type='ml.c6i.2xlarge')}")
# TGI (LLMs) on AWS Inferentia2 (Neuron)
print(f"TGI Neuron: {image_uris.retrieve(framework='huggingface-llm-neuronx', region=region, image_scope='inference', instance_type='ml.inf2.xlarge')}")

For PyTorch Training and PyTorch Inference DLCs, there is no such utility.