Title

AWS re:Invent 2022 - AI parallelism: How Amazon Search scales deep-learning training (CMP209)

Summary

James Park, a solutions architect at AWS, and RJ, an engineer with Amazon Search, discuss how Amazon scales deep-learning training for search.
They cover the growth of deep learning, particularly in NLP, and the challenges of training increasingly large models with limited hardware scalability.
Distributed training strategies like data parallelism, pipeline parallelism, and tensor parallelism are essential for handling large models.
AWS offers a range of services and infrastructure options to support different customer needs, including SageMaker for end-to-end ML workflows and EC2 instances for framework-level work.
Amazon Search's M5 team focuses on training large language models considering multi-modality, multi-locality, multilingual aspects, multi-entity relationships, and multi-task learning.
The M5 team uses PyTorch and DeepSpeed for training, with AWS hardware like P3DN, P4DN, and G4DN instances.
They emphasize the importance of reproducibility, reliability, and debuggability in their ML workflows.
Model vending and inference optimization are crucial for deploying models to production.
The M5 team has successfully trained and converged 100 billion parameter models, running 10K experiments per month, and achieved sub-10-millisecond latency for 1 billion parameter encoder models on both GPU and Inferentia hardware.

The exponential growth in model size requires innovative distributed training strategies to manage the computational load effectively.
AWS's comprehensive suite of services and infrastructure caters to a wide range of ML needs, from high-level APIs to low-level framework and infrastructure management.
The M5 team's approach to handling large language models is holistic, considering not just the technical aspects of training but also the practical implications for deployment across different locales and languages.
Reproducibility, reliability, and debuggability are critical for maintaining high experiment velocity and ensuring that ML models can be trained and iterated upon efficiently.
Collaboration with AWS AI, AWS Deep Engine Sciences, AWS Batch, NVIDIA, MetaPyTorch, and Amazon FSx teams has been instrumental in the M5 team's success, highlighting the importance of cross-team collaboration in tackling complex ML challenges.
The M5 team's ability to train large models quickly and cost-effectively, with high experiment throughput and low-latency inference, demonstrates the potential for large-scale AI applications in real-world scenarios, such as Amazon's search functionality.