Lawrence Livermore National Laboratory (LLNL) and artificial intelligence computer company Cerebras Systems have integrated the world’s largest computer chip into the National Nuclear Security Administration’s (NNSA’s) Lassen system, upgrading the top-tier supercomputer with cutting-edge AI technology.
Technicians recently completed connecting the Silicon Valley-based company’s massive, 1.2 trillion transistor Wafer-Scale Engine (WSE) chip – designed specifically for machine learning and AI applications – to the 23-petaflop Lassen, along with the CS-1 accelerator hardware system that houses the chip. Lassen is the unclassified companion to the IBM/NVIDIA system Sierra and is currently No. 14 on the latest Top 500 List of the world’s most powerful supercomputers.
The successful pairing of Cerebras’ machine learning system with Lassen’s world-class simulation capability makes LLNL the first institution to integrate the AI platform with a large-scale supercomputer and creates a radically new type of computing solution, enabling researchers to investigate novel approaches to predictive modeling. Users gained access to the system in July and have begun work on initial AI models.
Funded by the NNSA’s Advanced Simulation and Computing (ASC) program, the platform aims to accelerate solutions for crucial Department of Energy and NNSA national security mission applications by orders of magnitude over the coming decade, researchers said. Fusion implosion experiments performed at the National Ignition Facility, materials science and rapid design of new prescription drugs for COVID-19 and cancer (through the Accelerating Therapeutic Opportunities in Medicine, or ATOM project) are all early applications.
“The addition of an AI component to a world-class supercomputer like Lassen will allow our scientists and physicists to leverage the large investments in AI being made by industry to improve both the accuracy of our simulations as well as the productivity of the end user in solving mission critical problems,” said Thuc Hoang, acting director of the ASC program. “We are excited to be collaborating with Cerebras in putting this pioneering technology to work for NNSA.”
Billed as the “world’s fastest AI supercomputer,” the CS-1 system runs on the WSE chip, which consists of 400,000 AI-optimized cores, 18 gigabytes of on-chip memory and 100 petabits per second of on-chip network bandwidth. The upgrade marks the first time LLNL has a high-performance computing (HPC) resource that has included AI-specific hardware and effectively creates the world’s first computer system designed for “cognitive simulation” (CogSim) – a term used by LLNL scientists to describe the combination of traditional HPC simulations with AI.
“While Moore’s Law is not yet dead, we can see that it’s slowing down. The historic cadence of hardware improvements has slowed or even ceased, but the demand for computing has not,” said Chief Technology Officer for Livermore Computing Bronis R. de Supinski, who led the CS-1 procurement effort. “We need new answers for how to improve our ability to meet our mission requirements and to respond to ever-increasing computational requirements. Cognitive simulation is an approach that we believe could drive continued exponential capability improvements, and a system-level heterogeneous approach based on novel architectures such as the Cerebras CS-1 are an important part of achieving those improvements.”
Combining the CS-1 with Lassen allows LLNL to explore heterogeneity, wherein the supercomputer includes computer elements with different specializations that each contribute to working on a given job, de Supinski said. This would allow operations such as data generation and error correction to run concurrently, resulting in a much more integrated, efficient and cost-effective solution to scientific problems, according to Lab researchers.
“Having a heterogeneous system is motivating us to identify where in our applications multiple pieces can be executed simultaneously,” said LLNL computer scientist Ian Karlin. “For many of our cognitive simulation workloads, we will run the machine learning parts on the Cerebras hardware, while the HPC simulation piece continues on the GPUs, improving time-to-solution.”
The Cerebras WSE processor is optimized for deep learning. It is more than 56 times larger than the largest graphics processing unit (GPU), and contains 78 times more compute cores. With 18 gigabytes of on-chip memory, the WSE has 3,000 times more on-chip memory and more than 10,000 times more memory bandwidth than GPUs. The WSE chip is packaged in the CS-1 system, which includes water-cooling and more than 1.2 Terabits per second of I/O bandwidth. The platform is linked to Lassen through the supercomputer’s existing Mellanox InfiniBand communication network.
“Cerebras is proud to be partnering with LLNL to push the boundaries of supercomputing. It is every startup’s dream to help take their industry into new and uncharted territory,” said Andrew Feldman, founder and CEO of Cerebras Systems. “The CS-1 is the fastest AI computer in existence, and by integrating into Lassen, one of the world’s fastest supercomputers, a new and unique class of solutions emerge. This integrated AI + supercomputer powerhouse combines the best of supercomputing horsepower with the best in AI-specific computation to create a workflow aimed at solving the most challenging problems in science and technology.”
Lab computer scientists said the system’s capabilities will enable them to skip unnecessary processing in workflows and accelerate deep learning neural networks. Wafer-scale integration will minimize the limits of communication for scaling up the training of neural networks and how often researchers need to “slice and dice” problems into smaller jobs, said LLNL computer scientist Brian Van Essen, who heads the Lab’s efforts in large-scale neural networks.
“Either we can do scientific exploration in a shorter amount of compute time, or we could go more in-depth in the areas where the science is less certain, using more compute time but getting a better answer,” Van Essen said. “Having this system be monitored by the cognitive framework can allow us to improve the efficiency of our physicists. Rather than babysitting the simulations to look for errors, they can hopefully start to look more into the scientific questions that are being illuminated by this type of system.”
A research team led by Van Essen has selected two AI models from a Laboratory Director’s Initiative on Cognitive Simulation to run on the CS-1 system. LLNL researchers said that their preliminary work focuses on learning from up to five billion simulated laser implosion images for optimization of fusion targets used in experiments at the National Ignition Facility (NIF), with the goal of reaching high energy output and robust fusion implosions for stockpile stewardship applications.
LLNL physicist and Cognitive Simulation lead Brian Spears said NIF scientists will use the CS-1 to embed dramatically faster AI-based physics models directly inside Lassen’s simulation codes, resulting in significantly more detailed physics in less time. Additionally, the expanded compute load will allow scientists to merge enormous simulation volume with hyper-precision experimental data, Spears said.
“Pairing the AI power of the CS-1 with the precision simulation of Lassen creates a CogSim computer that kicks open new doors for inertial confinement fusion (ICF) experiments at the National Ignition Facility,” Spears said “Now, we can combine billions of simulated images with NIF’s amazing X-ray and neutron camera output to build improved predictions of future fusion designs. And, we’re pretty sure we’re just scratching the surface – giving ourselves a sneak peek at what future AI-accelerated HPC platforms can deliver and what we will demand of them.”
To ensure a successful integration, LLNL and Cerebras are engaging in an Artificial Intelligence Center of Excellence (AICoE) over the coming years to determine the optimal parameters that will allow cognitive simulation to work with the Lab’s workloads. Depending on the results, LLNL could add more CS-1 systems, both to Lassen and other supercomputing platforms.
“I’m super excited,” said Van Essen, who is leading the AICoE effort. “This is why we’re here, to do this kind of science. I’m a computer architect by training, so the opportunity to build this kind of system and be the first one to deploy these things at this scale is absolutely invigorating. Integrating and coupling it into a system like Lassen gives us a really unique opportunity to be the first to explore this kind of framework.”
