Supported frameworks and Amazon Regions - Amazon SageMaker AI
Services or capabilities described in Amazon Web Services documentation might vary by Region. To see the differences applicable to the China Regions, see Getting Started with Amazon Web Services in China (PDF).

Supported frameworks and Amazon Regions

Before using SageMaker smart sifting data loader, check if your framework of choice is supported, that the instance types are available in your Amazon account, and that your Amazon account is in one of the supported Amazon Regions.

Note

SageMaker smart sifting supports PyTorch model training with traditional data parallelism and distributed data parallelism, which makes model replicas in all GPU workers and uses the AllReduce operation. It doesn’t work with model parallelism techniques, including sharded data parallelism. Because SageMaker smart sifting works for data parallelism jobs, make sure that the model you train fits in each GPU memory.

Supported Frameworks

SageMaker smart sifting supports the following deep learning frameworks and is available through Amazon Deep Learning Containers.

Topics

PyTorch

Framework Framework version Deep Learning Container URI
PyTorch 2.1.0

763104351884.dkr.ecr.region.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-sagemaker

For more information about the pre-built containers, see SageMaker AI Framework Containers in the Amazon Deep Learning Containers GitHub repository.

Amazon Web Services Regions

The containers packaged with the SageMaker smart sifting library are available in the Amazon Web Services Regions where Amazon Deep Learning Containers are in service.

Instance types

You can use SageMaker smart sifting for any PyTorch training jobs on any instance types. We recommend that you use P4d, P4de, or P5 instances.