# Optimizing Data Distribution and Kernel Performance for Efficient Training of Chemistry Foundation Models: A Case Study with MACE

arXiv: https://arxiv.org/abs/2504.10700 · HTML: https://mlsanity.com/paper/2504.10700/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 45 references that are in the map.
- First posted on arXiv 14 Apr 2025.

## Builds on (most cited first, 14 of 14)

- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction](https://mlsanity.com/paper/2310.16802/) (2023)
- [Evaluation of the MACE Force Field Architecture: from Medicinal Chemistry to Materials Science](https://mlsanity.com/paper/2305.14247/) (2023)
- [Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN](https://mlsanity.com/paper/2406.12909/) (2024)

## Built on by (most cited first, 1 of 1)

- [The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products](https://mlsanity.com/paper/2506.13523/) (2025)

## Abstract (arXiv)

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.