Skip to main content
OSTI.GOVU.S. Department of Energy Office of Scientific and Technical Information
U.S. Department of Energy
Office of Scientific and Technical Information
 

Acceleration of Graph Neural Network-Based Prediction Models in Chemistry via Co-Design Optimization on Intelligence Processing Units

Journal Article · · Journal of Chemical Information and Modeling
 [1];  [2];  [2];  [2];  [3];  [1];  [1];  [1];  [1];  [2];  [4];  [5];  [2]
  1. Graphcore, Cambridge (United Kingdom)
  2. Pacific Northwest National Laboratory (PNNL), Richland, WA (United States)
  3. Univ. of Washington, Seattle, WA (United States)
  4. IBM Research, Yorktown Heights, NY (United States)
  5. Pacific Northwest National Laboratory (PNNL), Richland, WA (United States); Univ. of Washington, Seattle, WA (United States)
+ Show Author Affiliations

Atomic structure prediction and associated property calculations are the bedrock of chemical physics. Since high-fidelity ab initio modeling techniques for computing the structure and properties can be prohibitively expensive, this motivates the development of machine-learning (ML) models that make these predictions more efficiently. Training graph neural networks over large atomistic databases introduces unique computational challenges such as the need to process millions of small graphs with variable size and support communication patterns that are distinct from learning over large graphs such as social networks. We demonstrate a novel hardware-software co-design approach to scale up the training of atomistic graph neural networks (GNN) for structure and property prediction. First, to eliminate redundant computation and memory associated with alternative padding techniques and to improve throughput via minimizing communication, we formulate the effective coalescing of the batches of variable-size atomistic graphs as the bin packing problem and introduce a hardware-agnostic algorithm to pack these batches. In addition, we propose hardware-specific optimizations including a planner and vectorization for the gather-scatter operations targeted for Graphcore’s Intelligence Processing Unit (IPU), as well as model-specific optimizations such as merged communication collectives and optimized softplus. Putting these all together, we demonstrate the effectiveness of the proposed co-design approach by providing an implementation of a well-established atomistic GNN on the Graphcore IPUs. We evaluate the training performance on multiple atomistic graph databases with varying degrees of graph counts, sizes and sparsity. Here, we demonstrate that such a co-design approach can reduce the training time of atomistic GNNs and can improve the performance by up to 1.5× compared to the baseline implementation of the model on the IPUs. Additionally, we compare our IPU implementation with a Nvidia GPU-based implementation and show that our atomistic GNN implementation on the IPUs can run 1.8× faster on average compared to the execution time on the GPUs.

Research Organization:
Pacific Northwest National Laboratory (PNNL), Richland, WA (United States)
Sponsoring Organization:
USDOE Office of Science (SC), Basic Energy Sciences (BES). Chemical Sciences, Geosciences & Biosciences Division (CSGB); USDOE National Nuclear Security Administration (NNSA); USDOE Office of Science (SC), Advanced Scientific Computing Research (ASCR)
Grant/Contract Number:
AC05-76RL01830
OSTI ID:
2323423
Report Number(s):
PNNL-SA--193670
Journal Information:
Journal of Chemical Information and Modeling, Journal Name: Journal of Chemical Information and Modeling Journal Issue: 5 Vol. 64; ISSN 1549-9596
Publisher:
American Chemical SocietyCopyright Statement
Country of Publication:
United States
Language:
English

References (56)

A fast non-negativity-constrained least squares algorithm journal September 1997
First Principles Neural Network Potentials for Reactive Simulations of Large Molecular and Condensed Systems journal August 2017
Multilevelk-way Partitioning Scheme for Irregular Graphs journal January 1998
Combinatorial Optimization: Theory and Algorithms book January 2012
A simple proof of the inequality MFFD(L)≤71/60 OPT(L) + 1,L for the MFFD bin-packing algorithm journal July 1995
A 7160 theorem for bin packing journal October 1985
Predicting Drug–Target Interaction Using a Novel Graph Neural Network with 3D Structure-Embedded Graph Representation journal August 2019
SchNetPack: A Deep Learning Toolbox For Atomistic Systems journal November 2018
Machine Learning Prediction of Nine Molecular Properties Based on the SMILES Representation of the QM9 Quantum-Chemistry Dataset journal November 2020
3D-Scaffold: A Deep Learning Framework to Generate 3D Coordinates of Drug-like Molecules with Desired Scaffolds journal October 2021
The Rise of Neural Networks for Materials and Chemical Dynamics journal July 2021
Open Catalyst 2020 (OC20) Dataset and Community Challenges journal May 2021
Molecular Transformer: A Model for Uncertainty-Calibrated Chemical Reaction Prediction journal August 2019
Development of a “First Principles” Water Potential with Flexible Monomers: Dimer Potential Energy Surface, VRT Spectrum, and Second Virial Coefficient journal November 2013
Development of a “First Principles” Water Potential with Flexible Monomers. II: Trimer Potential Energy Surface, Third Virial Coefficient, and Small Clusters journal March 2014
The Role of Machine Learning in the Understanding and Design of Materials journal November 2020
E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials journal May 2022
Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements journal May 2022
polyBERT: a chemical language model to enable fully machine-driven ultrafast polymer informatics journal July 2023
Recent advances and applications of machine learning in solid-state materials science journal August 2019
Accelerated antimicrobial discovery via deep generative models and molecular dynamics simulations journal March 2021
Enabling deeper learning on big data for materials informatics applications journal February 2021
On scientific understanding with artificial intelligence journal October 2022
Molecular contrastive learning of representations via graph neural networks journal March 2022
CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling journal September 2023
A universal graph deep learning interatomic potential for the periodic table journal November 2022
Quantum chemistry structures and properties of 134 kilo molecules journal August 2014
ANI-1: an extensible neural network potential with DFT accuracy at force field computational cost journal January 2017
Hierarchical modeling of molecular energies using a deep neural network journal June 2018
SchNet – A deep learning architecture for molecules and materials journal June 2018
Atlas of putative minima and low-lying energy networks of water clusters n = 3–25 journal December 2019
A look inside the black box: Using graph-theoretical descriptors to interpret a Continuous-Filter Convolutional Neural Network (CF-CNN) trained on the global and local minimum energy structures of neutral water clusters journal July 2020
Informing geometric deep learning with electronic interactions to accelerate quantum chemistry journal July 2022
Deep dive into machine learning density functional theory for materials science and chemistry journal April 2022
Protein Data Bank (PDB): Database of Three-Dimensional Structural Information of Biological Macromolecules journal November 1998
Scaling Irregular Applications through Data Aggregation and Software Multithreading
  • Morari, Alessandro; Tumeo, Antonino; Chavarria-Miranda, Daniel
  • 2014 IEEE International Parallel & Distributed Processing Symposium (IPDPS), 2014 IEEE 28th International Parallel and Distributed Processing Symposium https://doi.org/10.1109/IPDPS.2014.117
conference May 2014
Accelerating Scientific Applications With SambaNova Reconfigurable Dataflow Architecture journal March 2021
AWB-GCN: A Graph Convolutional Network Accelerator with Runtime Workload Rebalancing conference October 2020
Reducing Communication in Graph Neural Network Training conference November 2020
Machine learning of accurate energy-conserving molecular force fields journal May 2017
In-Datacenter Performance Analysis of a Tensor Processing Unit conference January 2017
Dgl-Ke conference July 2020
Distributed Hybrid CPU and GPU training for Graph Neural Networks on Billion-Scale Heterogeneous Graphs conference August 2022
A simple on-line bin-packing algorithm journal July 1985
Exploring the GDB-13 chemical space using deep generative models journal March 2019
Foundation Models of Scientific Knowledge for Chemistry: Opportunities, Challenges and Lessons Learned conference January 2022
Neural Message Passing for Quantum Chemistry preprint January 2017
Attention Is All You Need preprint January 2017
Fast Graph Representation Learning with PyTorch Geometric preprint January 2019
Dissecting the Graphcore IPU Architecture via Microbenchmarking preprint January 2019
Directional Message Passing for Molecular Graphs preprint January 2020
Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures preprint January 2020
HydroNet: Benchmark Tasks for Preserving Intermolecular Interactions and Structural Motifs in Predictive and Generative Models for Molecular Data preprint January 2020
Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance preprint January 2021
Large-scale graph representation learning with very deep GNNs and self-supervision preprint January 2021
ChemBERTa-2: Towards Chemical Foundation Models preprint January 2022

Similar Records

Related Subjects

37 INORGANIC, ORGANIC, PHYSICAL, AND ANALYTICAL CHEMISTRY
active learning
artificial intelligence
circuits
cluster chemistry
computational chemistry
graph neural networks
hardware accelerators
high-performance computing
mathematical methods
neural networks
optimization
potential energy