How Amazon Search leverages PyTorch and its eco systems to build and deploy LLM into production

### Overview Presented at the PyTorch Mini Summit within the Open Source Summit (OSS), this talk highlights how AWS and the Amazon Search team utilize PyTorch and AWS infrastructure to build, train, and deploy large language models (LLMs) into production. Shubhakar Matkone (Principal GTM Specialist at AWS) begins by detailing AWS's strategic collaboration with Meta, the PyTorch Foundation, and its managed ML ecosystem. Aji Joseph (Principal Engineer at Amazon Search) then explains the architecture, operational lifecycle, and optimization techniques behind Amazon's "M5" model framework that powers search and shopping experiences across amazon.com.

### Key Topics - **AWS & PyTorch Ecosystem**: Strategic partnership with Meta, contributions to TorchServe, PyTorch Foundation governance, integration with Amazon SageMaker, SageMaker JumpStart, and custom silicon (AWS Trainium and Inferentia). - **The M5 Model Lifecycle**: Amazon Search’s M5 initiative, encompassing multimodal (images, unstructured text, tables, videos), multilingual, multi-locale, multi-entity, and multitask representations. The lifecycle involves pre-training a foundation model, multitask pre-training, model distillation/quantization, serialization, and distribution to internal partner teams for task-specific fine-tuning. - **Training Velocity & Reliability**: Driving ML innovation through hypothesis-driven development, running up to 10,000 experiment jobs monthly, achieving 100% reproducibility, automated recovery from hardware failures during multi-week training, and streaming data across AWS regions via a C++-based S3 CRT data loader. - **Decoupled Inference & Optimization**: Transitioning from PyTorch eager mode to serialized `torch.script` modules to decouple research code from production deployment, optimizing through pruning, quantization (FP16, Bfloat16, INT8), and compilers/SDKs like AWS Neuron and TensorRT.

### Technical Details - **Compute and Silicon**: Training utilizes mixed hardware environments, including AWS Trainium, NVIDIA GPUs (P4de/P4dn A100 instances, P3dn V100 instances), and Intel dl1 accelerators. Inferentia2 and GPUs run production inference workloads. - **Training Frameworks**: Built on DeepSpeed utilizing ZeRO optimizations alongside PyTorch Distributed. The team adopted Bfloat16 to stabilize large model training while converging encoder models up to 100 billion parameters. - **Data Ingestion**: To address regional compute availability constraints, the team streams data from petabyte-scale storage (such as 300 TB datasets) directly into CPU memory on GPU instances for on-the-fly pre-processing (e.g., resizing, color augmentation) using an optimized C++ AWS S3 client integrated with PyTorch DataLoader, achieving at least 10x performance over the standard S3 boto client. - **Inference Metrics & Pipeline**: Model serving achieves sub-10ms latency for a 1-billion-parameter encoder model running on AWS Inferentia2 and GPU instances. TorchScript handles core graph serialization, while auxiliary pre- and post-processing code handles branching constraints for non-batch real-time text and batch multimodal inference.

### Takeaways Amazon Search demonstrates that deploying LLMs at scale requires decoupling model development from serving via robust serialization (`torch.script`), investing heavily in experiment reproducibility, and architecting data loaders to follow geographically dispersed compute. Moving forward, the team is targeting PyTorch 2.0 compiler integrations, FP8 compute formats, and sparse computation techniques to further drive down training and inference costs across industry partnerships.