Your browser doesn't support javascript.
loading
Acceleration of Graph Neural Network-Based Prediction Models in Chemistry via Co-Design Optimization on Intelligence Processing Units.
Helal, Hatem; Firoz, Jesun; Bilbrey, Jenna A; Sprueill, Henry; Herman, Kristina M; Krell, Mario Michael; Murray, Tom; Roldan, Manuel Lopez; Kraus, Mike; Li, Ang; Das, Payel; Xantheas, Sotiris S; Choudhury, Sutanay.
Afiliação
  • Helal H; Graphcore, Kett House, Station Rd, Cambridge CB1 2JH, U.K.
  • Firoz J; Advanced Computing, Mathematics and Data Division, Pacific Northwest National Laboratory, 1100 Dexter Ave N, Seattle, Washington 98109, United States.
  • Bilbrey JA; Artificial Intelligence and Data Analytics Division, Pacific Northwest National Laboratory, 902 Battelle Boulevard, Richland, Washington 99352, United States.
  • Sprueill H; Artificial Intelligence and Data Analytics Division, Pacific Northwest National Laboratory, 902 Battelle Boulevard, Richland, Washington 99352, United States.
  • Herman KM; Department of Chemistry, University of Washington, Seattle, Washington 98185, United States.
  • Krell MM; Graphcore, Kett House, Station Rd, Cambridge CB1 2JH, U.K.
  • Murray T; Graphcore, Kett House, Station Rd, Cambridge CB1 2JH, U.K.
  • Roldan ML; Graphcore, Kett House, Station Rd, Cambridge CB1 2JH, U.K.
  • Kraus M; Graphcore, Kett House, Station Rd, Cambridge CB1 2JH, U.K.
  • Li A; Advanced Computing, Mathematics and Data Division, Pacific Northwest National Laboratory, 902 Battelle Boulevard, Richland, Washington 99352, United States.
  • Das P; IBM Research, Yorktown Heights, New York 10598, United States.
  • Xantheas SS; Department of Chemistry, University of Washington, Seattle, Washington 98185, United States.
  • Choudhury S; Advanced Computing, Mathematics and Data Division, Pacific Northwest National Laboratory, 902 Battelle Boulevard, Richland, Washington 99352, United States.
J Chem Inf Model ; 64(5): 1568-1580, 2024 03 11.
Article em En | MEDLINE | ID: mdl-38382011
ABSTRACT
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 codesign 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 codesign 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. We demonstrate that such a codesign approach can reduce the training time of atomistic GNNs and can improve their 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.
Assuntos

Texto completo: 1 Base de dados: MEDLINE Assunto principal: Redes Neurais de Computação / Aceleração Idioma: En Ano de publicação: 2024 Tipo de documento: Article

Texto completo: 1 Base de dados: MEDLINE Assunto principal: Redes Neurais de Computação / Aceleração Idioma: En Ano de publicação: 2024 Tipo de documento: Article