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

TL;DR

This paper develops and benchmarks GPU-accelerated coupled-cluster algorithms in Python using CuPy and PyTorch, demonstrating significant speedups and architecture-dependent performance improvements on Hopper and Grace Hopper GPUs.

Contribution

Introduces new batching algorithms and a generic tensor contraction protocol for efficient GPU-based CCSD calculations in Python, with comprehensive benchmarking on modern GPU architectures.

Findings

PyTorch outperforms CuPy by 20% on H100.

Both libraries perform similarly on GH200.

Achieved a 10-fold speedup over previous GPU implementation.

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.