Argonne is pleased to be partnering with Microsoft DeepSpeed in the development of the Aurora model for our core science project through DeepSpeed4Science open sourced technologies. Science-driven large-language models will be developed at Argonne and trained on the Aurora Exascale system (https://www.anl.gov/aurora). The targeting models are in the trillion-parameter class, pretrained on a large corpus of general text and code as well as expertly curated large-scale (trillions of tokens) scientific domain datasets in microbiology, biochemistry, ecology, materials, chemistry, mathematics, physics and cosmology. DeepSpeed4Science technology is critical to reach the training performance requirements and scalability needed for these models to be feasibly trained. Argonne and partners recently formed the Trillion Parameter Consortium to bring together global-scale stakeholders interested in advancing large-scale scientific discovery. Microsoft’s DeepSpeed4Science team are founding members of the consortium.

Fusion energy represents a promising source of clean energy for the future, harnessing the power processes that fuel the sun and stars. However, plasma disruptions in tokamaks present significant challenges to the consistent and safe delivery of fusion energy. The INFUSE project, representing a partnership between scientists from Princeton University, DOE Argonne National Lab, DOE Brookhaven National Lab, and Microsoft, addresses this challenge by leveraging AI and advanced transformer models to enhance plasma disruption forecasting and control capabilities in magnetically-confined tokamak experiments. DeepSpeed4Science open source library plays a key role in the project by enabling efficient distributed training of these large transformer models, tailored for intricate time series forecasting tasks.

The DeepSpeed4Science team is excited to welcome AMD as a founding partner. We are looking forward to closely collaborating with AMD and their Research and Advanced Development (RAD) team to accelerate AI4Science models and workflows on world-class AMD hardware and software. AMD RAD is a leading research group with a stellar track record of partnering with the U.S. Department of Energy (DOE) National Laboratories and other organizations for more than a decade to drive critical advances in next-generation high-performance computing technologies and scientific breakthroughs. For example, AMD collaborated with the DOE National Laboratories on the FastForward, FastFoward-2, DesignForward, DesignForward-2, and PathForward programs to research and develop hardware and software technologies for exascale computing that are now found in some of the world’s fastest and most power-efficient supercomputers. 

AMD offers a wide variety of compute and network technology solutions with unmatched performance for critical workloads across AI/ML and HPC. The AMD ROCm™ open software platform is optimized for major HPC and machine learning frameworks, including PyTorch and DeepSpeed. According to the June 2023 Top500/Green500 list, systems with AMD EPYC™ Processors and/or AMD Instinct™ Accelerators power twelve of the twenty fastest supercomputers and thirteen of the twenty most energy-efficient supercomputers in the world. This includes ORNL’s Frontier supercomputer, the world’s fastest supercomputer on the 2023 Top500 list and the first to break the coveted exascale barrier, and CSC’s LUMI supercomputer, the fastest supercomputer in Europe. The Frontier and LUMI systems also take the top two spots in the latest HPL-MxP mixed-precision benchmark, which highlights the convergence of HPC and AI workloads. The Frontier supercomputer is being used to build a suite of open foundation large language models (LLMs) called FORGE for cutting edge scientific discovery. FORGE includes models with up to 26 billion parameters using 257 billion tokens from over 200 million scientific articles. Furthermore, the LUMI Supercomputer is being used to train a GPT-3 LLM with 13 billion parameters. The Allen Institute is collaborating with AMD to use LUMI to train a new, open LLM called Open Language Model (OLMo), for scientific discovery. This LLM is intended to benefit the research community by providing access and education for all aspects of model creation.

The DeepSpeed4Science team looks forward to closely collaborating with AMD to explore their broad portfolio of compute and network hardware including CPUs, GPUs, AI inference engines, FPGAs, Adaptive SoCs, and Smart NICs in order to deliver adaptable and accelerated systems that drive the AI4Science future. For example, high performance and energy efficiency may be achieved with ML-based surrogate models by running scientific simulations on AMD CPUs and GPUs, training/retraining on AMD GPUs, performing inference on AI Engines, and using FPGAs for data processing. Efficient scale-out performance can be achieved by offloading common distributed deep learning primitives on AMD Smart NICs.

The AMD RAD team has been actively developing and investigating MLIR-based compiler passes, smart runtimes for heterogeneous architectures, fast communication primitives, and optimized neural network libraries for AMD’s heterogeneous hardware portfolio. These techniques and technologies can help optimize foundational AI4Science models on systems with AMD hardware. AMD RAD has also released low-overhead, open-source performance monitoring and analysis tools that will be critical to design, optimize, and deploy AI4Science workloads. We look forward to further exploring these research topics by closely collaborating with AMD researchers and engineers to build new foundational AI4Science models and co-design deep learning system architecture technologies. Together, we plan to advance scientific knowledge and accelerate breakthroughs across various domains, such as material sciences, computational chemistry, health sciences, manufacturing, molecular dynamics, high-energy physics, cosmology, and more.