Efficient Coupled-Cluster Python Frameworks for Next-Generation GPUs: A Comparative Study of CuPy and PyTorch on the Hopper and Grace Hopper Architecture

Abstract

In this work, we introduce new batching algorithms to effectively handle large contractions encountered in coupled-cluster singles and doubles (CCSD) implementations in Python on the Video Random Access Memory (VRAM) of graphical processing units (GPUs), thereby improving performance. Specifically, we benchmark the performance of the CuPy and PyTorch libraries on a single NVIDIA Hopper (H100) and the Grace Hopper (GH200) architectures. We begin by optimizing the particle-particle ladder bottleneck contraction in CCSD using an asymmetric and dynamic splitting recipe, and then move toward a generic tensor contraction protocol that enables tensor contractions to be performed almost exclusively on GPUs. We benchmark our new, fully generic GPU-accelerated coupled-cluster implementations for various molecular systems and basis-set sizes, using both the CuPy and PyTorch libraries. While PyTorch outperforms CuPy on H100 by approximately 20\%, both perform similarly on the GH200 architecture. Compared to our initial GPU implementation [J. Chem. Theory Comput. 2024, 20, 3, 1130--1142], we achieve a 10-fold speedup. In molecular CCSD calculations, we report additional speedups between 3 and 16 for a single CCSD iteration using Cholesky-decomposed electron repulsion integrals compared to our original GPU-CPU hybrid implementation.