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

Area of Science:

• Computational Physics and Chemistry
• Materials Science
• Machine Learning

Background:

• High-fidelity ab initio methods for atomic structure and property prediction are computationally expensive.
• Machine learning (ML) models, particularly graph neural networks (GNNs), offer a more efficient alternative.
• Training GNNs on large atomistic databases presents unique computational challenges related to variable graph sizes and communication patterns.

Purpose of the Study:

• To develop and demonstrate a novel hardware-software codesign approach to scale up the training of atomistic GNNs.
• To improve the efficiency and performance of GNN-based atomic structure and property prediction.
• To optimize GNN training for specialized hardware like Graphcore's Intelligence Processing Unit (IPU).

Main Methods:

• Formulated batching of variable-size atomistic graphs as a bin packing problem with a hardware-agnostic algorithm to minimize redundancy and communication.
• Implemented hardware-specific optimizations for Graphcore IPUs, including a planner, vectorization for gather-scatter operations, and model-specific optimizations (merged communication collectives, optimized softplus).
• Deployed a well-established atomistic GNN model on Graphcore IPUs and evaluated performance on diverse atomistic graph databases.

Main Results:

• The codesign approach reduced training time for atomistic GNNs by up to 1.5× compared to baseline IPU implementations.
• Achieved an average speedup of 1.8× on IPUs compared to a Nvidia GPU-based implementation for the atomistic GNN.
• Demonstrated improved performance and efficiency across atomistic graph databases with varying characteristics (counts, sizes, sparsity).

Conclusions:

• The proposed hardware-software codesign effectively scales up the training of atomistic GNNs for structure and property prediction.
• This approach significantly enhances computational efficiency, leading to faster training times and improved model performance.
• The optimized IPU implementation shows superior performance compared to GPU-based solutions, highlighting the potential of specialized hardware for scientific ML.