Cerebras: Wafer Scale Engine Outperforms Frontier Supercomputer on Molecular Dynamics Simulations

Print Friendly, PDF & Email

SUNNYVALE, Calif. – May 15, 2024 – Accelerated generative AI chip company Cerebras Systems, in collaboration with researchers from Sandia, Lawrence Livermore, and Los Alamos National Laboratories, said it has acheved a breakthrough in molecular dynamics (MD) simulations using the second generation Cerebras Wafer Scale Engine (WSE-2). Researchers performed atomic scale simulations at the millisecond scale – 179x faster than on the world’s leading supercomputer, Frontier, which is built with 39,000 GPUs.

“Existing supercomputers have been limited to simulating materials at the atomic scale at a rate of 2-3 simulated microseconds per month, restricting our understanding of how materials evolve and behave over longer periods,” Cerebras said. “This breakthrough achieved by Cerebras and its collaborators at national laboratories has shattered this barrier. By harnessing the power of the Cerebras WSE-2, the processor at the heart of the Cerebras CS-2 system, researchers can now simulate materials for milliseconds – an astounding leap that opens up entirely new vistas in materials science.”

“The partnership between the NNSA laboratories and Cerebras Systems is part of the Advanced Memory Technology (AMT) program, which aims to accelerate exascale supercomputers by 40x as early as 2025. With Cerebras’ currently deployed wafer-scale computers, the teams achieved this materials science breakthrough and a speedup that exceeded the goal of the AMT program by more than 4X,” said James H. Laros III, Distinguished Member of Technical Staff and AMT program lead. “This experience bodes well for future impacts to our program and potential scientific advances.”

The research team innovated across both hardware and software to overcome the limitations of today’s supercomputers. By mapping individual atoms onto the WSE-2’s nearly one million cores and enabling efficient communication between neighboring cores, the system simulated 270,000 timesteps per second across 800,000 atoms – a staggering 179-fold speedup over Frontier, the world’s leading supercomputer. This breakthrough allows researchers to gain unprecedented insights into the long-term behavior and future evolution of materials at the atomic scale.

“This work changes the landscape of what is possible with molecular dynamics simulations,” said Michael James, Chief Architect of Advanced Technologies and co-founder of Cerebras Systems. “Simulations that would have taken a year on a traditional supercomputer can now be completed in just two days. Scientists will now be able to explore previously inaccessible phenomena across a wide range of domains.”

Long timescale simulations will allow scientists to explore previously inaccessible phenomena across a wide range of domains:

  • Materials scientists can now study the long-term behavior of complex materials, such as the evolution of grain boundaries in metals, leading to the development of stronger, more resilient materials.
  • Pharmaceutical researchers can simulate protein folding and drug-target interactions over physiologically relevant timescales, accelerating the discovery of life-saving therapies.
  • Renewable energy experts can optimize catalytic reactions and design more efficient energy storage systems by simulating atomic-scale processes over extended durations.

To achieve this remarkable feat, the Cerebras team employed a novel mapping scheme that assigns each atom to a single core on the WSE-2. The cores are organized in a 2D grid, with the physical simulation domain mapped onto this grid to preserve locality. Cores communicate with their neighbors to exchange atom information, allowing for efficient parallel processing. The WSE-2’s unique architecture, with its high memory bandwidth and low-latency communication fabric, enables this fine-grained parallelism to be exploited effectively.

“The NNSA’s Advanced Memory Technologies program started 1.5 years ago with the goal of 40x performance improvement on critical NNSA applications over what can be achieved on exascale systems. We all had our doubts on achieving this goal within the short timeframe, but Cerebras’s technology and team has helped us exceed this goal by demonstrating unprecedented 179x performance improvement on MD simulations,” Siva Rajamanickam, Sandia National Laboratories, Principal Member of Technical Staff. “These results open up new opportunities for materials research and science discoveries beyond what we envisioned. We are excited to continue this collaboration with Cerebras and explore new frontiers in science.”

As Cerebras continues to push the boundaries of high-performance computing with its wafer-scale technology, even more groundbreaking advancements are on the horizon.

“I have been working in atomistic simulation of materials for more than 20 years. During that time, I have participated in massive improvements in both the size and accuracy of the simulations.  However, despite all this, we have been unable to increase the actual simulation rate.  The wall-clock time required to run simulations has barely budged in the last 15 years,” Aidan Thompson, Sandia National Laboratories, Distinguished Member of Technical Staff. “With the Cerebras Wafer-Scale Engine, we can all of a sudden drive at hypersonic speeds.  This joint DOE-Cerebras team has accomplished something unprecedented and highly disruptive. I look forward to seeing how this transforms all kinds of scientific research in the near future.”

For more information, see https://arxiv.org/abs/2405.07898.

Speak Your Mind