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

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.