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

A PORTABLE AND VERSATILE LIMITED-MEMORY BFGS IMPLEMENTATION IN PETSC/TAO

arXiv:2608.00196v1 [cs.DC] 31 Jul 2026

HANSOL SUH∗ , TOBIN ISAAC∗ , ALP DENER∗ , TODD MUNSON∗ , HONG ZHANG∗ , AND RICHARD TRAN MILLS∗ 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 B0 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. Key words. limited-memory BFGS, quasi-Newton methods, numerical optimization, performance portability, GPU computing, PETSc/TAO AMS subject classifications. 65K05, 90C30, 90C53, 65Y05, 65Y10

1. Introduction. The limited-memory BFGS (L-BFGS) Hessian update scheme [13] is an important tool in numerical optimization: it is the workhorse of quasiNewton unconstrained and bound-constrained smooth optimization. It is implemented in PETSc/TAO (the Portable, Extensible Toolkit for Scientific Computation and the Toolkit for Advanced Optimization [2, 3]) as a particular instance of a limited-memory variable metric operator, which can be used in quasi-Newton methods to solve nonlinear systems of equations and unconstrained or bound-constrained optimization problems. The BFGS update to a base Hessian B0 or Hessian inverse H0 = B0−1 can also appear in other quasi-Newton optimization algorithms, such as quasi-Newton Krylov methods, globalized by linesearch or trust-region methods. All of these optimizers are available in PETSc/TAO. 2. Background. L-BFGS is sometimes taken to mean a complete optimization algorithm for the unconstrained minimization of a smooth function f (x). In this work we take L-BFGS to mean only the approximation scheme Hk ≈ (∇2 f (xk ))−1 for the inverse Hessian based on a sequence {xk } of solution vectors and the corresponding sequence {gk } of gradients. In this scheme, H0 ∈ RN ×N is taken to be a base symmetric positive definite (SPD) matrix (often a scaled identity but with other common choices as well; see section 3). Given two successive approximate solutions and their gradients, (xk−1 , gk−1 ) and (xk , gk ), the full BFGS update is defined by H0full ≜ H0 and (2.1)

full T Hkfull ≜ (I − Φk−1 )T Hk−1 (I − Φk−1 ) + sk−1 d−1 k−1 sk−1 ,

∗ Mathematics and Computer Science Division, Argonne National Laboratory, Lemont, IL 60439 (hsuh@anl.gov, toby.isaac@gmail.com, alp.dener@gmail.com, tmunson@anl.gov, hongzhang@anl.gov, rtmills@anl.gov).

1

2

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

where (2.2)

sk−1 ≜ xk − xk−1 ,

yk−1 ≜ gk − gk−1 ,

(2.3)

dk−1 ≜ sTk−1 yk−1 ,

T Φk−1 ≜ yk−1 d−1 k−1 sk−1 .

For large problems it is not efficient to form Hkfull explicitly, as this requires O(N 2 ) work and storage. It is more common to choose a fixed history size m ≪ N and implement an approximation Hkm where only the last m pairs of s and y vectors are kept. Formally, the L-BFGS approximate inverse Hessian Hkm is defined by ( H0 , m = 0 or k = 0, m (2.4) Hk ≜ m−1 −1 T T (I − Φk−1 ) Hk−1 (I − Φk−1 ) + sk−1 dk−1 sk−1 otherwise. Letting Bkm ≜ (Hkm )−1 be the corresponding approximate Hessian, one can show that (2.5)

Bkm = B0 +

k−1 X

−1 T T yi d−1 i yi − pm−(k−i),i fm−(k−i),i pm−(k−i),i ,

i=max{0,k−m}

where (2.6)

pm,k ≜ Bkm sk ,

fm,k ≜ sTk Bkm sk .

In this work we compare the efficiency of different approaches of implementing (2.4). We assume that an implementation B of L-BFGS must provide the following operations required of all quasi-Newton Hessian approximations: m m 1. B.update(xk , gk ): Update B from representing Bk−1 /Hk−1 to Bkm /Hkm . m 2. B.solve(u, v): Compute v = Hk u. 3. B.mult(u, v): Compute v = Bkm u. Different optimization algorithms use these core operations in different ways. An important one is unconstrained quasi-Newton optimization globalized by a line search. Algorithm 1: Quasi-Newton optimization with line search Until convergence g ← ∇f (x) 3 B.update(x, g) 4 B.solve(g, p) 5 x ← x − αp, where α is determined by line search on f (x − αp).

1

2

In algorithm 1, B.mult is not used, but a more complex algorithm like a boundconstrained quasi-Newton method globalized by a trust region may require B.mult and may use B.solve(v, p) with a generic right-hand side vector v unrelated to the gradient. 3. L-BFGS use cases in PETSc/TAO. PETSc/TAO is used to solve problems of vastly different sizes and supports MPI parallelism and GPU acceleration [14]. While the extensible software pattern of the library allows for an algorithm like L-BFGS to have multiple implementations for different architectures, the maintenance and tuning of a large number of implementations become a burden for the library developers, and users may not know how to choose the best implementation for their problem. It

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

3

is important, then, to provide a default implementation of L-BFGS that is portable (the same algorithm achieves close-to-optimal performance on different architectures) and versatile (different problems achieve close-to-optimal performance on the same architecture). With that in mind, we describe what kinds of performance, portability, and versatility matter when we assess L-BFGS algorithms. The time to update and solve is the most relevant performance metric. An L-BFGS approximation can be used by any optimization algorithm that needs an approximate Hessian, but algorithm 1 is the most common choice, so when comparing the performance of different implementations of L-BFGS, we focus on the sum of the time of lines 3 and 4, which we refer to together as the quasi-Newton update. The location and movement of vectors are the most important determinants of performance. All our algorithms are built from vector-vector and matrix-vector operations with low arithmetic intensity, so memory bandwidth and the latency of synchronization are critical. In PETSc/TAO there are three important cases for where the RN vectors reside, corresponding to different implementations of the Vec interface: 1. In the host memory of a single process (VECSEQ). No synchronization between memory domains is required, so here latency is negligible, and the primary concerns for L-BFGS are the sizes and bandwidths of the cache hierarchy and how well algorithms exploit locality. 2. In the device memory of a single GPU (VECSEQCUDA, VECSEQHIP). In addition to bandwidth and cache concerns, reduction operations across multiple threadblocks within GPU memory incur a synchronization latency that elementwise vector updates do not: this manifests as lower bandwidths for inner products at small vector sizes in streaming benchmarks such as [1, Figure 1] measuring performance on the NVIDIA A100. Potentially more significant, however, are latencies that are caused by the programming model. PETSc’s function VecDot() writes the value of an inner product to a scalar on the host: this introduces latency from the transfer of the value from device to host memory and from the pipeline stall caused by the host thread blocking before it launches subsequent kernels to the stream of computation. This source of latency can be avoided in other programming models that return a reference to the value in device memory, such as pytorch’s tensor.dot(). At the time of writing, a similar ManagedMemory interface is in development for PETSc/TAO [10], but that feature is not yet represented in the performance results in this work. 3. Distributed across multiple memory domains, whether host or device memory (VECMPI, VECMPICUDA, VECMPIHIP). Here the interconnect between memory domains is a significant source of latency in reductions and broadcasts regardless of the programming model. The base inverse Hessian H0 may be cheap or expensive, constant or variable. While it is common to choose H0 = I in the absence of better information, in physicsbased simulations H0 may be a relatively expensive preconditioner, such as a multigrid preconditioner for elasticity [5, 4]. In such cases the application of H0 takes the majority of the time in B.solve(), and extraneous applications of H0 should be avoided. Although the analysis of quasi-Newton methods is often simplified by the base matrix H0 being constant, in practice variable metric versions of quasi-Newton methods are used, where H0 is actually H0,k , which is adapted from H0,k−1 based on new data. The main implication for L-BFGS implementations is that products with H0,k cannot be reused after updates and must be recomputed. If the variability is only

4

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

T scalar variability (if, for instance, H0,k = (sTk−1 yk−1 )(yk−1 yk−1 )−1 I as suggested in [15]), then such quantities can be reused with the appropriate rescaling, but reuse is not possible for more general variable metric methods. Work comparing optimization algorithms from PETSc/TAO has demonstrated the superiority of diagonal variable metric methods over constant and scalar variable metric methods for a wide range of problems [7].

4. Storage of vectors. As described in section 2, L-BFGS uses sliding intervals of past vectors. To describe the intervals used in these algorithms, we define (4.1)

mk ≜ min{m, k},

(4.2)

ι(m, k) ≜ {k − mk , . . . , k − 1},

(4.3)

ω(m, k) ≜ {0, . . . , mk − 1},

and we use subscripts to indicate index sets, as in xι(m,k) ≜ {xi }i∈ι(m,k) . We will continue to use N as the size of the vectors. When vectors are spread across host processes using MPI, we use n for length of the largest subvector assigned to any one process (n is the more useful value for describing the time complexity of the algorithms we will present). An implementation of Hkm must store the latest vectors sι(m,k) and yι(m,k) , and some algorithms use other auxiliary vectors with matching numbering. For these vectors it is common to use arrays of vectors kept in history order : for example, an m array s̃ of m vectors with sι(m,k) stored in s̃ω(m,k) . When updating from Hk−1 to Hkm it is important to maintain history order by relabeling the vectors and not by copying so that the time complexity is O(n), as shown in algorithm 2, and not O(mn). (In algorithm 2 and subsequent algorithms, we annotate nontrivial operations with the corresponding BLAS-like routine and input sizes.) Algorithm 2: Store vector sk in s̃, history-order storage with history size m if k ≥ m then s̃temp = s̃0 3 for j ∈ {0, . . . , m − 2} in order do s̃j = s̃j+1 4 s̃m−1 = s̃temp

1

2

5

s̃min{m−1,k} ← sk

// reference, no copy // ” // ” // copy(n)

In section 6 we will present algorithms that format L-BFGS operations in terms of Level 2 BLAS and Level 3 BLAS instead of Level 1 BLAS (or equivalent cuBLAS or hipBLAS routines). Arrays of vectors are incompatible with these dense routines, so we must use matrices of column vectors for the history vectors: for example, S = [S0 | · · · |Sm−1 ] ∈ RN ×m to store the s vectors. It would be convenient if the column vectors of S were in history order (sι(m,k) stored in Sω(m,k) ), but maintaining this invariant would have O(mn) time complexity. Instead, we store all history vectors in cyclic order, with sk copied into Sk mod m , so that updating S has the same complexity as algorithm 2. An alternative to cyclic ordering would be to have extra columns in S. If S had m + r columns, then sι(m,k) could be stored in a sliding interval of columns, where the active columns would have to be copied back to the start of the matrix every r + 1 steps. This would give updating S an amortized complexity of O((1 + m−1 r+1 )n) per

5

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

step. The drawback of this approach is the extra spatial complexity, which is why we did not implement it. In PETSc/TAO, the column-vector matrices are instances of dense implementations of the Mat interface1 that use the same type of memory as the vectors; and when there are multiple MPI ranks, they are row-partitioned to match the vectors. 5. Storage of inner-product matrices. In addition to the history vectors, the algorithms in section 6 use the inner products between different sets of history vectors. Several, for example, use the values ri,j ≜ sTi yj for i, j ∈ ι(m, k). As with the basis vectors, we must store these values in a dense matrix R ∈ Rm×m to take advantage of the appropriate linear algebra routines. There are several design choices in how R is stored that impact performance. Host vs. device memory. R resides in the same type of memory as the RN vectors and column-vector matrices like S described in section 4. This means that products such as V ← SR can be computed without transferring R between the host and the device. It also means that computing a matrix of inner products such as R ← S T Y does not synchronize the host and the device. Multiple MPI ranks: duplicate vs. single storage. In PETSc/TAO, VecDot() computes a dot product using an allreduce operation so that the result is duplicated on every MPI rank, but MatTransposeMatMult() computes a matrix of inner products using a reduce operation and stores each inner product in just one location.2 L-BFGS algorithms could be implemented with matrices like R being duplicated on each MPI rank, or with R being stored on a single rank: the choice affects the communication complexity of products with column-vector matrices. • If R is duplicated on each process, R ← S T Y has the communication complexity of an allreduce of size m2 , and V ← SR is communication-free. • If R is stored once, R ← S T Y has the communication complexity of a reduce of size m2 and V ← SR the complexity of a bcast of size m2 . In PETSc/TAO, we choose to store R once on the first MPI rank because this allows us to use more of the existing Mat routines. We recognize, however, that in a selfconsistent MPI implementation T (reduce) + T (bcast) ≥ T (allreduce) [12], and for small message sizes many MPI implementations will have T (reduce) + T (bcast) ≈ 2T (allreduce). Cyclic ordering. Like matrices of column vectors, matrices of inner products in cyclic order (ri,j are stored in Ri mod m,j mod P m for i, j ∈ ι(m, k)). Because of the matching indexing, a product such as vi ← j∈ι(m,k) ri,j sj can be accomplished in a single gemm operation. The same P is not true of products that are triangular with respect to history order, like vi ← j∈ι(m,k), j≥i ri,j sj . We introduce htriuk (R) ∈ Rm×m to refer to the entries in R that are upper-triangular with respect to history order (see also Figure 1): given i, j ∈ ι(m, k), ( (5.1)

htriuk (R)i mod m,j mod m ≜

Ri mod m,j mod m , i ≥ j, 0 otherwise.

Matrix solves with htriuk (R) are accomplished by using three BLAS-like routines, as demonstrated in trsv cyclic (algorithm 3). 1 MATSEQDENSE, MATSEQDENSECUDA, MATSEQDENSEHIP, MATMPIDENSE, MATMPIDENSEHIP. 2 In all cases, GPU-aware MPI is used whenever possible [16].

MATMPIDENSECUDA,

6

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Sk mod m = sk−m 

         

···

···

new

old

        

 (1) trsv         

  (3) trsv        

(2) gemv

S(k−1) mod m = sk−1 S

htriuk (R)

Figure 1: Cyclic ordering. Left: cyclic ordering of the history vectors sι(m,k) in matrix S. Right: the history-order upper-triangular storage of ri,j = sTi yj innerproducts stored in the R matrix, annotated with the steps used to apply htriuk (R)−1 in algorithm 3.

Algorithm 3: trsv cyclic Data: R ∈ Rm×m , u ∈ Rm Result: u ← htriuk (R)−1 u 1 i ← k mod m 2 ιnew ← {0, . . . , i − 1} 3 ιold ← {i, . . . , m − 1} −1 4 uιnew ← triu(Rιnew ,ιnew ) uιnew 5 if k ≥ m then 6 uιold ← uιold − Rιold ,ιnew uιnew 7 uιold ← triu(Rιold ,ιold )−1 uιold

// trsv(k̂) // gemv(m − k̂, k̂) // trsv(m − k̂)

Algorithm 3 requires the same work as a single trsv or GPU-based equivalent, but when the data resides on a GPU algorithm 3 is implemented using three sequential kernel launches instead of one. We call this an inplace strategy. Because triangular operations appear frequently in L-BFGS algorithms, we also considered storing R and related matrices in history order, which we call the reorder strategy, to make triangular operations as fast as possible. In our performance measurements, this was never beneficial: the speedup of triangular operations was outweighed by the cost of maintaining history order. Consequently, all the data presented in this paper were collected with the inplace strategy. 6. L-BFGS algorithms. The dense storage of column-vector and inner-product matrices described in sections 4 and 5 can be used by multiple different algorithms to implement L-BFGS. We start by describing the common features of all of them. Table 1 shows that all implementations store the base Hessian B0 (which can be any SPD PETSc Mat, including a MATSHELL that wraps user-supplied code). If B0 = γI or B0 is diagonal, then H0 will be a direct inverse; otherwise H0 is the application of a preconditioner (PCApply()) or an iterative method (KSPSolve()).

7

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

Table 1: Common data structures across all L-BFGS implementations m k B0 xprev g prev S Y d

history size, constant number of completed updates, initially 0 SPD RN ×N matrix RN vector, the latest solution value, initially x0 RN vector, the latest gradient, initially g0 RN ×m cyclic-order column-vector matrix [S0 | · · · |Sm−1 ] that stores sι(m,k) RN ×m cyclic-order column-vector matrix [Y0 | · · · |Ym−1 ] that stores yι(m,k) Rm cyclic-order vector of inner products, di = SiT Yi

Each implementation stores the last x and g vectors used to update B and the sι(m,k) and yι(m,k) vectors. Algorithm 4 shows that all implementations use a curvature condition sTk yk ≥ ϵtol ∥yk ∥2 to ensure a valid update, so the sTk yk values are stored in a vector d for reuse. Algorithm 4 also shows how we use · · · = B.(. . . ) to indicate which variables in the algorithm are references to data structures of the implementation and how we annotate steps with parallel communication with MPI-like routines and sizes. Algorithm 4: B.update(x, g) m, k, B0 , xprev , g prev , S, Y, d = B.(. . . ) if B0 is variable metric then B0 .update(x, g) prev 3 s←x−x // axpy(n) prev 4 y ←g−g // axpy(n)    T δ s 5 ← T y // 2 dot(n), allreduce(2) η y 6 if δ > ϵtol η then 7 i ← k mod m 8 Si ← s // copy(n) 9 Yi ← y // copy(n) 10 di ← δ 11 k ←k+1 12 B.update solve(x, g) 13 if B.mult is used then B.update mult(x, g) 1 2

xprev ← x prev 15 g ←g 14

// copy(n) // copy(n)

We note that in this scheme the data structures used by B.mult are not eagerly updated to enhance the performance of algorithms like algorithm 1 that do not use it. The cost of B.update, exclusive of B.update solve and B.update mult, is (6.1) T (B.update) = 2T (axpy(n)) + 2T (dot(n)) + 4T (copy(n)) + T (allreduce(2)) + O(1). 6.1. The recursive approach. We first describe the best performance that can be achieved with recursive algorithms that use only Level 1 BLAS. 6.1.1. Recursive B.solve. An approach to B.solve that directly translates (2.4) into a recursive algorithm has two main advantages: no additional data structures

8

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

are needed, so B.update solve is a no-op, and the algorithm does not change if B0 varies from iteration to iteration. Algorithm 5: Recursive::B.solve(v, z) m, k, B0 , S, Y, d = B.(. . . ) a ∈ Rmk 3 z ←v 4 for j ∈ ι(m, k) in reverse order do 5 i ← j mod m 6 ai ← (SiT z)/di 7 z ← z − ai yi

1 2

z ← B0−1 z 9 for j ∈ ι(m, k) in order do 10 i ← j mod m 11 β ← (YiT z)/di 12 z ← z + (ai − β)Si

// temporary work vector // copy(n)

// dot(n), allreduce(1) // axpy(n) // B0 .solve()

8

// dot(n), allreduce(1) // axpy(n)

Algorithm 5 shows that the main disadvantage of this approach is the for-loops that cannot be unrolled because of data dependencies, which results in poor temporal locality of memory access and reduction operations that cannot be coalesced. The cost of Recursive::B.solve is (6.2) T (Recursive :: B.solve) = T (B0 .solve) + T (copy(n)) + 2mk {T (axpy(n)) + T (dot(n)) + T (allreduce(1))} + O(mk ). 6.1.2. Recursive B.mult. Equation (2.5) shows that Bkm can be expressed nonrecursively as a rank-2mk update of B0 : this form is good for performance because the updates can be performed in parallel. In particular, this means that 2mk separate reductions can be combined into a single larger reduction, as described in algorithm 6, where the column-vector matrix P stores the p vectors (2.6) in cyclic order.3 Algorithm 6: Recursive::B.mult(v, z) m, k, B0 , S, Y, d, P, f = B.(. . . ) a, b ∈ Rmk 3 ω ← ω(m, k) 4 z ← B0 v    T a Y 5 ← ωT v b P    ω −1   a diag(dω ) 0 a 6 ← b 0 diag(f ) b   ω   a −Pω 7 z ← z + Yω b 1 2

// temporary work vectors // active history indices // B0 .mult() // 2mk dot(n), allreduce(2mk ) // 2 diags(mk ) // 2mk axpy(n)

3 Algorithm 6 shows that we use diagm and diags to annotate the cost for elementwise multiplication

and division, respectively, even though these are not BLAS routines.

9

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

Table 2: Data structures for B.mult for the recursive approach P̊ P f

RN ×m matrix to store B0 S RN ×m cyclic-order column-vector matrix [P0 | · · · |Pm−1 ] that stores {pm−(k−i),i }i∈ι(m,k) (2.6) Rm cyclic-order vector of inner products, fi = PiT Si

The cost of this algorithm is (6.3) T (Recursive :: B.mult) = T (B0 .mult) + 2mk {T (axpy(n)) + T (dot(n))} + T (allreduce(2mk )) + O(mk ). The main drawback is that the vectors {pm−(k−i),i }i∈ι(m,k) are defined recursively and must be computed in B.update mult. These vectors also have limited reuse: after m updating from Bk−1 to Bkm , the p vectors cannot be reused if k > m (because dropping the oldest p vector from the recursion changes all subsequent vectors), and they cannot be reused if H0 has changed. To handle these two cases that invalidate the p vectors separately, a recursive L-BFGS implementation can store (Table 2) the p vectors as well as the auxiliary vectors p̊ι(m,k) ≜ {B0 si }i∈ι(m,k) ,

(6.4)

which need to be updated only if B0 has changed, and which do not need to be computed at all if B0 is a scaled identity matrix. We give Recursive::update mult as algorithm 11 in the appendix. We can summarize its performance when k ≥ m as (6.5) T (Recursive :: B.update mult) = mupdate T (B0 .mult) +mT (axpy(n))+m2 T (dot(n))+(m2 −m)T (axpy(n))+(2m−1)T (allreduce(O(m))) + O(m2 ), where (6.6)

  0, mupdate = 1,   m,

B0 = B0,k = γk I, B0 is constant, otherwise.

Regardless of the value of mupdate , Recursive::B.update mult has O(m2 n) time complexity just in its BLAS routines, as opposed to O(mn) for Recursive::B.solve. Thus it is important to avoid eagerly calling this routine in B.update. 6.2. The compact dense approach. Byrd, Nocedal, and Schnabel [6] presented a compact dense representation of Hkm and developed an approach to using it in algorithm 1 in a way that is work-equivalent to using the recursive approach from the preceding section. Here we present this approach, summarizing their analysis for the case B0 = γI and analyzing the case when B0 is a general variable-metric matrix. We define a new set of vectors qι (m, k) ≜ {H0 yi }i∈ι(m,k) and let Q ∈ RN ×m be a cyclic-order column-vector matrix for storing them. We define the inner-product matrices as (6.7)

R ≜ htriuk (S T Y ),

D ≜ diag(S T Y ),

Z ≜ QT Y = Y T H0 Y.

10

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Table 3: Data structures for B.solve for the compact dense approach R Q Z acache bcache

Rm×m matrix to store htriuk (S T Y ) RN ×m matrix to store H0 Y Rm×m matrix to store Y T H0 Y Rm vector to cache the product S T g prev Rm vector to cache the product QT g prev

With these matrices, the compact dense representation of Hkm is (6.8)

 Hkm = H0 + S

  R−T (Z + D)R−1 Q −R−1

−R−T 0

 ST . QT



6.2.1. Compact dense B.solve. The inner matrix in the product in (6.8) can be applied to a vector in a handful of BLAS operations (including trsv cyclic (algorithm 3) and its transpose) using O(m2 ) work. The majority of the work in B.solve is in the products with S, Q, and their transposes, (6.9) T (CompactDense :: B.solve) = T (B0 .solve()) + 4T (gemv(n, mk )) + T (reduce(2mk )) + T (bcast(2mk )) + O(m2k ), where reductions and broadcasts can be coalesced because they occur in parallel. In B.update solve the Q, R, and Z matrices must be updated. Like the P̊ matrix in the recursive approach, updating Q can be avoided if H0 = γI. R can be updated with a single gemv(n, mk − 1) between yk and the older columns of S (the diagonal value sTk yk was computed in B.update). Because Z is symmetric, it can also be updated with a single gemv(n, mk ) between yk and the active columns of Q. If we sum the BLAS work just described in updating and solving using the compact dense approach (using ∼ 2mn flops for gemv(n, m) and ∼ 2n flops for dot(n)), we get ∼ (12mk − 2)n flops, more than the ∼ 8mk n flops for B.solve in the recursive approach. Byrd et al. solved this problem by recognizing that in the quasi-Newton update in algorithm 1 the same gradient vector gk is an argument to both B.update and B.solve. If the S T gk and QT gk products are cached during B.update(xk , gk ) (Table 3), then they do not need to be recomputed in B.solve(gk , p), as shown in algorithm 7. When we are given gk+1 in the next update, S T yk = S T gk+1 − S T gk can be computed when S T gk+1 is computed with only O(m) additional work. The details of how this is accomplished are in algorithm 13 in the appendix. Here we summarize the performance of B.update solve in isolation when k ≥ m as (6.10) T (CompactDense :: update solve) = mupdate T (B0 .solve) + T (gemm(mupdate , mupdate , n)) + 2T (gemv(n, m)) + T (dot(n)) + T (reduce(m2update )) + 2T (reduce(m)) + T (reduce(1)) + O(m2 ), where mupdate is the same as (6.6). When the S T gk and QT gk products are cached, the complexity of the quasi-Newton update using algorithms 7 and 13 becomes, for

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

11

Algorithm 7: CompactDense::B.solve(v, z) m, k, B0 , S, Y, d, Q, R, Z, g prev , acache , bcache = B.(. . . ) m 2 a, b, c, e ∈ R k // temporary work vectors 3 ω ← ω(m, k) // active history indices −1 4 (α, Q̃) = if B0 = βI then (β , Y ) else (1, Q) // αQ̃ = Q −1 5 z ← B0 v // B0 .solve() prev 6 if v ̸= g     then SωT a 7 v // 2 gemv(n, mk ), reduce(2mk ) ← b αQ̃Tω 1

else (a, b) ← (acache , bcache ) ω ω −1 9 a ← −htriuk (Rω,ω ) a      c Z + diag(dω ) I a 10 ← ω,ω e I 0 b −T 11 c ← −htriuk (Rω,ω ) c     c 12 z ← z + Sω αQ̃ω e 8

// 2 copy(mk ) // trsv cyclic(mk ), scal(mk ) // gemv(mk , mk ), diagm(mk ), axpy(mk ) // trsv cyclic(mk ), scal(mk ) // 2 gemv(n, mk ), bcast(2mk )

k ≥ m, (6.11) T (CompactDense :: update solve(xk , gk ) + CompactDense :: solve(gk , p)) = (1 + mupdate )T (B0 .solve) + T (gemm(mupdate , mupdate , n)) + 4T (gemv(n, m) + T (dot(n)) + T (reduce(m2update )) + T (reduce(2m)) + T (bcast(2m)) + T (reduce(1)) + O(m2 ). When we compare this to T (Recursive :: B.solve) (6.2), we see the majority of the BLAS operations are now in Levels 2 and 3 BLAS, and we also see that the number of parallel reductions and broadcasts has been reduced from O(m) to O(1), which should translate to superior machine performance. Whether the compact dense approach is superior, however, depends on B0 and mupdate . If B0 = γI, then mupdate = 0, and the BLAS work in (6.11) is ∼ (8m + 2)n flops. Byrd et al. further optimized this algorithm by reusing the cached inner products to check the curvature condition in algorithm 4 and the stopping conditions for the algorithm, so that the BLAS work in the compact dense approach can be further reduced to the optimal ∼ 8mn flops in this case. If B0 is constant but not diagonal, then mupdate = 1, and the BLAS work is ∼ (8m + 4)n, which is still nearly optimal. However, the number of applications of B0 .solve is now 2 instead of 1, so the compact dense approach will be worse if B0 .solve is a relatively expensive operation. If B0 is variable-metric, then mupdate = m, and the Q and Z matrices must be recomputed at each iteration, leading to m applications of B0 .solve and ∼ (2m2 + 8m + 2)n flops of BLAS work. For large m the compact dense approach could require so much extra work that its superior machine performance does not make it faster than the recursive approach. 6.2.2. Compact dense B.mult. The Sherman–Morrison–Woodbury formula can be used on (6.8) to compute a compact dense representation of Bkm . Using P̊ = B0 S

12

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Table 4: Data structures for B.mult for the compact dense approach P̊ F̊ L J

RN ×m matrix to store B0 S Rm×m matrix to store P̊ T S Rm×m matrix to store hstrilk (S T Y ) Rm×m matrix to store the Cholesky factor JJ T = F̊ + L diag(d)−1 LT

as previously defined and letting F̊ ≜ P̊ T S = S T B0 S, we have (6.12)

Bkm = B0 −

 Y

  −D P̊ L

LT F̊

−1  T  Y , P̊ T

where L is the history-order strictly lower-triangular portion of S T Y , L = hstrilk (S T Y ) ≜ S T Y − htriuk (S T Y ). The Cholesky factorization JJ T ≜ F̊ + LDLT is guaranteed to exist [6, Theorem 2.4], which allows us to write (6.12) as         I D−1 LT I 0 I 0 D−1 0 −Y T m (6.13) Bk = B0 − Y P̊ . 0 I 0 I 0 J −T J −1 −L I P̊ T If these matrices are stored (Table 4) and updated in B.update mult, then (6.13) shows how to efficiently implement this representation of Bkm in B.mult (algorithm 8). Algorithm 8: CompactDense::B.mult(v, z) m, k, B0 , S, Y, d, P̊ , L, J = B.(. . . ) a, b ∈ Rmk // temporary work vectors 3 ω ← ω(m, k) // active history indices 4 (α, P̃ ) = if B0 = βI then (β, S) else (1, P̊ ) // αP̃ = P̊ 5 z ← B0 v // B0 .mult()     −YωT a 6 ← v // 2 gemv(n, mk ), reduce(2mk ) b αP̃ωT −1 7 a ← diag(dω ) a // diags(mk ) −T −1 8 b ← Jω,ω Jω,ω (b − Lω,ω a) // gemv(mk , mk ), 2 trsv(mk ) −1 T 9 a ← a + diag(dω ) Lω,ω b // gemv(mk , mk ), diags(mk )     a 10 z ← z − Yω // 2 gemv(n, mk ), bcast(2mk ) αP̃ω b 1

2

The time complexity of algorithm 8 is (6.14) T (CompactDense :: B.mult) = T (B0 .mult) + 4T (gemv(n, mk )) + T (reduce(2mk )) + T (bcast(2mk )) + O(m2k ). Equation (6.12) is a rank-2m update; but unlike the p vectors in the rank-2m update defining Bkm (2.5), the p̊ vectors are not recursively defined. The cost of updating P̊ was discussed in subsection 6.1. The cost of updating F̊ is dependent on

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

13

Table 5: Data structures for B.solve for the intermediate dense approach R cache

a

Rm×m matrix to store htriuk (S T Y ) Rm vector to cache the product S T g prev

the variability of B0 in the same way as the Z matrix in Table 3, and only one row of L needs to be updated with mk − 1 new inner products per iteration. Once these matrices are computed, the J matrix can be computed in O(m3k ) work. Therefore the cost of B.update mult, given in the appendix as algorithm 14, is for k ≥ m, (6.15) T (CompactDense :: update mult) = mupdate T (B0 .mult) + T (gemv(n, m − 1)) + T (reduce(m − 1)) + T (gemm(max{1, mupdate }, m, n)) + T (reduce(max{1, mupdate }m)) + O(m3 ). If B0 = γI or if B0 is constant, then this is superior to the complexity of Recursive::B.update mult (6.5) because it requires only O(mn) BLAS work as opposed to O(m2 n). If B0 is a general variable-metric matrix, however, the compact dense approach also requires O(m2 n) BLAS work. 6.3. Intermediate dense approach. If we multiply out the blocks in the compact dense representation (6.8), we see (6.16)

Hkm = H0 + SR−T (Y T H0 Y + D)R−1 S T − SR−T QT − QR−1 S T = (I − SR−T Y T )H0 (I − Y R−1 S T ) + SR−T DR−1 S T ,

which looks like the sum of an oblique projection of H0 plus a correction, similar to the original definition in (2.1). In this form H0 appears only once, and there is no need for the Q = H0 Y matrix of column vectors nor the Z = Y T H0 Y inner product matrix. An implementation of L-BFGS that does not require those data structures avoids the drawbacks of the compact dense approach discussed in subsection 6.2. 6.3.1. Intermediate dense B.solve. The last line in (6.16) can be reorganized into a form we call the intermediate dense form,        I 0 H0 0 I −Y I (6.17) Hkm = I SR−T , 0 D 0 I −Y T I R−1 S T which shows the steps involved in computing z ← Hkm v. This approach only requires that the R = htriuk (S T Y ) matrix be kept up to date (Table 5); but in order to optimize for the quasi-Newton update, the implementation should cache the S T gk product in B.update mult for reuse in B.mult, using the same caching strategy as the compact dense approach in subsection 6.2. The result is algorithm 9. The complexity of algorithm 9 applied to a generic vector v is (6.18) T (IntermediateDense :: B.solve) = T (B0 .solve) + 4T (gemv(n, mk )) + 2T (reduce(mk )) + 2T (bcast(mk )) + O(m2k ). The complexity of B.update solve (algorithm 15 in the appendix) is T (IntermediateDense :: B.update solve) = T (gemv(n, mk )) + T (dot(n)) + T (reduce(n, mk )) + T (reduce(1)) + O(mk ).

14

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Algorithm 9: IntermediateDense::B.solve(v, z) m, k, B0 , S, Y, d, R, g prev , acache = B.(. . . ) m 2 a∈R k 3 ω ← ω(m, k) prev 4 if v ̸= g then 5 a ← SωT v 1

a ← acache ω

else a ← htriuk (Rω,ω )−1 a 8 z ← v − Yω a −1 9 z ← B0 z 10 a ← diag(dω )a T 11 a ← a − Yω z −T 12 a ← htriuk (Rω,ω ) a 13 z ← z + Sω a 6

// temporary work vector // active history indices // gemv(n, mk ), reduce(mk ) // copy(mk ) // trsv cyclic(mk ) // gemv(n, mk ), bcast(mk ) // B0 .solve() // diagm(mk ) // gemv(n, mk ), reduce(mk ) // trsv cyclic(mk ) // gemv(n, mk ), bcast(mk )

7

When the S T gk product is cached, the cost of the quasi-Newton update is T (IntermediateDense :: B.solve(xk , gk ) + IntermediateDense :: B.update solve(gk , p)) = T (B0 .solve) + 4T (gemv(n, mk )) + T (dot(n)) + 2T (reduce(mk )) + 2T (bcast(mk )) + T (reduce(1)) + O(m2k ). Compared with the cost of T (Recursive :: B.solve) (6.2), this uses only ∼ 2n additional flops, and the majority of the BLAS work is in Level 2 BLAS, like the compact dense approach. The reductions and broadcasts are not as coalesced as in the compact dense approach (6.11), because the intermediate formulation is essentially two sequential rank-m updates, as opposed to one rank-2m update, but the performance is not affected by an expensive or variable B0 . 6.3.2. Non-existence of an intermediate dense B.mult. There cannot be an analog to the intermediate dense form (6.17) for Bkm that avoids the drawbacks of the compact dense approach to implementing B.mult. By an analog we mean an algorithm to compute z ← Bkm x exactly that uses only O(1) applications of B0 .mult and B0 .solve for every choice of B0 , S, and Y . Consider the case where S = Y = U is a set of orthonormal vectors. Then (6.12) simplifies to  −1  T    −I 0 U m Bk = B0 − U B0 U 0 U T B0 U U T B0 = B0 + U U T − B0 U (U T B0 U )−1 U T B0 . If we then consider the product U T H0 Bkm H0 U , we see   U T H0 Bkm H0 U = U T H0 B0 + U U T − B0 U (U T B0 U )−1 U T B0 H0 U = U T H0 U + (U T H0 U )2 − (U T B0 U )−1 , which when rearranged yields (U T B0 U )−1 = U T H0 U + (U T H0 U )2 − U T H0 Bkm H0 U.

15

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

Table 6: Quasi-Newton update performance (B.update(xk , gk ) + B.solve(gk , p)). In the reduce and bcast columns, an entry “a (b)” denotes a collective operations that coalesce into b synchronizations (one latency each); b is omitted when a = b. B0−1

BLAS memops

MPI synchronizations

Lvl. 1

Lvl. 2

Lvl. 3

allr.

reduce

bcast

B.update

0

∼ 18n

0

0

1

0

0

Recursive

1

∼ 10mn

0

0

2m

0

0

CompactDense B0,k = γk I B0 constant B0 = B0,k var.

1 2 1+m

∼ 4n ∼ 4n ∼ 2n

∼ 4mn ∼ 4mn ∼ 4mn

0 ∼ 2mn

0 0 0

4 (1) 4 (1) 4 (1)

2 (1) 2 (1) 2 (1)

Intermediate

1

∼ 2n

∼ 4mn

0

0

3 (2)

2

The existence of an algorithm with the desired properties would thus imply the existence of an algorithm that exactly computes z ← (U T B0 U )−1 x using O(1) applications of B0 .mult and B0 .solve for every SPD B0 and every set of orthonormal vectors U . The O(1) limit allows neither the formation of (U T B0 U ) for factorization nor for the solution via an iterative method that could take up to m iterations, so this is generally impossible. Because of this negative result, CompactDense::B.mult (algorithm 8) is the only dense algorithm for B.mult that we consider. 7. Quasi-Newton update performance comparison. 7.1. Performance model. Throughout we analyze the limited-memory BFGS (L-BFGS) method at an iteration k ≥ m, so that all the history vectors (sk−m , yk−m ), . . . , (sk−1 , yk−1 ) are present. The L-BFGS inverse Hessian Hk is a rank-2m update of H0 , so there exists a matrix Uk ∈ Rn×2m such that Hk = H0 + Uk UkT . If we did not count the setup time to compute Uk , this representation would always be the most efficient for computing p ← Hk g, because the update term Uk UkT g is computed in only two dense matrix-vector products and requires reading only one additional matrix. With this representation in mind, we define the effective bandwidth of an L-BFGS implementation as (7.1)

Be ≜

wn(2m + 2) . Tupdate + Tsolve

7.2. Performance measurements. Compute environments. In this paper we present CPU- and GPU-based results on OLCF’s Frontier and ALCF’s Polaris. Frontier is an HPE Cray EX machine. Each node is equipped with one 64-core AMD EPYC 7A53 (Trento) CPU and 4 AMD Instinct MI250X cards, each with 2 Graphics Compute Dies (GCDs) for a total of 8 GCDs per node. The CPUs and GPUs are connected via Infinity Fabric. Polaris is an HPE Apollo 6500 machine, equipped with 32 core AMD EPYC 7543P Milan, with 4 NVIDIA A100 GPU cards, per node. The CPUs and GPUs are connected via PCIe-4.0, and the GPUs are connected via NVLink. For Polaris, both single- and four-GPU regimes were tested. However, for Frontier, despite each node offering up to eight GCDs (usable as eight separate GPUs), this paper only investigates code

16

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Table 7: Environments for each machine machine

modules and environment

Frontier

PrgEnv-gnu/8.6.0, rocm/6.3.1, craype-x86-trento, cray-mpich/8.1.31, craype-accel-amd-gfx90a PrgEnv-nvidia/8.6.0, nvidia/24.11, craype-x86-milan

Polaris

efficiency on a single GCD. Table 7 describes the modules and environments used on each machine. Solve and update step performance. We measure the timings of update and solve when the right-hand side is gk , which we call Tstep . The goal is to compare the implementation to the performance model in (7.1). The program that generated these timings is included in the PETSc test suite.4 The value Tstep is the average iteration time of a loop that repeatedly calls Update(xk , gk ) followed by Solve(gk ) (but without a linesearch or even any object evaluation).The loop only includes steps for which k ≥ m so that the full history size is used, because the initial steps for k < m will have different performance characteristics, making it harder to compare the measurements with the performance model. For performance testing, we used two types of H0 : a constant diagonal matrix, αI, and a diagonal matrix. For scaling H0 , we used two types: no scaling and diagonal. Diagonal scaling uses a full-memory restricted Broyden update formula to construct a diagonal matrix for H0 initialization [11]. In all, we conducted tests on two possible combinations: constant diagonal matrix without scaling and diagonal matrix with diagonal scaling. The former represents the simplest strategy for H0 , while the latter represents the most complex case, with which we illustrate shortfalls of the compact dense method and superiority of our intermediate dense strategy. Mult and update step performance. We perform the same test for Mult. We use the same setup as Solve, including the PETSc test suite, loop structure, and types of H0 matrices. We note that Mult timings vary markedly between the Recursive and compact dense methods. We therefore show two representative cases: 4 MPI ranks on CPUs and 4 MPI ranks on GPUs. 7.3. CPU. In this subsection we present performance measurements for the CPU case. We have constructed time/efficiency plots that compare Tstep for different L-BFGS algorithms at different problem sizes n and different history sizes m with different strategies for H0 . Efficiency is measured by using the effective bandwidth Be . Solid lines connect different measurements with the same problem size n at history sizes m ∈ {5, 10, 20, 50}; the problem size n is annotated, and the history size m is indicated with the marker size. We note that all CPU performance measurements were conducted on ALCF’s Polaris. We present four plots measuring Solve() and one plot measuring Mult() done on the CPU. We have MPI rank counts of 1 and 4; and, for each given MPI rank, we present both constant diagonal H0 without scaling and diagonal H0 with diagonal scaling. We choose 4 MPI ranks because L3 cache can become saturated with the problem size at this rank count. For Mult(), we present one plot with 4 MPI ranks, 4 ksp ksp utils lmvm tests-solve performance 0.

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

17

with constant diagonal H0 without scaling. Be (GB/s) 30 1 B 2 S

104

20

105

103

10

1 B 4 S

106

method recursive comp. dense dense history size m 5 10 20 50

107 n = 102

0

10−5

10−4

10−3

10−2

10−1

100

Tstep (s)

Figure 2: BFGS Solve comparison, constant diagonal H0 without scaling. Be is the effective bandwidth, as defined in (7.1). Polaris (EPYC 7543P), 1 MPI CPU process

In Figure 2, one can observe that at n = 102 the problem size is so small that the latency dominates. The efficiency of the recursive algorithm is nearly constant, while both dense algorithms exhibit nearly constant runtime. At n = 104 , the latency is no longer dominant, and all the problem data fits into the 32 MiB L3 cache for each of the history sizes tested. At this problem size the 256 KiB L2 cache is only ≈ 3× the size of a single vector. We speculate that the intermediate dense algorithm performs better than the recursive algorithm for larger values of m because it makes better use of this cache. At n = 106 , the 32 MiB L3 cache is only ≈ 4× bigger than a single vector. For the recursive algorithm, if its rank-1 updates keep the vectors being modified in cache, it is predicted that the recursive algorithm will have the same effective bandwidth as the intermediate dense algorithm, which matches our measurements. At n = 107 a single vector does not fit in the L3 cache, so we predict that the recursive algorithm will behave poorly compared with the dense algorithm, as the rank-1 update no longer keeps the vectors being modified in cache, which also matches our measurements. We measured the bandwidth to be approximately 43 GB/s for a single core. One can observe similar behavior for 4 MPI ranks. In Figures 3 and 5, we note that, for n = 106 , the recursive algorithm very slightly outperforms both dense algorithms. We observed that when n = 106 , Update() was slightly more expensive for dense algorithms, possibly due to extra memory movement to assemble appropriate internal objects, while such extra cost can be avoided for recursive algorithms. Regardless, we still observe that the dense algorithm outperforms the recursive algorithm for the largest problem size of n = 107 . When H0 is not constant, the compact-dense formulation performs substantially worse. For a diagonal H0 with diagonal scaling, Figures 4 and 5 show that the compact dense algorithm is the slowest of the three across all problem and history

18

1

B 2 S H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

80 Be (GB/s) 60 105

1 B 4 S

40 104 n = 106

method recursive comp. dense dense history size m 5 10 20 50

20 103 107

0

10−4

10−3

10−2

10−1

100 Tstep (s)

Figure 3: BFGS Solve comparison, constant diagonal H0 without scaling. Polaris (EPYC 7543P), 4 MPI CPU processes

Be (GB/s) 25 1 B 2 S

20

n = 104 105 106

15

107

3

10

1 B 4 S

10

method recursive comp. dense dense history size m 5 10 20 50

5 102

0

10−5

10−4

10−3

10−2

10−1

100

Tstep (s)

Figure 4: BFGS Solve comparison, diagonal H0 with diagonal scaling. Polaris (EPYC 7543P), 1 MPI CPU process

sizes, for both 1 and 4 MPI processes. This behavior is expected: as summarized in Table 6, a variable H0 forces the compact dense method to perform m additional evaluations of H0 and ∼ (2m2 + 8m + 2)n flops of BLAS work. By contrast, the intermediate-dense formulation performs much better than the compact-dense one and typically outperforms the recursive algorithm. The exceptions occur with 4 MPI processes for n = 105 with m = 50 and for n = 106 , where the recursive method is slightly faster. We hypothesize that storing the R matrix on the first MPI rank (reduce and bcast), rather than duplicating it on each process (allreduce),

1 B 2 S

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

60 Be (GB/s) 1 B 4 S

40 104

105 n = 106

20

107

19

method recursive comp. dense dense history size m 5 10 20 50

103

0

10−4

10−3

10−2

10−1

100

Tstep (s)

Figure 5: BFGS Solve comparison, diagonal H0 with diagonal scaling. Polaris (EPYC 7543P), 4 MPI CPU processes

introduces overhead that slows the intermediate dense method in these cases. Even so, for the largest problem size, n = 107 , the intermediate dense formulation remains faster than the recursive algorithm. We also discuss Mult() timings. We present Mult() timings with diagonal H0 without scaling in Figure 6, with 4 CPU MPI ranks. We note that for all cases, dense versions significantly outperform the recursive version, as expected. We note that the speedup is much more pronounced for higher history size m, with almost an order of magnitude speedup for m = 50 for n > 103 . This speedup is expected because this problem type has O(mn) BLAS work, compared with O(m2 n) BLAS work for the recursive method. 7.4. GPU. In this subsection we present performance measurements for the GPU; notation and plotting conventions are identical to Subsection 7.3, except for using circle markers instead of square markers. We present four plots measuring Solve() timings and one plot measuring Mult() timings on the GPUs of ALCF’s Polaris machine. As in the CPU experiments, for each GPU MPI rank, we consider two setups: constant diagonal H0 without scaling and diagonal H0 with diagonal scaling. In addition, we present one plot measuring Solve() conducted on OLCF’s Frontier, with constant diagonal H0 without scaling. In Figure 7 one can observe that for all problem sizes, dense algorithms vastly outperform the recursive algorithm, as expected. Since H0 is constant with no scaling, the compact dense and intermediate dense algorithms’ performances are similar, although intermediate dense is slightly faster at smaller problem sizes. The recursive algorithm exhibits an almost constant effective bandwidth across m, for each problem size, consistent with its low arithmetic intensity: runtime is bounded by memory bandwidth and synchronization latency rather than compute. We find that for constant diagonal H0 without scaling, the highest speedup factor of 7.35 is achieved, with n = 105 with m = 50, and the speedup factor for the biggest problem size is 2.36 for n = 107 with m = 50.

20

1 B 4 S

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Be (GB/s)

method Recursive Dense history size m 5 10 20 50

30 104

105

20 106

10 n = 107

103

0

10−4

10−3

10−2

10−1

100

101

Tstep (s)

Figure 6: BFGS Mult comparison, diagonal H0 without scaling, Polaris (EPYC 7543P), 4 MPI CPU processes 1 B 2 S 1 B 2 S

600 Be (GB/s) 400 7

n = 10 106

200

1 B 4 S

method recursive comp. dense dense history size m 5 10 20 50

105

0

10−3

10−2

Tstep (s)

Figure 7: BFGS Solve comparison, constant diagonal H0 without scaling. Polaris A100, 1 GPU

Similar to the CPU-based results, we observe the compact dense approach exhibiting much worse performance when H0 is not constant. This is not surprising, because different architecture does not change the fundamental additional algorithmic overhead of having to recompute over all the basis vectors. The compact dense algorithm, which performed much worse than the recursive algorithm on the CPU, performs on-par with the recursive method on GPUs, even outperforming it for the small problem size of n = 105 . We suspect that, despite its higher computational burden, the compact dense algorithm with variable H0 incurs lower synchronization latency and places less pressure on memory bandwidth, which may explain its slightly better performance

1 B 2 S

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

500 Be (GB/s) 400

1 B 4 S

n = 106

300

107

105

200

21

method recursive comp. dense dense history size m 5 10 20 50

100 104

0

10−3

10−2

10−1

1 B 2 S 1 B 2 0S

Tstep (s)

Figure 8: BFGS Solve comparison, constant diagonal H without scaling. Frontier MI250X, 1 GPU

500 Be (GB/s) 400

1 B 4 S

300 106

n = 107

200 100 0

method recursive comp. dense dense history size m 5 10 20 50

105

10−3

10−2

Tstep (s)

Figure 9: BFGS Solve comparison, diagonal H0 with diagonal scaling. Polaris A100, 1 GPU

relative to the recursive algorithm. We find that for variable H0 with diagonal scaling, the highest speedup factor of 4.73 is achieved, with n = 105 with history size m = 50, and the speedup factor for largest problem size is 2.16 for n = 107 with m = 50. One can observe similar behavior for 4 GPUs. In Figures 10 and 11, we note that the intermediate dense algorithm greatly outperforms the recursive algorithm. For constant diagonal H0 without scaling, intermediate dense and compact dense algorithms show similar results, as also observed in CPU-based results and results from one GPU card. But for diagonal H0 with diagonal scaling, while the intermediate dense algorithm still greatly outperforms the recursive algorithm, the compact dense

22

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS 1 B 2 S

2,500 Be (GB/s) 2,000 1,500

108

1 B 4 S

n = 107

1,000 500 0

method recursive comp. dense dense history size m 5 10 20 50

106

10−3

10−2

10−1 Tstep (s)

Figure 10: BFGS Solve comparison, constant diagonal H0 without scaling. Polaris 1 B A100, 4 GPUs 2 S

Be (GB/s) 2,000

1 B 4 S

1,500 108

1,000

n = 107

method recursive comp. dense dense history size m 5 10 20 50

500 106

0

10−3

10−2

10−1 Tstep (s)

Figure 11: BFGS Solve comparison, diagonal H0 with diagonal scaling. Polaris A100, 4 GPUs

algorithm performs slightly worse than the recursive one. For the Mult() timings in Figure 12, we observe that the dense method outperforms the recursive method by roughly an order of magnitude as the history size m increases. As noted in Subsection 7.3, this gap is expected because the recursive method requires O(m2 n) BLAS work for update, compared with O(mn) work for the dense method. 8. Performance on a convolutional neural network training benchmark. In this section we evaluate the performance of the dense BFGS solver against the baseline recursive BFGS implementation on a supervised learning problem: training a

1 B 4 S

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

1,000 Be (GB/s)

23

method Recursive Dense history size m 5 10 20 50

800 600 400

n = 107

200 106

0

10−2

10−1

Tstep (s)

Figure 12: BFGS Mult comparison, diagonal H0 without scaling. Polaris A100, 4 GPUs

convolutional neural network on the MNIST image classification benchmark. We integrated PETSc/TAO’s BFGS solvers with PyTorch, enabling their use for training a neural network (NN). The MNIST training set contains 60,000 grayscale images, each of size 28 × 28 pixels, labeled across ten classes. The objective function minimizes the cross-entropy loss between the model predictions and the ground-truth labels. We use a simple CNN with two convolutional layers followed by two fully connected layers. It contains 1,199,882 trainable parameters, resulting in a BFGS Hessian matrix of size 1,199,882 × 1,199,882. Training was performed with a mini-batch size of 1024, and the maximum number of BFGS iterations per mini-batch was set to 20. Table 8 reports the solver times for different history sizes. The solver time includes all operations performed by TAO except for the loss function evaluation (forward propagation) and the gradient computation (backpropagation). As shown in Table 8, the dense BFGS solver consistently outperforms the recursive implementation across all history sizes, achieving up to a 4.7% speedup.

Table 8: Comparison of recursive and dense BFGS under different history sizes on the MNIST benchmark. Values are mean±std of per epoch time cost (in seconds) over three repeated runs. Improvement = (Meanrecursive − Meandense )/Meanrecursive . History Size

Recursive BFGS

Dense BFGS

Improvement

5 10 20

8.882±0.11 9.111±0.09 9.272±0.09

8.804±0.10 8.825±0.09 8.838±0.07

0.9% 3.1% 4.7%

9. Generalization to other quasi-Newton methods.

24

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

9.1. Davidon–Fletcher–Powell. As indicated in Equation (6.13), CompactDen se::B.mult requires a Cholesky factorization. Although we found Cholesky variant to be faster than recursive version, as noted in Figures 6 and 12, for certain H0 scaling strategies, one may encounter zero pivot during Cholesky decomposition. 5 When Cholesky decomposition becomes numerically unstable, yet the computational efficiency of the dense method remains essential, the Davidon–Fletcher–Powell (DFP) method provides a viable alternative. It is well known that the DFP method is dual to BFGS method [8]. Therefore, from BFGS, by exchanging sk with yk , Bk with Hk , and Lk with RkT , one can obtain the DFP method. [9, Appendix D] With this, one can substitute CompactDense::B.mult with DenseDFP::B.solve and avoid numerically unstable Cholesky decomposition. Such a hybrid approach is implemented in MATLMVMDQN matrix type. However, this practice comes at the cost of a slight duality gap. 9.2. Broyden family. PETSc/TAO provides several types of limited-memory variable-metric (LMVM) matrices. In developing our dense BFGS formulations, we identified reusable design patterns common to the MATLMVM family, enabling the techniques used for dense BFGS to be applied to other MATLMVM variants as well. Consequently, the following matrices in PETSc/TAO now support dense formulations aimed at performance portability: Broyden, Bad Broyden, Symmetric Broyden, Symmetric Bad Broyden, and SR1. 6 7 . For brevity, we do not report benchmarks for these variants. 10. Discussion. Quasi-Newton Krylov trust region. This work has focused on L-BFGS globalized by a line search, which requires only one application of Hk per iteration. When using L-BFGS within a trust-region optimization method, an approximation to the forward Hessian B(xk ) is also needed for the Steihaug-CG bounded Krylov subspace iteration. Both the recursive and dense formulations of L-BFGS require additional setup computations to construct the data structures needed to efficiently apply Hk−1 to a vector. One approach that would avoid these costs would be to perform Steihaug-CG with a limited memory L-DFP approximation of Bk preconditioned by the L-BFGS Hk approximation. This approach would require minimal setup costs. It is not implemented in PETSc/TAO but could be the subject of further investigation. 11. Conclusion. We have introduced a new L-BFGS representation—the intermediate dense formulation—and implemented it in PETSc/TAO alongside the compact dense formulation of [6]. The intermediate dense method preserves the main performance benefits of dense formulations (improved locality and reduced synchronization relative to the recursive two-loop recursion) while avoiding a key limitation of compact dense methods in variable-metric settings: extra passes over the history and repeated applications of a changing base inverse Hessian H0 . We have evaluated three PETSc/TAO implementations (recursive, compact dense, and intermediate dense) on CPUs and GPUs on the ALCF Polaris and OLCF Frontier. Results show that dense methods consistently outperform the recursive approach at large problem sizes, that compact dense is most effective when H0 is constant or inexpensive, and that the intermediate dense approach provides the most robust performance across architectures 5 -mat lmvm J0 mat type diagonal -mat lmvm scale type diagonal -mat lmvm alpha 0.25 -mat lmvm theta 1.0 for the test used in the paper 6 MATLMVMBROYDEN, MATLMVMBADBROYDEN, MATLMVMSYMBROYDEN, MATLMVMSYMBADBROYDEN, MATLMVMSR1. 7 -mat lmvm mult algorithm <recursive, compact dense, dense> options flag can be used

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

25

and base-matrix choices. On the MNIST neural network training benchmark, we have demonstrated end-to-end gains from the dense implementation, achieving up to a 4.7% solver-time speedup over recursive BFGS. Acknowledgments. This material is based upon work supported by the U.S. Department of Energy (DOE) under Contract No. DE-AC02-06CH11357. This material is based upon work partially supported by the U.S. DOE, Office of Science, Office of Advanced Scientific Computing Research by the Competitive Portfolios for Advanced Scientific Computing Research Program, by the Scientific Discovery through Advanced Computing (SciDAC) Program through the FASTMath Institute, and by the Exascale Computing Project (ECP). This research used supporting resources at the Argonne and the Oak Ridge Leadership Computing Facilities. The Argonne Leadership Computing Facility at Argonne National Laboratory is supported by the Office of Science of the U.S. DOE under Contract No. DE-AC02-06CH11357. The Oak Ridge Leadership Computing Facility at the Oak Ridge National Laboratory is supported by the Office of Science of the U.S. DOE under Contract No. DE-AC05-00OR22725. REFERENCES [1] H. Anzt, Y. M. Tsai, A. Abdelfattah, T. Cojean, and J. Dongarra, Evaluating the performance of NVIDIA’s A100 Ampere GPU for sparse and batched computations, in 2020 IEEE/ACM Performance Modeling, Benchmarking and Simulation of High Performance Computer Systems (PMBS), IEEE, Nov. 2020, https://doi.org/10.1109/pmbs51919.2020. 00009. [2] S. Balay, S. Abhyankar, M. F. Adams, S. Benson, J. Brown, P. Brune, K. Buschelman, E. M. Constantinescu, L. Dalcin, A. Dener, V. Eijkhout, J. Faibussowitsch, W. D. Gropp, V. Hapla, T. Isaac, P. Jolivet, D. Karpeev, D. Kaushik, M. G. Knepley, F. Kong, S. Kruger, D. A. May, L. C. McInnes, R. T. Mills, L. Mitchell, T. Munson, J. E. Roman, K. Rupp, P. Sanan, J. Sarich, B. F. Smith, S. Zampini, H. Zhang, H. Zhang, and J. Zhang, PETSc Web page. https://petsc.org/, 2023, https://petsc.org/. [3] S. Balay, S. Abhyankar, M. F. Adams, J. Brown, P. Brune, K. Buschelman, E. M. Constantinescu, A. Dener, J. Faibussowitsch, W. D. Gropp, et al., PETSc/TAO users manual revision 3.22, tech. report, Argonne National Laboratory (ANL), Argonne, IL (United States), 09 2024, https://doi.org/10.2172/2476320, https://www.osti.gov/biblio/ 2476320. [4] J. Brown, V. Barra, N. Beams, L. Ghaffari, M. Knepley, W. Moses, R. Shakeri, K. Stengel, J. L. Thompson, and J. Zhang, Performance portable solid mechanics via matrix-free p-multigrid, 2022, https://arxiv.org/abs/2204.01722. [5] J. Brown and P. Brune, Low-rank quasi-Newton updates for robust Jacobian lagging in Newton-type methods, in International Conference on Mathematics and Computational Methods Applied to Nuclear Science and Engineering, 2013, pp. 2554–2565. [6] R. H. Byrd, J. Nocedal, and R. B. Schnabel, Representations of quasi-newton matrices and their use in limited memory methods, Mathematical Programming, 63 (1994), pp. 129–156, https://doi.org/10.1007/bf01582063. [7] A. Dener and T. Munson, Accelerating limited-memory quasi-newton convergence for largescale optimization, in Lecture Notes in Computer Science, Springer International Publishing, 2019, pp. 495–507, https://doi.org/10.1007/978-3-030-22744-9 39. [8] J. B. Erway, V. Jain, and R. F. Marcia, Shifted limited-memory DFP systems, in 2013 Asilomar Conference on Signals, Systems and Computers, IEEE, 2013, pp. 1033–1037. [9] J. B. Erway and R. F. Marcia, On solving large-scale limited-memory quasi-newton equations, Linear Algebra and its Applications, 515 (2017), pp. 196–225, https://doi.org/10.1016/j.laa. 2016.11.003. [10] J. Faibussowitsch, M. F. Adams, R. T. Mills, S. Zampini, and J. Zhang, Safe, seamless, and scalable integration of asynchronous GPU streams in PETSc, 2023, https://arxiv.org/ abs/2306.17801. [11] J. C. Gilbert and C. Lemaréchal, Some numerical experiments with variable-storage quasiNewton algorithms, Mathematical Programming, 45 (1989), pp. 407–435.

26

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

[12] J. Larsson Traff, W. D. Gropp, and R. Thakur, Self-consistent MPI performance guidelines, IEEE Transactions on Parallel and Distributed Systems, 21 (2010), pp. 698–709, https: //doi.org/10.1109/tpds.2009.120, http://dx.doi.org/10.1109/TPDS.2009.120. [13] D. C. Liu and J. Nocedal, On the limited memory BFGS method for large scale optimization, Mathematical Programming, 45 (1989), pp. 503–528, https://doi.org/10.1007/bf01589116. [14] R. T. Mills, M. F. Adams, S. Balay, J. Brown, A. Dener, M. Knepley, S. E. Kruger, H. Morgan, T. Munson, K. Rupp, B. F. Smith, S. Zampini, H. Zhang, and J. Zhang, Toward performance-portable PETSc for GPU-based exascale systems, Parallel Computing, 108 (2021), p. 102831, https://doi.org/10.1016/j.parco.2021.102831. [15] J. Nocedal and S. J. Wright, Numerical Optimization, Springer New York, 2006, https: //doi.org/10.1007/978-0-387-40065-5. [16] J. Zhang, J. Brown, S. Balay, J. Faibussowitsch, M. Knepley, O. Marin, R. T. Mills, T. Munson, B. F. Smith, and S. Zampini, The PetscSF scalable communication layer, IEEE Transactions on Parallel and Distributed Systems, 33 (2022), pp. 842–853, https: //doi.org/10.1109/tpds.2021.3084070, http://dx.doi.org/10.1109/TPDS.2021.3084070.

Appendix A. Recursive algorithms. In this appendix we present auxiliary algorithms for the recursive algorithm. Algorithm 10: update P0(m, k, B0 , S, P̊ ) Data: m, k: history size and current iterate Data: S ∈ Rn×m column-vector matrix, Si mod m = si , i ∈ ι(m, k) Data: P̊ ∈ Rn×m column-vector matrix, P̊i mod m = p̊i , i ∈ ι(m, k − 1) Result: α ∈ R, P̃ ∈ Rn×m such that αP̃i mod m = p̊i , i ∈ ι(m, k) 1 if B0 = γI then 2 return (γ, S) // return S by reference, no copy if B0 has changed then for j ∈ ι(m, k − 1) do 5 i ← j mod m 6 P̊i ← B0 Si

3

4

i ← (k − 1) mod m 8 P̊i ← B0 Si 9 return (1, P̊ )

// B0 .mult()

7

// B0 .mult() // return P̊ by reference, no copy

27

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

Algorithm 11: Recursive::B.update mult(x, g) m, k, B0 , S, Y, d, P, P̊ , f = B.(. . . ) 2 (α, P̃ ) ← update P0(m, k, B0 , S, P̊ ) 3 if B0 has changed or k > m then 4 ι ← ι(m, k) 1

5 6

else ι ← {k − 1}

for j ∈ ι in order do i ← j mod m 9 m̃ ← m − (k − j) 10 a, b ∈ Rm̃ 11 ω ← {ℓ mod m : ℓ ∈ ι(m̃, j)} 12 Pi ← αP̃i    T a 13 ← Yω Pω Si b    −1   a diag(dω ) 0 a 14 ← b 0 diag(fω ) b     a 15 Pi ← Pi + Yω Pω b 16 fi ← SiT Pi

// algorithm 10 // all p vectors must be recomputed // only pm−1,k−1 must be computed

7

8

// temporary work vectors // axpy(n) // 2m̃ dot(n), allreduce(2m̃) // 2 diags(m̃) // 2m̃ axpy(n) // dot(n), allreduce(1)

Appendix B. Compact dense algorithms. In this appendix we present auxiliary algorithms for the compact dense algorithm. Algorithm 12: update Atb(m, k, β, A, bold , bnew , c, e) Data: (m, k): history size and current iterate Data: β ∈ R Data: A: RN ×m cyclic-order column-vector matrix Data: bold , bnew : Rn vectors Data: c: Rm vector where cj mod m = βATj mod m bold for j ∈ ι(m, k − 1) Result: c: Rm vector where cj mod m = βATj mod m bnew for j ∈ ι(m, k) Result: e: Rm vector where cj mod m = βATj mod m (bnew − bold ) for j ∈ ι(m, k) 1 i ← (k − 1) mod m 2 ω ← ω(m, k) // ATi bold has not been computed yet, compute it now T 3 ci ← βAi bold // dot(n), reduce(1) 4 eω = −cω // axpy(mk ) T 5 cω ← βAω bnew // gemv(n, mk ), reduce(mk ) 6 eω ← eω + cω // axpy(mk )

28

H. SUH, T. ISAAC, A. DENER, T. MUNSON, H. ZHANG, R. T. MILLS

Algorithm 13: CompactDense::B.update solve(x, g) m, k, B0 , S, Y, d, Q, R, Z, g prev , acache , bcache = B.(. . . ) 2 (α, Q̃) ← update Q(m, k, B0 , Y, Q) // equivalent to algorithm 10, αQ̃ = Q 3 i ← (k − 1) mod m // newest history index 4 ω ← {ℓ mod m : ℓ ∈ ι(m, k)} // active history indices 5 ω− ← {ℓ mod m : ℓ ∈ ι(m − 1, k − 1)} // active indices except the newest prev cache , g, a , R:,i ) // R and acache are up to date 6 update Atb(m, k, 1, S, g 7 if B0 did not change then 8 update Atb(m, k, α, Q̃, g prev , g, bcache , Z:,i ) 1

else if B0 updated by a scalar factor H0,k ← γH0,k−1 then Zω− ,ω− ← γZω− ,ω− // scal((mk − 1)2 ) cache cache // scal(mk − 1) 11 bω− ← γbω− 12 update Atb(m, k, β, Q̃, g prev , g, bcache , Z:,i ) 9

10

else Qω ← B0 Yω 15 Zω,ω ← QTω Yω ← QTω g prev 16 bcache ω

13

14

17

// mk B0 .mult() // gemm(mk , mk , n), reduce(m2k ) // gemv(n, mk ), reduce(mk )

Zi,ω− ← Zω− ,i

// copy(m̃ − 1)

Algorithm 14: CompactDense::B.update mult(x, g) m, k, B0 , S, Y, d, P̊ , L, F̊ , J = B.(. . . ) M ∈ Rmk ×mk // temporary work matrix 3 i ← (k − 1) mod m 4 ω ← ω(m, k) // active history indices 5 ω− ← {ℓ mod m : ℓ ∈ ι(m − 1, k − 1)} // active indices except the newest 6 (α, P̃ ) ← update P0(m, k, B0 , S, P̊ ) // algorithm 10 T 7 Li,ω− ← Si Yω− // gemv(n, mk − 1), reduce(mk − 1) 8 Lω− ,i ← 0 // zeroing this column ensures L = hstrilk (S T Y ) 9 if B0 did not change then // gemv(n, mk ), reduce(mk ) 10 Fω,i ← αP̃ωT Si 11 Fi,ω− ← Fω− ,i // copy(mk − 1) 1

2

else if B0 updated by a scalar factor B0,k ← γB0,k−1 then Fω− ,ω− ← γFω− ,ω− // scal((mk − 1)2 ) T 14 Fω,i ← αP̃ω Si // gemv(n, mk ), reduce(mk ) 15 Fi,ω− ← Fω− ,i // copy(mk − 1)

12

13

16 17

else Fω,ω ← αP̃ωT Sω

M ← Lω,ω diag(dω )−1 LTω,ω 19 Jω,ω ← Cholesky(M + F̊ω,ω ) 18

// gemm(mk , mk , mk ), reduce(m2k ) // mk diags(mk ), gemm(mk , mk , mk ) // axpy(m2k ), potrf(mk )

Appendix C. Intermediate dense algorithms. In this appendix section, we present auxiliary algorithms for intermediate dense algorithm.

PORTABLE AND VERSATILE L-BFGS IN PETSC/TAO

Algorithm 15: IntermediateDense::B.update solve(g) m, k, S, Y, R, d, g prev , acache = B.(. . . ) 2 i ← k mod m 3 l ← max{m − 1, k − 1} // store S T g as part of computing r = S T y = S T (g − g prev ) prev 4 update Atb(m, k, 1, S, g , g, acache , R:,i ) // the upper triangle of R with respect to history index (R(i mod m,j mod m) , j ≥ i) now contains current values of S T Y

1

29

Related documents

• Nexus: Depth-Adaptive KV-Cache Splicing and Retrieval-Decoupled Tool Routing for Agentic LLMs on Unified Memory#658973
• Nexus: Depth-Adaptive KV-Cache Splicing and Retrieval-Decoupled Tool Routing for Agentic LLMs on Unified Memory#656837
• Sangam: Efficiently Serving Diffusion LLMs with the AR Stack#343436
• HyperParallel-FSDP: Topology-Aware Fully Sharded Training with Layout-Driven Muon on Ascend SuperPods#1035990
• HyperParallel-FSDP: Topology-Aware Fully Sharded Training with Layout-Driven Muon on Ascend SuperPods#1006819