Distributed AI on AWS: Reference Architectures & Examples

This repository contains reference architectures and examples for distributed AI training and inference on Amazon SageMaker HyperPod, AWS ParallelCluster, AWS Parallel Computing Service (AWS PCS), and Amazon Elastic Kubernetes Service (Amazon EKS). The examples cover different model families and sizes, training frameworks and parallel optimizations (PyTorch DDP/FSDP, Megatron-LM, NeMo), and serving engines (vLLM, SGLang, NVIDIA Dynamo).

The major components of this repository are:

├── architectures/               # Cluster reference architectures (CloudFormation, Terraform)
├── ami/                         # Scripts to create Amazon Machine Images (Packer/Ansible)
├── examples/                    # Runnable training, inference, and use-case examples
├── validation/                  # Environment and cluster health validation tools
├── observability/               # Monitoring, metrics exporters, and profiling
└── micro-benchmarks/            # Micro-benchmarks (NCCL, NCCOM, NVSHMEM, etc.)

Workshops

You can follow the workshops below to train models on AWS. Each walks through several examples and shares practical guidance on operating a cluster for LLM training.

Blog

Posts about distributed AI on AWS are published at https://awslabs.github.io/awsome-distributed-ai/. The Hugo source lives on the content branch.

Blog content is editorially curated by AWS authors. Code samples in this repo (architectures/, examples/, etc.) accept external contributions as usual; see CONTRIBUTING.md.

Architectures

Each subdirectory under architectures/ is a deployable cluster architecture or a shared building block.

See docs/efa-cheatsheet.md for EFA tuning and the recommended environment variables.

Custom Amazon Machine Images

Custom machine images can be built using Packer for AWS ParallelCluster, Amazon EKS, and plain EC2. These images are based on Ansible roles and playbooks.

Examples

Examples live under examples/ and are organized along two axes:

• examples/training/ and examples/inference/ are framework-centric: the training or inference engine is the subject, and model variants underneath illustrate it (e.g. training/fsdp/, training/megatron-lm/, training/nemo/). Swapping the model gives "the same example with a different model."
• examples/use-cases/ is use-case-centric: a specific model or task is the subject and the framework is incidental (e.g. use-cases/detr-finetune/, use-cases/vjepa2/). Swapping the framework would still leave a recognizable demo.

Each example follows this general structure:

examples/
├── training/                   # framework-centric training/fine-tuning engines
│   └── <framework>/            # e.g. fsdp, deepspeed, megatron-lm, nemo, trl
│       └── <model>/            # e.g. llama3 (may be omitted for single-model cases)
│           ├── Dockerfile      # Container / environment setup
│           ├── README.md
│           ├── slurm/          # Slurm-specific launch scripts
│           └── kubernetes/     # Kubernetes manifests
├── inference/                  # framework-centric inference engines (vllm, …)
└── use-cases/                  # use-case-centric end-to-end demos
    └── <name>/                 # e.g. detr-finetune, esm2-hyperpod

The top-level directory for each example contains a general introduction and environment setup (Dockerfiles, training scripts, configs), while subdirectories provide service-specific launch instructions.

Browse examples/ to see the full list of frameworks, engines, and use cases.

Validation and Observability

Environment and cluster health validation tools live under validation/; monitoring stacks, metrics exporters, and profiling guides live under observability/.

Micro-benchmarks

Micro-benchmarks for evaluating network and communication performance are under micro-benchmarks/.

Contributors

Thanks to all the contributors for building, reviewing and testing.

Star History