Showing 1–10 of 10 results for author: Firoz, J

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

Authors: Jesun Firoz, Franco Pellegrini, Mario Geiger, Darren Hsu, Jenna A. Bilbrey, Han-Yi Chou, Maximilian Stadler, Markus Hoehnerbach, Tingyu Wang, Dejun Lin, Emine Kucukbenli, Henry W. Sprueill, Ilyes Batatia, Sotiris S. Xantheas, MalSoon Lee, Chris Mundy, Gabor Csanyi, Justin S. Smith, Ponnuswamy Sadayappan, Sutanay Choudhury

Abstract: 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 la… ▽ More 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. △ Less

Submitted 14 April, 2025; originally announced April 2025.

Comments: Accepted at The 34th ACM International Symposium on High-Performance Parallel and Distributed Computing (HPDC 2025)

Final Report for CHESS: Cloud, High-Performance Computing, and Edge for Science and Security

Authors: Nathan Tallent, Jan Strube, Luanzheng Guo, Hyungro Lee, Jesun Firoz, Sayan Ghosh, Bo Fang, Oceane Bel, Steven Spurgeon, Sarah Akers, Christina Doty, Erol Cromwell

Abstract: Automating the theory-experiment cycle requires effective distributed workflows that utilize a computing continuum spanning lab instruments, edge sensors, computing resources at multiple facilities, data sets distributed across multiple information sources, and potentially cloud. Unfortunately, the obvious methods for constructing continuum platforms, orchestrating workflow tasks, and curating dat… ▽ More Automating the theory-experiment cycle requires effective distributed workflows that utilize a computing continuum spanning lab instruments, edge sensors, computing resources at multiple facilities, data sets distributed across multiple information sources, and potentially cloud. Unfortunately, the obvious methods for constructing continuum platforms, orchestrating workflow tasks, and curating datasets over time fail to achieve scientific requirements for performance, energy, security, and reliability. Furthermore, achieving the best use of continuum resources depends upon the efficient composition and execution of workflow tasks, i.e., combinations of numerical solvers, data analytics, and machine learning. Pacific Northwest National Laboratory's LDRD "Cloud, High-Performance Computing (HPC), and Edge for Science and Security" (CHESS) has developed a set of interrelated capabilities for enabling distributed scientific workflows and curating datasets. This report describes the results and successes of CHESS from the perspective of open science. △ Less

Submitted 21 October, 2024; originally announced October 2024.

Report number: Pacific Northwest National Laboratory, PNNL-36859 ACM Class: C.2.4; C.4; D.1.3; J.2; K.6.4

Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry

Authors: Hatem Helal, Jesun Firoz, Jenna Bilbrey, Mario Michael Krell, Tom Murray, Ang Li, Sotiris Xantheas, Sutanay Choudhury

Abstract: Molecular property calculations are the bedrock of chemical physics. High-fidelity \textit{ab initio} modeling techniques for computing the molecular properties can be prohibitively expensive, and motivate the development of machine-learning models that make the same predictions more efficiently. Training graph neural networks over large molecular databases introduces unique computational challeng… ▽ More Molecular property calculations are the bedrock of chemical physics. High-fidelity \textit{ab initio} modeling techniques for computing the molecular properties can be prohibitively expensive, and motivate the development of machine-learning models that make the same predictions more efficiently. Training graph neural networks over large molecular databases introduces unique computational challenges such as the need to process millions of small graphs with variable size and support communication patterns that are distinct from learning over large graphs such as social networks. This paper demonstrates a novel hardware-software co-design approach to scale up the training of graph neural networks for molecular property prediction. We introduce an algorithm to coalesce the batches of molecular graphs into fixed size packs to eliminate redundant computation and memory associated with alternative padding techniques and improve throughput via minimizing communication. We demonstrate the effectiveness of our co-design approach by providing an implementation of a well-established molecular property prediction model on the Graphcore Intelligence Processing Units (IPU). We evaluate the training performance on multiple molecular graph databases with varying degrees of graph counts, sizes and sparsity. We demonstrate that such a co-design approach can reduce the training time of such molecular property prediction models from days to less than two hours, opening new possibilities for AI-driven scientific discovery. △ Less

Submitted 24 November, 2022; originally announced November 2022.

MSREP: A Fast yet Light Sparse Matrix Framework for Multi-GPU Systems

Authors: Jieyang Chen, Chenhao Xie, Jesun S Firoz, Jiajia Li, Shuaiwen Leon Song, Kevin Barker, Mark Raugas, Ang Li

Abstract: Sparse linear algebra kernels play a critical role in numerous applications, covering from exascale scientific simulation to large-scale data analytics. Offloading linear algebra kernels on one GPU will no longer be viable in these applications, simply because the rapidly growing data volume may exceed the memory capacity and computing power of a single GPU. Multi-GPU systems nowadays being ubiqui… ▽ More Sparse linear algebra kernels play a critical role in numerous applications, covering from exascale scientific simulation to large-scale data analytics. Offloading linear algebra kernels on one GPU will no longer be viable in these applications, simply because the rapidly growing data volume may exceed the memory capacity and computing power of a single GPU. Multi-GPU systems nowadays being ubiquitous in supercomputers and data-centers present great potentials in scaling up large sparse linear algebra kernels. In this work, we design a novel sparse matrix representation framework for multi-GPU systems called MSREP, to scale sparse linear algebra operations based on our augmented sparse matrix formats in a balanced pattern. Different from dense operations, sparsity significantly intensifies the difficulty of distributing the computation workload among multiple GPUs in a balanced manner. We enhance three mainstream sparse data formats -- CSR, CSC, and COO, to enable fine-grained data distribution. We take sparse matrix-vector multiplication (SpMV) as an example to demonstrate the efficiency of our MSREP framework. In addition, MSREP can be easily extended to support other sparse linear algebra kernels based on the three fundamental formats (i.e., CSR, CSC and COO). △ Less

Submitted 15 September, 2022; originally announced September 2022.

High-order Line Graphs of Non-uniform Hypergraphs: Algorithms, Applications, and Experimental Analysis

Authors: Xu T. Liu, Jesun Firoz, Sinan Aksoy, Ilya Amburg, Andrew Lumsdaine, Cliff Joslyn, Assefaw H. Gebremedhin, Brenda Praggastis

Abstract: Hypergraphs offer flexible and robust data representations for many applications, but methods that work directly on hypergraphs are not readily available and tend to be prohibitively expensive. Much of the current analysis of hypergraphs relies on first performing a graph expansion -- either based on the nodes (clique expansion), or on the edges (line graph) -- and then running standard graph anal… ▽ More Hypergraphs offer flexible and robust data representations for many applications, but methods that work directly on hypergraphs are not readily available and tend to be prohibitively expensive. Much of the current analysis of hypergraphs relies on first performing a graph expansion -- either based on the nodes (clique expansion), or on the edges (line graph) -- and then running standard graph analytics on the resulting representative graph. However, this approach suffers from massive space complexity and high computational cost with increasing hypergraph size. Here, we present efficient, parallel algorithms to accelerate and reduce the memory footprint of higher-order graph expansions of hypergraphs. Our results focus on the edge-based $s$-line graph expansion, but the methods we develop work for higher-order clique expansions as well. To the best of our knowledge, ours is the first framework to enable hypergraph spectral analysis of a large dataset on a single shared-memory machine. Our methods enable the analysis of datasets from many domains that previous graph-expansion-based models are unable to provide. The proposed $s$-line graph computation algorithms are orders of magnitude faster than state-of-the-art sparse general matrix-matrix multiplication methods, and obtain approximately $5-31{\times}$ speedup over a prior state-of-the-art heuristic-based algorithm for $s$-line graph computation. △ Less

Submitted 27 January, 2022; originally announced January 2022.

Comments: Accepted at "36th IEEE International Parallel & Distributed Processing Symposium (IPDPS '22)"

Report number: PNNL-SA-167812 MSC Class: 05C65(Primary); 05C85(Secondary); 68W10(Secondary) ACM Class: G.2.2

SpectralFly: Ramanujan Graphs as Flexible and Efficient Interconnection Networks

Authors: Stephen Young, Sinan Aksoy, Jesun Firoz, Roberto Gioiosa, Tobias Hagge, Mark Kempton, Juan Escobedo, Mark Raugas

Abstract: In recent years, graph theoretic considerations have become increasingly important in the design of HPC interconnection topologies. One approach is to seek optimal or near-optimal families of graphs with respect to a particular graph theoretic property, such as diameter. In this work, we consider topologies which optimize the spectral gap. We study a novel HPC topology, SpectralFly, designed aroun… ▽ More In recent years, graph theoretic considerations have become increasingly important in the design of HPC interconnection topologies. One approach is to seek optimal or near-optimal families of graphs with respect to a particular graph theoretic property, such as diameter. In this work, we consider topologies which optimize the spectral gap. We study a novel HPC topology, SpectralFly, designed around the Ramanujan graph construction of Lubotzky, Phillips, and Sarnak (LPS). We show combinatorial properties, such as diameter, bisection bandwidth, average path length, and resilience to link failure, of SpectralFly topologies are better than, or comparable to, similarly constrained DragonFly, SlimFly, and BundleFly topologies. Additionally, we simulate the performance of SpectralFly on a representative sample of micro-benchmarks using the Structure Simulation Toolkit Macroscale Element Library simulator and study cost-minimizing layouts, demonstrating considerable benefit of the SpectralFly topology. △ Less

Submitted 14 February, 2022; v1 submitted 23 April, 2021; originally announced April 2021.

Fast and Scalable Sparse Triangular Solver for Multi-GPU Based HPC Architectures

Authors: Chenhao Xie, Jieyang Chen, Jesun S Firoz, Jiajia Li, Shuaiwen Leon Song, Kevin Barker, Mark Raugas, Ang Li

Abstract: Designing efficient and scalable sparse linear algebra kernels on modern multi-GPU based HPC systems is a daunting task due to significant irregular memory references and workload imbalance across the GPUs. This is particularly the case for Sparse Triangular Solver (SpTRSV) which introduces additional two-dimensional computation dependencies among subsequent computation steps. Dependency informati… ▽ More Designing efficient and scalable sparse linear algebra kernels on modern multi-GPU based HPC systems is a daunting task due to significant irregular memory references and workload imbalance across the GPUs. This is particularly the case for Sparse Triangular Solver (SpTRSV) which introduces additional two-dimensional computation dependencies among subsequent computation steps. Dependency information is exchanged and shared among GPUs, thus warrant for efficient memory allocation, data partitioning, and workload distribution as well as fine-grained communication and synchronization support. In this work, we demonstrate that directly adopting unified memory can adversely affect the performance of SpTRSV on multi-GPU architectures, despite linking via fast interconnect like NVLinks and NVSwitches. Alternatively, we employ the latest NVSHMEM technology based on Partitioned Global Address Space programming model to enable efficient fine-grained communication and drastic synchronization overhead reduction. Furthermore, to handle workload imbalance, we propose a malleable task-pool execution model which can further enhance the utilization of GPUs. By applying these techniques, our experiments on the NVIDIA multi-GPU supernode V100-DGX-1 and DGX-2 systems demonstrate that our design can achieve on average 3.53x (up to 9.86x) speedup on a DGX-1 system and 3.66x (up to 9.64x) speedup on a DGX-2 system with 4-GPUs over the Unified-Memory design. The comprehensive sensitivity and scalability studies also show that the proposed zero-copy SpTRSV is able to fully utilize the computing and communication resources of the multi-GPU system. △ Less

Submitted 12 December, 2020; originally announced December 2020.

Parallel Algorithms and Heuristics for Efficient Computation of High-Order Line Graphs of Hypergraphs

Authors: Xu T. Liu, Jesun Firoz, Andrew Lumsdaine, Cliff Joslyn, Sinan Aksoy, Brenda Praggastis, Assefaw Gebremedhin

Abstract: This paper considers structures of systems beyond dyadic (pairwise) interactions and investigates mathematical modeling of multi-way interactions and connections as hypergraphs, where captured relationships among system entities are set-valued. To date, in most situations, entities in a hypergraph are considered connected as long as there is at least one common "neighbor". However, minimal commona… ▽ More This paper considers structures of systems beyond dyadic (pairwise) interactions and investigates mathematical modeling of multi-way interactions and connections as hypergraphs, where captured relationships among system entities are set-valued. To date, in most situations, entities in a hypergraph are considered connected as long as there is at least one common "neighbor". However, minimal commonality sometimes discards the "strength" of connections and interactions among groups. To this end, considering the "width" of a connection, referred to as the $s$-overlap of neighbors, provides more meaningful insights into how closely the communities or entities interact with each other. In addition, $s$-overlap computation is the fundamental kernel to construct the line graph of a hypergraph, a low-order approximation of the hypergraph which can carry significant information about the original hypergraph. Subsequent stages of a data analytics pipeline then can apply highly-tuned graph algorithms on the line graph to reveal important features. Given a hypergraph, computing the $s$-overlaps by exhaustively considering all pairwise entities can be computationally prohibitive. To tackle this challenge, we develop efficient algorithms to compute $s$-overlaps and the corresponding line graph of a hypergraph. We propose several heuristics to avoid execution of redundant work and improve performance of the $s$-overlap computation. Our parallel algorithm, combined with these heuristics, is orders of magnitude (more than $10\times$) faster than the naive algorithm in all cases and the SpGEMM algorithm with filtration in most cases (especially with large $s$ value). △ Less

Submitted 15 July, 2021; v1 submitted 22 October, 2020; originally announced October 2020.

Comments: 11 pages

Report number: PNNL-SA-164086 MSC Class: 05C65 ACM Class: G.2.2

The Anatomy of Large-Scale Distributed Graph Algorithms

Authors: Jesun Sahariar Firoz, Thejaka Amila Kanewala, Marcin Zalewski, Martina Barnas, Andrew Lumsdaine

Abstract: The increasing complexity of the software/hardware stack of modern supercomputers results in explosion of parameters. The performance analysis becomes a truly experimental science, even more challenging in the presence of massive irregularity and data dependency. We analyze how the existing body of research handles the experimental aspect in the context of distributed graph algorithms (DGAs). We d… ▽ More The increasing complexity of the software/hardware stack of modern supercomputers results in explosion of parameters. The performance analysis becomes a truly experimental science, even more challenging in the presence of massive irregularity and data dependency. We analyze how the existing body of research handles the experimental aspect in the context of distributed graph algorithms (DGAs). We distinguish algorithm-level contributions, often prioritized by authors, from runtime-level concerns that are harder to place. We show that the runtime is such an integral part of DGAs that experimental results are difficult to interpret and extrapolate without understanding the properties of the runtime used. We argue that in order to gain understanding about the impact of runtimes, more information needs to be gathered. To begin this process, we provide an initial set of recommendations for describing DGA results based on our analysis of the current state of the field. △ Less

Submitted 23 July, 2015; originally announced July 2015.

ACM Class: D.1.3

The 1.375 Approximation Algorithm for Sorting by Transpositions Can Run in $O(n\log n)$ Time

Authors: Jesun Sahariar Firoz, Masud Hasan, Ashik Zinnat Khan, M. Sohel Rahman

Abstract: Sorting a Permutation by Transpositions (SPbT) is an important problem in Bioinformtics. In this paper, we improve the running time of the best known approximation algorithm for SPbT. We use the permutation tree data structure of Feng and Zhu and improve the running time of the 1.375 Approximation Algorithm for SPbT of Elias and Hartman to $O(n\log n)$. The previous running time of EH algorithm… ▽ More Sorting a Permutation by Transpositions (SPbT) is an important problem in Bioinformtics. In this paper, we improve the running time of the best known approximation algorithm for SPbT. We use the permutation tree data structure of Feng and Zhu and improve the running time of the 1.375 Approximation Algorithm for SPbT of Elias and Hartman to $O(n\log n)$. The previous running time of EH algorithm was $O(n^2)$. △ Less

Submitted 17 October, 2009; originally announced October 2009.

Comments: 5 pages