# Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 52 references that are in the map.
- First posted on arXiv 4 Jul 2025.

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

- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian](https://mlsanity.com/paper/2306.04922/) (2023)
- [Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size](https://mlsanity.com/paper/2304.10061/) (2023)
- [Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations](https://mlsanity.com/paper/2203.09697/) (2022)
- [Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs](https://mlsanity.com/paper/2111.06483/) (2021)
- [A Data-Centric Approach to Extreme-Scale Ab initio Dissipative Quantum Transport Simulations](https://mlsanity.com/paper/1912.10024/) (2019)
- [Understanding Data Movement in Tightly Coupled Heterogeneous Systems: A Case Study with the Grace Hopper Superchip](https://mlsanity.com/paper/2408.11556/) (2024)

## Abstract (arXiv)

Equivariant Graph Neural Networks (eGNNs) trained on density-functional theory (DFT) data can potentially perform electronic structure prediction at unprecedented scales, enabling investigation of the electronic properties of materials with extended defects, interfaces, or exhibiting disordered phases. However, as interactions between atomic orbitals typically extend over 10+ angstroms, the graph representations required for this task tend to be densely connected, and the memory requirements to perform training and inference on these large structures can exceed the limits of modern GPUs. Here we present a distributed eGNN implementation which leverages direct GPU communication and introduce a partitioning strategy of the input graph to reduce the number of embedding exchanges between GPUs. Our implementation shows strong scaling up to 128 GPUs, and weak scaling up to 512 GPUs with 87% parallel efficiency for structures with 3,000 to 190,000 atoms on the Alps supercomputer.