Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915374270775296 |
|---|---|
| author | Kaniselvan, Manasa Maeder, Alexander Xia, Chen Hao Ziogas, Alexandros Nikolaos Luisier, Mathieu |
| author_facet | Kaniselvan, Manasa Maeder, Alexander Xia, Chen Hao Ziogas, Alexandros Nikolaos Luisier, Mathieu |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2507_03840 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction Kaniselvan, Manasa Maeder, Alexander Xia, Chen Hao Ziogas, Alexandros Nikolaos Luisier, Mathieu Machine Learning Materials Science Distributed, Parallel, and Cluster Computing Computational Physics 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. |
| title | Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction |
| topic | Machine Learning Materials Science Distributed, Parallel, and Cluster Computing Computational Physics |
| url | https://arxiv.org/abs/2507.03840 |