As the pandemic swept across the world, virtually every research supercomputer lit up to support Covid-19 investigations. But even as the world transformed, the fairly stable status quo of simulation-based scientific computing was itself beginning to more rapidly change with the burgeoning field of AI and AI-specific accelerators. The Argonne Leadership Computing Facility (ALCF) at Argonne National Laboratory is, perhaps, the exemplar of these dual shifts, hosting both a wide range of Covid-19 research and the ALCF AI Testbed, equipped with a cornucopia of novel AI accelerators. At the AI Systems Summit this week, Venkatram Vishwanath—data science team lead at ALCF—spoke to the intersection of advanced AI acceleration and Covid-19 research at the lab.

“We recently deployed the ThetaGPU and Polaris systems, which are HPE-Nvidia systems where we’re using A100 GPUs, and we are on our path to deploy Aurora, which is an Intel-HPE system and will have the Intel GPUs,” Vishwanath said. “As we march along these [systems], what we’ve also noticed is the growth of AI accelerators and the promise that they have shown for enterprise workloads.”

That promise, he said, made ALCF interested in how the accelerators could be used to tackle the “surge” in scientific machine learning, with AI being used to replace parts of simulations—or even entire simulations. “Our key goal is to really understand: how do you deploy and design infrastructure to facilitate these novel ways of how HPC and AI are being integrated for science?”

Argonne’s AI Testbed

To that end, ALCF began a pathfinding effort that has, so far, been reified through the AI Testbed. The testbed is currently home to five core systems: a CS-2 system from Cerebras; an MK1 system from Graphcore; a Groq system; a Gaudi system from Habana Labs; and a Dataflow system from SambaNova.

“As part of the testbed, we have a wide gamut of science applications that we have witnessed to run effectively on these systems,” Vishwanath continued, “that spans from understanding the drug response for cancer, various imaging science experiments, inference operations in tokamak reactor operations, protein-protein folding, among others.”

“As part of Covid-19, we are really interested in understanding how the SARS-CoV-2 spike protein interacts with the host—and this is a really complex, multiscale problem.”

Vishwanath went on to outline Argonne’s traditional computing workflow for these sorts of science problems. “We would run a large ensemble of jobs, which have different configurations,” he said. “We would then write this out to storage, analyze the data and visualize the results—and then try to understand, perhaps there were some configurations that did not produce results that we wanted, and we would stop these, or we would find new configurations that we should have explored and kickstart new simulations.”

“The key goal is to see: can we do this more intelligently?”

Working smarter (and harder)

To that end, ALCF researchers worked in early 2021 to see if they could combine simulations with online learning to steer simulations towards the configurations they needed to explore.

“We are running, in this case, the OpenMM simulations on the ThetaGPU nodes,” Vishwanath said. “We have this coupled with the [Cerebras] CS-1 system that is on the ALCF AI testbed. We stream the data and the trajectories from Theta-GPU to the CS1. We run online training on the CS-1, where we are trying to use the trajectories to learn this embedding space. We then run our inference on one of the ThetaGPU nodes to identify outliers, which helps us either stop existing simulations or seed new ones as part of this.”

Operating this workflow on a 131,000 system that Vishwanath described as “very critical for Covid-19 research,” he said that “we observe a close to 3.7-fold improvement … where we obtain close to 83,000 samples per second on a CS-1 compared to 2,200 samples per second on our DGX.” (ThetaGPU consists of 24 Nvidia DGX A100 nodes.)

“The other thing to note is that by using the CS-1 accelerator in this case, you are able to completely overlap the training with the compute runs,” he continued. “So the overhead in terms of the overall workflow is mitigated.” This translated to sampling a SARS-CoV-2 spike protein conformation similar to a native folded state (“very key”) in 18 hours of wall clock time. “If we were to run an equilibrium simulation on the A100s … it would have taken us approximately 923 hours,” Vishwanath said. “We were able to explore this space and achieve a 50× improvement over an equilibrium ensemble simulation.”

Taking AI acceleration to “new dimensions”

That specific work was completed in 2021—but Vishwanath said that since then, ALCF has “taken it to completely new dimensions.” Specifically, the researchers began looking at how to integrate SARS-CoV-2 research data from cryogenic electron microscopy (cryo-EM) experiments, all-atom molecular dynamics (MD) simulations and fluctuating finite-element analysis (FFEM).

“How do you integrate cryo-EM, FFEA and the all-atom MD simulations to really give you the long timescales that are needed as part of this research?” Vishwanath asked.

The ALCF research team designed a workflow where the cryo-EM data was used as an initial best guess, which was then fed to the FFEA process, which helped identify key variables using hierarchical AI. This was then used to steer all-atom MD simulations, which themselves could be used to help researchers design subsequent cryo-EM experiments.

“In comparison to the previous work, we have, now, three different types of AI that are coupled in,” Vishwanath said. “Some for surrogate modeling, some for intelligent steering of the simulation and for inferencing, among others. And this workflow is also geographically distributed, and it spans components that can run on the CPU, that can run on the GPU … [and] the accelerators that I’ve been describing here.”

This workflow—which enabled research that was nominated for the Gordon Bell Special Prize for High Performance Computing-Based Covid-19 Research in 2021—spanned the Perlmutter system (National Energy Research Scientific Computing Center), the Frontera system (Texas Advanced Computing Center), the Summit system (Oak Ridge Leadership Computing Facility), both Theta and ThetaGPU (ALCF) and the AI Testbed.

Vishwanath said that this geographically and technologically distributed workflow was managed using Balsam. “We were able to map a really diverse set of workflows, such as NAMD 3.0 [which] runs effectively on the GPUs to leverage the resources, FFEA that ran on the CPUs, and the deep learning training and inference jobs that ran on the accelerators.”

Once again, he said, Cerebras hardware—specifically, the Wafer-Scale Engine that powers its systems—proved its mettle in the process. (In the interim, Argonne had upgraded its CS-1 system to a CS-2.)

“What we observe is that out of the box, we get about 100× improvement on the Wafer-Scale Engine over a single V100 GPU,” he said. “We expect this to improve as we start optimizing this model on the Wafer-Scale Engine.”

Vishwanath closed by reaffirming how seriously Argonne takes its experiments with novel AI accelerators. “I’ll say that AI accelerators are really promising for scientific machine learning and we’ve observed significant speedups for a wide gamut of science,” he said, adding that Argonne expected such accelerators to play a “key role in future system design.”

To learn more about this research, read additional coverage from HPCwire here.