exa-AMD: An Exascale-Ready Framework for Accelerating the Discovery and Design of Functional Materials

TL;DR

exa-AMD tackles the combinatorial explosion in multinary materials search by integrating an end-to-end, modular workflow that couples rapid ML screening with rigorous first-principles validation on exascale HPC. The framework automates structure generation, ML-based stability screening, DFT relaxation, and convex-hull analysis, all orchestrated by Parsl to achieve scalable, heterogeneous-resource execution. Key contributions include a large seed/template pool, two-generation structure generation strategies, FPGA-accelerated ML screening, and demonstrable exascale-ready performance on complex ternary and quaternary systems, yielding updated phase diagrams and new candidate materials. The work enables rapid exploration of vast chemical spaces with minimal manual intervention, supports reproducibility through open-source code and documentation, and lays the groundwork for future enhancements like adaptive genetic algorithms and ML potentials to access novel structure motifs.

Abstract

We present exa-AMD, an open-source, high-performance framework designed for accelerated materials discovery on modern supercomputers. exa-AMD overcomes key computational bottlenecks in large-scale structure prediction through task-based parallelization, adaptive load balancing, and optimized data management for CPU and GPU architectures. The framework automates the end-to-end workflow, from generating candidate structures to evaluating formation energies and updating phase diagrams. Its modular design allows users to easily replace or extend components with custom machine learning models, alternative initial structure templates, and future structure generators, enabling flexible integration with emerging AI approaches. We demonstrate strong scaling across high-performance computing platforms and highlight applications to Na-B-C, Ce-Co-B, and Fe-Co-Zr systems, establishing exa-AMD as a robust and exascale-ready tool for accelerating the discovery and design of functional materials. exa-AMD is publicly available on GitHub, with detailed documentation and reproducible test cases to support community engagement and collaborative research.

Figures (9)

• Figure 1: Schematic workflow consists of five major steps: (1) crystal structures construction; (2) formation energy prediction using ML models; (3) structure selection; (4) first-principles calculations; (5) post-processing.
• Figure 2: Illustration of the two primary high-throughput structure generation methods. (a) Element substitution: systematic replacement of atomic species at fixed crystallographic sites. (b) Lattice-volume scaling: uniform expansion or contraction of the unit cell volume by applying scaling factors.
• Figure 3: Example mapping of exa-AMD stages to hardware via Parsl executors: CPU nodes (blue) and GPU nodes (green).
• Figure 4: Strong scaling of exa-AMD workflow for the Na-B-C system on NERSC's Perlmutter supercomputer. Wall-clock times are shown for both GPU and CPU nodes (1, 2, 4, 8, 16, 32). The workflow exhibits near-linear speed-up and substantial GPU acceleration at all node counts.
• Figure 5: Benchmark results for the Ce-Co-B system on NERSC's Perlmutter supercomputer, executed on GPU architectures for 4 to 256 nodes, and on CPU for 4 to 128 nodes. Performance demonstrates efficient strong scaling and consistent GPU advantage for large-scale campaigns.