AstraZeneca improved life sciences model training time with AMD Instinct MI300X GPUs

In collaboration with AstraZeneca, the AMD Silo AI team successfully ported three life sciences models to run on AMD Instinct™ MI300X GPUs. Prior to this effort, the workloads used AstraZeneca’s environment based on NVIDIA GPUs, which were replaced by the higher memory bandwidth and newer generation MI300X GPU.

“Working with the AMD Silo AI team has been easy and efficient. Together, we have worked as one team to push the boundaries of model performance” says Dr. Per Hillertz, One IT Site Lead, AstraZeneca. 

The AMD Silo AI team helped AstraZeneca switch to an AMD Instinct™ MI300X node on the TensorWave by leveraging the AMD ROCm™ open software stack. Out-of-the-box performance was further improved through additional optimization.

Models Used and Key Improvements 

The project focused on three models: REINVENT4, SemlaFlow, and SwinUNETR. With additional optimization the following results were achieved:

• SemlaFlow is a graph neural network model for generating 3D molecular structures, widely used in drug discovery. With optimizations from the AMD Silo AI team, training time was cut by 49%¹
• REINVENT4 is a de novo design model for creating new molecules. Using its transfer learning capabilities, we trained models to generate molecules similar to specific inputs. AMD Silo AI optimizations improved speed by an average of 41% across four (4) configurations, compared to the same models without optimizations².
• SwinUNETR is a deep learning model for 3D medical image segmentation, such as detecting tumors in CT or MRI scans. By tuning data loading algorithms and using the latest PyTorch attention features, we reduced training time by up to 1.8x in a baseline setup³.

Impact and Future Collaboration

All three models— SemlaFlow, REINVENT4, and SwinUNETR—ran seamlessly on AMD Instinct™ MI300X with no code changes or extra engineering effort. The AMD ROCm™ Open software stack was a drop-in replacement for NVIDIA CUDA-based recipes, where we were able to quickly replace CUDA packages with ROCm™ software parallels, making it a straightforward and painless experience. The PyTorch-based AI code, previously run on the NVIDIA Tesla V100 GPU, didn’t require any changes to run on AMD Instinct™ MI300X and ran out-of-the-box, making it an effortless and pleasing experience. AstraZeneca’s researchers appreciated the AMD software stack’s intuitiveness and ease of use. “We were able to achieve these results without any significant changes to our software. Our results showed substantial improvements in the speed of training,” says Dr. Hannes Loeffler, Machine Learning Researcher, AstraZeneca. The smooth transition to MI300X was further enhanced by the performance gains from AMD Silo AI optimization.

Building on the success of our initial collaboration, AstraZeneca and AMD Silo AI are now exploring deeper co-creation opportunities. Together, we aim to accelerate AstraZeneca's scientific and product AI-driven drug development workflows improving speed, streamlining processes, and enhancing accuracy.

Open ecosystems in AI facilitate interoperability across the AI stack, from hardware to applications. Deploying an open approach enables investment flexibility, speed of innovation and adoption of new technologies, while avoiding vendor lock-in and allowing for the broader community to drive the direction of technology. The tools and approaches applied here are not exclusive to life sciences and pharmaceuticals but are adaptable and scalable across other industries and use cases, such as automotive, robotics, and weather forecasting.

AstraZeneca: AstraZeneca (LSE/STO/Nasdaq: AZN) is a global, science-led biopharmaceutical company that focuses on the discovery, development, and commercialisation of prescription medicines in Oncology, Rare Diseases, and BioPharmaceuticals, including Cardiovascular, Renal & Metabolism, and Respiratory & Immunology. Based in Cambridge, UK, AstraZeneca’s innovative medicines are sold in more than 125 countries and used by millions of patients worldwide. Please visit astrazeneca.com and follow the Company on social media @AstraZeneca.

AMD Silo AI:  Pushes the Frontier of AI on AMD Compute Platforms.  AMD Silo AI is a leading AI lab that helps customers develop and deploy advanced AI models and solutions optimized for leadership compute platforms.  As a global AI center of excellence with over 300 AI scientists and 125+ PhDs, they combine deep scientific knowledge with practical technology understanding to help organizations integrate, deploy and scale AI effectively. Our work spans from groundbreaking research to enterprise-ready AI solutions.