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 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.

Recently, graph neural network (GNN) has shown significant strength in predicting the quantum mechanical properties of molecules. Based on GNN, the DimeNet++ leverages both distance information of atomic pairs and angle information of atomic triplets via message passing mechanism to predict quantum mechanical properties of molecules and has achieved state-of-the-art results. However, there are more than 10 thousand operators in DimeNet++, which results in low GPU utilization and large CPU launch overhead. The extensive time taken for the training of DimeNet++ is a significant drawback. The training period of DimeNet++ exceeds one month on a single NVIDIA A100 GPU. A common method for reducing training time involves employing data parallelism, which equally distributes the global batch across each GPU. However, data-parallel task partitioning, by default, does not consider load imbalance within the batch. This load imbalance leads to considerable synchronization overhead in a multi-GPU setting, reducing the overall efficiency of the parallelism. For the strong-scaling scenario, it results in 32% of the compute resource being wasted. In light of these observations, we propose a novel approach, FastDimeNet++, which delivers high GPU utilization, low CPU overhead, and extensive scalability, achieved through a series of optimization strategies. These include (i) a communication-free load-balancing sampler, (ii) computation graph reconstruction, and (iii) kernel fusion and redundancy bypass. Our experiments demonstrate that FastDimeNet++ achieves a GPU utilization rate of approximately 88% based on a mini-batch size of 4. Furthermore, we scale FastDimeNet++ to 512 GPUs, reaching 2.8 PetaFLOPS. In the MLPerf HPC V1.0, the winning DimeNet++ submission required a total training time of 111.86 minutes, whereas FastDimeNet++ introduced for the MLPerf HPC V2.0 required just 21.93 minutes, demonstrating a significant performance improvement of over 5 ×.

We propose the geometry-informed neural operator (GINO), a highly efficient approach for learning the solution operator of large-scale partial differential equations with varying geometries. GINO uses a signed distance function (SDF) representation of the input shape and neural operators based on graph and Fourier architectures to learn the solution operator. The graph neural operator handles irregular grids and transforms them into and from regular latent grids on which Fourier neural operator can be efficiently applied. We provide an efficient implementation of GINO using an optimized hashing approach, which allows efficient learning in a shared, compressed latent space with reduced computation and memory costs. GINO is discretization-invariant, meaning the trained model can be applied to arbitrary discretizations of the continuous domain and applies to any shape or resolution. To empirically validate the performance of our method on large-scale simulation, we generate the industry-standard aerodynamics dataset of 3D vehicle geometries with Reynolds numbers as high as five million. For this large-scale 3D fluid simulation, numerical methods are expensive to compute surface pressure. We successfully trained GINO to predict the pressure on car surfaces using only five hundred data points. The cost-accuracy experiments show a 26,000x speed-up compared to optimized GPU-based computational fluid dynamics (CFD) simulators on computing the drag coefficient. When tested on new combinations of geometries and boundary conditions (inlet velocities), GINO obtains a one-fourth reduction in error rate compared to deep neural network approaches.

The interdependence between nodes in graphs is key to improve class prediction on nodes, utilized in approaches like Label Probagation (LP) or in Graph Neural Networks (GNNs). Nonetheless, uncertainty estimation for non-independent node-level predictions is under-explored. In this work, we explore uncertainty quantification for node classification in three ways: (1) We derive three axioms explicitly characterizing the expected predictive uncertainty behavior in homophilic attributed graphs.(2) We propose a new model Graph Posterior Network (GPN) which explicitly performs Bayesian posterior updates for predictions on interdependent nodes. GPN provably obeys the proposed axioms. (3) We extensively evaluate GPN and a strong set of baselines on semi-supervised node classification including detection of anomalous features, and detection of left-out classes. GPN outperforms existing approaches for uncertainty estimation in the experiments.