Amazon Deep Learning Containers for PyTorch 2.7 ARM64 Training on EC2
Amazon Deep Learning Containers (DLCs)
This release includes a container image for Training on GPU, optimized for performance and scale on Amazon EC2. The image provides stable versions of NVIDIA CUDA, cuDNN, NCCL, and other components. All software components in this image are scanned for security vulnerabilities and updated or patched in accordance with Amazon Security best practices.
A list of available containers can be found in our documentation
Release Notes
Introduced containers for PyTorch 2.7 for Training on EC2. For details about this release, check out our GitHub release tag
. This image should be used with the G5g instance type
, which is powered by Graviton CPUs and NVIDIA T4G Tensor Core GPUs. Please refer to the official PyTorch 2.7.0 release notes here
. This image includes the following libraries:
CUDA 12.8.0
cuDNN 9.8.0.87
NCCL 2.26.2
EFA installer 1.40.0 (with Amazon OFI NCCL embedded)
Transformer Engine 2.0
Flash Attention 2.7.3
GDRCopy 2.5
Please note that EFA, Transformer Engine, Flash Attention, and GDRCopy have not been tested because of lack of hardware support.
The Dockerfile can be found here
.
For latest updates, please refer to the aws/deep-learning-containers GitHub repo.
Security Advisory
Amazon recommends that customers monitor critical security updates in the Amazon Security Bulletin
Python 3.12 Support
Python 3.12 is supported in the PyTorch ARM64 Training containers.
GPU Instance Type support
The containers support the Graviton GPU instance type G5g and contain the following software components for GPU support:
CUDA 12.8
cuDNN 9.8.0.87+cuda12.8
NCCL 2.26.2+cuda12.8
Amazon Regions support
The containers are available in the following regions:
Region |
Code |
|---|---|
US East (Ohio) |
us-east-2 |
US East (N. Virginia) |
us-east-1 |
US West (Oregon) |
us-west-2 |
US West (N. California) |
us-west-1 |
AF South (Cape Town) |
af-south-1 |
Asia Pacific (Hong Kong) |
ap-east-1 |
Asia Pacific (Hyderabad) |
ap-south-2 |
Asia Pacific (Mumbai) |
ap-south-1 |
Asia Pacific (Osaka) |
ap-northeast-3 |
Asia Pacific (Seoul) |
ap-northeast-2 |
Asia Pacific (Tokyo) |
ap-northeast-1 |
Asia Pacific (Melbourne) |
ap-southeast-4 |
Asia Pacific (Jakarta) |
ap-southeast-3 |
Asia Pacific (Sydney) |
ap-southeast-2 |
Asia Pacific (Singapore) |
ap-southeast-1 |
Asia Pacific (Malaysia) |
ap-southeast-5 |
Asia Pacific (Thailand) |
ap-southeast-7 |
Mexico (Central) |
mx-central-1 |
Canada (Central) |
ca-central-1 |
Canada (Calgary) |
ca-west-1 |
EU (Zurich) |
eu-central-2 |
EU (Frankfurt) |
eu-central-1 |
EU (Ireland) |
eu-west-1 |
EU (London) |
eu-west-2 |
EU( Paris) |
eu-west-3 |
EU (Spain) |
eu-south-2 |
EU (Milan) |
eu-south-1 |
EU (Stockholm) |
eu-north-1 |
Israel (Tel Aviv) |
il-central-1 |
Middle East (Bahrain) |
me-south-1 |
Middle East (UAE) |
me-central-1 |
SA (Sau Paulo) |
sa-east-1 |
China (Beijing) |
cn-north-1 |
China (Ningxia) |
cn-northwest-1 |
Build and Test
Built on: c6g.12xlarge
Tested on: c8g.4xlarge, t4g.2xlarge, r8g.2xlarge, m7g.4xlarge, g5g.16xlarge
Known Issues
-
There is no official Triton
distribution for ARM64/aarch64 yet, so some torch.compile workloads will fail with: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised: RuntimeError: Cannot find a working triton installation. More information on installing Triton can be found at https://github.com/openai/triton See GitHub issue
: Passing device_id to torch.distributed.init_process_group() results in NCCL randomly hanging during communications.