This repository provides a comprehensive collection of reference architectures, best practices, and test cases for distributed machine learning model training on AWS. It targets ML engineers and researchers seeking to optimize large-scale training on services like Amazon SageMaker Hyperpod, AWS ParallelCluster, AWS Batch, and Amazon EKS, offering practical examples and validation tools.

How It Works

The project structures distributed training solutions into distinct categories: reference architectures (CloudFormation templates for S3, VPC, and compute clusters), custom AMIs and containers, and detailed test cases. Test cases are organized by framework (PyTorch, JAX) and parallelization strategy (DDP, FSDP, Megatron-LM, NeMo, Trainium), including scripts, configurations, and validation utilities for performance monitoring and troubleshooting.

Quick Start & Requirements

• Installation: Primarily involves deploying CloudFormation templates for infrastructure setup.
• Prerequisites: AWS account, CloudFormation, Docker, Python, and potentially specific ML frameworks and libraries depending on the test case. EFA (Elastic Fabric Adapter) is recommended for high-performance networking.
• Resources: Requires AWS infrastructure provisioning (VPC, S3, compute instances).
• Documentation: Workshops are available for SageMaker HyperPod and AWS ParallelCluster.

Highlighted Details

• Supports multiple AWS compute services: SageMaker HyperPod, ParallelCluster, Batch, EKS.
• Extensive test cases cover various PyTorch distributed training strategies (FSDP, Megatron-LM, NeMo) and AWS Trainium.
• Includes validation scripts and tools for performance monitoring (e.g., EFA Prometheus exporter, nsight profiling).
• CI integration with pytest is provided for automated testing.

Maintenance & Community

The repository is part of AWS Samples, indicating official backing. It acknowledges contributions from various individuals.

Licensing & Compatibility

The repository is licensed under the Apache-2.0 License, permitting commercial use and integration with closed-source projects.

Limitations & Caveats

The repository is undergoing a major refactoring, with test cases being a particular focus. Users preferring the previous structure should refer to v1.1.0. Architectures are designed to work with specific S3 buckets and VPCs created via provided templates.