arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2406.12909cs.LGphysics.comp-ph
收藏

面向原子材料建模的可扩展、可信且节能的预测图基础模型训练:基于HydraGNN的案例研究

Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN

Oak Ridge National Laboratory(橡树岭国家实验室);Lawrence Berkeley National Laboratory(劳伦斯伯克利国家实验室);AMD Research(AMD研究院)

Massimiliano Lupo Pasini, Jong Youl Choi, Kshitij Mehta, Pei Zhang, David Rogers, Jonghyun Bae, Khaled Z. Ibrahim, Ashwin M. Aji, Karl W. Schulz, Jorda Polo, Prasanna Balaprakash

AI总结:

本文基于多头图卷积架构HydraGNN,通过多任务学习训练可扩展、可信且节能的预测图基础模型,利用超1.54亿原子结构在数万GPU上实现近线性强扩展性能。

摘要:

我们展示了使用HydraGNN(一种多头图卷积神经网络架构)开发和训练可扩展、可信且节能的预测图基础模型(GFMs)的工作。HydraGNN在训练规模和数据多样性两方面拓展了图神经网络(GNN)计算的边界。它对消息传递算法进行了抽象,允许对定义GNN中近邻卷积的算法创新进行复现和比较。这项工作讨论了一系列优化措施,这些优化使得GFMs训练能够扩展到数万个GPU,并处理包含数亿个图的数据集。我们的GFMs采用多任务学习(MTL),同时学习原子结构的图级和节点级属性,例如能量和原子力。我们使用超过1.54亿个原子结构进行训练,在两台最先进的美国能源部(US-DOE)超级计算机上展示了该方法的性能并分享了经验教训,这两台计算机分别是国家能源研究科学计算中心的Perlmutter千万亿次级系统和橡树岭领导力计算设施的Frontier百亿亿次级系统。HydraGNN架构使GFM能够在Perlmutter上使用超过2,000个GPU、在Frontier上使用16,000个GPU实现近线性的强扩展性能。

英文摘要:

We present our work on developing and training scalable, trustworthy, and energy-efficient predictive graph foundation models (GFMs) using HydraGNN, a multi-headed graph convolutional neural network architecture. HydraGNN expands the boundaries of graph neural network (GNN) computations in both training scale and data diversity. It abstracts over message passing algorithms, allowing both reproduction of and comparison across algorithmic innovations that define nearest-neighbor convolution in GNNs. This work discusses a series of optimizations that have allowed scaling up the GFMs training to tens of thousands of GPUs on datasets consisting of hundreds of millions of graphs. Our GFMs use multi-task learning (MTL) to simultaneously learn graph-level and node-level properties of atomistic structures, such as energy and atomic forces. Using over 154 million atomistic structures for training, we illustrate the performance of our approach along with the lessons learned on two state-of-the-art United States Department of Energy (US-DOE) supercomputers, namely the Perlmutter petascale system at the National Energy Research Scientific Computing Center and the Frontier exascale system at Oak Ridge Leadership Computing Facility. The HydraGNN architecture enables the GFM to achieve near-linear strong scaling performance using more than 2,000 GPUs on Perlmutter and 16,000 GPUs on Frontier.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑