Optimizing Data Distribution and Kernel Performance for Efficient Training of Chemistry Foundation Models: A Case Study with MACE
作者: Jesun Firoz, Franco Pellegrini, Mario Geiger, Darren Hsu, Jenna A. Bilbrey, Han-Yi Chou, Maximilian Stadler, Markus Hoehnerbach, Tingyu Wang, Dejun Lin, Emine Kucukbenli, Henry W. Sprueill, Ilyes Batatia, Sotiris S. Xantheas, MalSoon Lee, Chris Mundy, Gabor Csanyi, Justin S. Smith, Ponnuswamy Sadayappan, Sutanay Choudhury
分类: cs.DC, cs.AI
发布日期: 2025-04-14
备注: Accepted at The 34th ACM International Symposium on High-Performance Parallel and Distributed Computing (HPDC 2025)
💡 一句话要点
优化数据分布与内核性能，加速化学基础模型MACE的训练。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 化学基础模型 图神经网络 数据分布优化 内核优化 多目标装箱问题
📋 核心要点
- 现有化学基础模型训练面临数据分布不均和计算效率低下的问题，尤其是在处理大量不同大小的分子图时。
- 论文提出一种结合数据分布优化和内核优化的方法，通过多目标装箱算法平衡数据负载，并优化MACE中的关键张量收缩内核。
- 实验结果表明，该方法显著加速了MACE的训练过程，在大型数据集和多GPU环境下，每个epoch的训练时间缩短至原来的六分之一。
📝 摘要（中文）
化学基础模型(CFMs)利用作用于3D分子图结构的图神经网络(GNNs)，正成为计算化学家和材料科学家不可或缺的工具。这些模型有助于理解物质和发现新的分子与材料。与处理大型同构图的GNN不同，CFM使用的GNN处理大量不同大小的几何图，需要与为大型同构GNN开发的不同优化策略。本文针对最先进的CFM——MACE，提出了对CFM训练的两个关键阶段的优化：数据分布和模型训练。我们通过将数据分布中的负载平衡问题建模为多目标装箱问题来解决这一挑战。我们提出了一种高效、快速和实用的迭代算法，确保高效的数据分布。在训练阶段，我们将对称张量收缩确定为MACE中的关键计算内核，并优化该内核以提高整体性能。我们平衡的数据分布和内核优化的组合方法显著增强了MACE的训练过程。实验结果表明，在拥有260万个样本的数据集上，使用740个GPU进行训练时，每个epoch的执行时间从12分钟大幅缩短到2分钟。
🔬 方法详解
问题定义：化学基础模型（CFMs）的训练，特别是像MACE这样的模型，需要处理大量的、大小不一的分子图数据。传统的数据分配方法容易导致负载不均衡，某些GPU需要处理更多的数据，从而成为训练的瓶颈。此外，MACE模型中的对称张量收缩操作是计算密集型的，其效率直接影响整体训练速度。
核心思路：论文的核心思路是双管齐下：首先，通过优化数据分布，确保每个GPU的负载尽可能均衡；其次，通过优化MACE模型中的关键计算内核（即对称张量收缩），提高计算效率。数据分布优化被建模为一个多目标装箱问题，旨在最小化负载不均衡。内核优化则针对MACE的特定计算模式进行定制。
技术框架：整体框架包含两个主要阶段：数据预处理和模型训练。在数据预处理阶段，使用提出的迭代算法解决多目标装箱问题，将数据分配到不同的GPU上。在模型训练阶段，优化后的对称张量收缩内核被集成到MACE模型中，加速训练过程。
关键创新：论文的关键创新在于将数据分布问题建模为多目标装箱问题，并提出了一种高效的迭代算法来解决它。这种方法能够有效地平衡各个GPU的负载，避免了传统数据分配方法中的瓶颈。此外，针对MACE模型的特定计算模式，对对称张量收缩内核进行了优化，进一步提高了计算效率。
关键设计：数据分布优化采用迭代算法，旨在最小化负载不均衡。具体而言，算法尝试将数据样本分配到不同的“箱子”（即GPU），并根据负载情况进行调整，直到达到一个较好的平衡状态。对称张量收缩内核的优化细节未知，但可以推测涉及内存访问模式的优化、并行计算的优化，以及可能的硬件加速技术。
🖼️ 关键图片
📊 实验亮点
实验结果表明，通过结合优化的数据分布和内核优化，MACE模型的训练速度得到了显著提升。在拥有260万个样本的数据集上，使用740个GPU进行训练时，每个epoch的执行时间从12分钟大幅缩短到2分钟，实现了6倍的加速。这表明该方法在实际应用中具有很高的价值。
🎯 应用场景
该研究成果可广泛应用于计算化学、材料科学等领域，加速新分子和新材料的发现过程。通过提高化学基础模型的训练效率，研究人员可以更快地探索化学空间，筛选具有特定性质的分子，从而加速药物研发、催化剂设计等过程。此外，该方法也可推广到其他需要处理大量小图数据的图神经网络应用中。
📄 摘要（原文）
Chemistry Foundation Models (CFMs) that leverage Graph Neural Networks (GNNs) operating on 3D molecular graph structures are becoming indispensable tools for computational chemists and materials scientists. These models facilitate the understanding of matter and the discovery of new molecules and materials. In contrast to GNNs operating on a large homogeneous graphs, GNNs used by CFMs process a large number of geometric graphs of varying sizes, requiring different optimization strategies than those developed for large homogeneous GNNs. This paper presents optimizations for two critical phases of CFM training: data distribution and model training, targeting MACE - a state-of-the-art CFM. We address the challenge of load balancing in data distribution by formulating it as a multi-objective bin packing problem. We propose an iterative algorithm that provides a highly effective, fast, and practical solution, ensuring efficient data distribution. For the training phase, we identify symmetric tensor contraction as the key computational kernel in MACE and optimize this kernel to improve the overall performance. Our combined approach of balanced data distribution and kernel optimization significantly enhances the training process of MACE. Experimental results demonstrate a substantial speedup, reducing per-epoch execution time for training from 12 to 2 minutes on 740 GPUs with a 2.6M sample dataset.