Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data
Massimiliano Lupo Pasini and Jong Youl Choi and Pei Zhang and Kshitij Mehta and Rylie Weaver and Ashwin M. Aji and Karl W. Schulz and Jorda Polo and Prasanna Balaprakash
使用图神经网络的图形基础模型可以进行可持续、高效的原子建模。 为了应对在预训练过程中处理多源、多保真数据的挑战,最近的研究采用了多任务学习,其中共享消息传递层最初处理输入原子结构,无论来源如何,然后将它们路由到多个解码头,预测数据特定的输出。 这种方法稳定了预训练,增强了模型对未勘探化学区域的可转移性。 大约400万个结构的初步结果令人鼓舞,但关于更大,更多样化的数据集和超级计算机上的可扩展性的可推广性仍然存在问题。 我们提出了一种多任务并行方法,通过 GPU 加速将每个头分布在计算资源上。 在开源HydraGNN架构中实现,我们的方法在来自五个数据集的2400多万个结构上进行了训练,并在Perlmutter,Aurora和Frontier超级计算机上进行了测试,展示了所有三种高度异构计算架构的高效扩展。
Graph foundation models using graph neural networks promise sustainable, efficient atomistic modeling. To tackle challenges of processing multi-source, multi-fidelity data during pre-training, recent studies employ multi-task learning, in which shared message passing layers initially process input atomistic structures regardless of source, then route them to multiple decoding heads that predict data-specific outputs. This approach stabilizes pre-training and enhances a model's transferability to...
