A Portable and Versatile Limited-Memory BFGS Implementation in PETSc/TAO

2026-07-31

AI Summary

This work addresses the limitations of traditional L-BFGS algorithms in large-scale heterogeneous systems, where high synchronization overhead and poor cache locality hinder performance. Within the PETSc/TAO framework, the authors propose a novel intermediate dense representation of L-BFGS that combines the low memory access cost of recursive formulations with the low synchronization latency of compact representations. They design an efficient history vector management scheme and an update strategy tailored for variable-metric methods. The resulting implementation significantly enhances portability and computational efficiency across CPU and GPU platforms, demonstrating reduced memory traffic and accelerated convergence in single-node experiments on the Polaris and Frontier supercomputers.

Abstract

The limited-memory BFGS (L-BFGS) Hessian update scheme is the critical kernel in many quasi-Newton optimization algorithms. The most common approach to implementing L-BFGS uses $2m$ sequential rank-1 updates as part of solving a linear system when there are $m$ history steps. The performance of this approach suffers when the latency of synchronization is significant, and its poor temporal locality increases the memory traffic when vectors do not fit in cache. The compact dense representation of L-BFGS results in an approach that has minimal synchronization latency and better temporal locality, but it requires an additional pass over the basis vectors and an additional basis that must be recomputed when the $B_0$ matrix changes as in variable-metric methods. In the Portable Extensible Toolkit for Scientific Computation and the Toolkit for Advanced Optimization (PETSc/TAO), we have implemented an intermediate dense formulation of BFGS that retains most of the good characteristics of both the recursive and compact dense approaches. We report single-node performance tests of these implementations on the U.S. Department of Energy's Polaris and Frontier machines, testing both GPU-based and CPU-based computations.

Problem

• L-BFGS
• synchronization latency
• temporal locality
• memory traffic
• quasi-Newton

Innovation

• L-BFGS
• dense formulation
• temporal locality
• synchronization latency
• PETSc/TAO