面向原子材料建模的可扩展、可信赖且高效的预测图构件基础模型：HydraGNN案例研究
机器学习
2025-04-03 v4 计算物理
摘要
本文介绍了开发和训练可扩展、可信赖且能效的预测图构件基础模型（GFMs）的工作，采用 HydraGNN—a 多头图卷积神经网络架构。HydraGNN 扩展了图神经网络（GNN）计算在训练规模和数据多样性方面的边界。它抽象了信息传递算法，允许复现和比较那些定义图神经网络中最近邻卷积的算法创新。这一工作讨论了一系列优化措施，使 GFMs 的训练规模达到数万块 GPU，数据集包含数亿个图。我们的 GFMs 使用多任务学习（MTL）同时学习原子结构的图级和节点级属性，如能量和原子力。我们使用超过 1.54 亿个原子结构进行训练，展示了我们的做法在美国能源部（US-DOE）两台领先的超级计算机上的性能：位于国家能源研究科学计算中心的 Perlmutter 超算规模系统，以及橡树会计计算设施的 Frontier 超算规模系统。HydraGNN 架构使 GFM 能在 Perlmutter 上使用超过 2000 块 GPU 实现接近线性强度缩放，在 Frontier 上使用 16000 块 GPU 实现强度缩放。
引用
@article{arxiv.2406.12909,
title = {Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN},
author = {Massimiliano Lupo Pasini and Jong Youl Choi and Kshitij Mehta and Pei Zhang and David Rogers and Jonghyun Bae and Khaled Z. Ibrahim and Ashwin M. Aji and Karl W. Schulz and Jorda Polo and Prasanna Balaprakash},
journal= {arXiv preprint arXiv:2406.12909},
year = {2025}
}
备注
51 pages, 32 figures