Skip to content

Optimizer: avoid heap allocations per evaluation and per iteration in L-BFGS/BFGS - #3445

Open
jachymb wants to merge 2 commits into
stan-dev:developfrom
jachymb:feature/optimizer-fewer-allocations
Open

jachymb wants to merge 2 commits into
stan-dev:developfrom
jachymb:feature/optimizer-fewer-allocations

Conversation

@jachymb

@jachymb jachymb commented Oct 18, 2026

Copy link
Copy Markdown

A relatively small, but noticeable memory improvement for the optimizers.

AI use disclosure: This was developed with the help of Claude Code (Opus 5.5, model ID claude-opus-5-5).

Submission Checklist

  • Run unit tests: ./runTests.py src/test/unit
  • Run cpplint: make cpplint
  • Declare copyright holder and open-source license: see below

Summary

A refactor of src/stan/optimization/ that keeps every iterate and result:

  • ModelAdaptor::operator()(x, f, g) called stan::model::log_prob_grad, which allocates a std::vector<var> per call, after copying x into a std::vector<double>, and whose gradient was then copied into g. It now keeps the std::vector<var> as a member, calls model.log_prob<true, jacobian> on it and writes the adjoints into g directly. Errors are handled as before: std::domain_error returns 1 with the message, other exceptions propagate, autodiff memory is recovered on every path, a non-finite gradient returns 3 (g then holds the whole gradient, where develop stopped writing at the first non-finite entry).
  • s and y of BFGSMinimizer::step and the two-loop workspace of LBFGSUpdate::search_direction are members instead of locals, and WolfeLineSearch no longer copies the previous gradient. That copy reached the caller only through the gradient output of a failed search, which no caller reads.

This removes one heap allocation per gradient evaluation, one per line search, and two per iteration (three for L-BFGS).

Intended Effect

Less optimizer overhead per evaluation and per iteration, which matters for cheap models with many iterations. Measured in process against develop's code (-O3 -DNDEBUG, 3 runs of interleaved rounds), branch / develop time:

WSL2, g++ 12.2 MSYS2 MinGW-w64, g++ 15.2
L-BFGS iteration on a cheap quadratic, $n = 2$–$10$ 0.66–0.85 0.49–0.63
the same, $n = 100 – 100000$ 0.68–0.96 0.80–0.98

The same harness gives 0.92–1.06 (WSL2) and 0.98–1.10 (MinGW) for code that is identical on both sides. ModelAdaptor::operator() saves one std::vector<var> allocation and two std::vector<double> copies per gradient evaluation; on two 2-parameter models its measured ratio, 0.90–0.96 (WSL2), is within that noise.

How to Verify

The refactor keeps the values and return codes the optimizer uses, so no test fails on develop, and it adds none. Existing tests that reach the changed code pass unchanged (bfgs_exponential_boundary_test.cpp runs the optimizer until its line search fails): bfgs_linesearch_test.cpp, bfgs_minimizer_test.cpp, bfgs_test.cpp, bfgs_exponential_boundary_test.cpp, lbfgs_update_test.cpp, services/optimize/{bfgs,lbfgs,bfgs_jacobian,lbfgs_jacobian}_test.cpp, and the Pathfinder tests services/pathfinder/{eight_schools,normal_glm}_test.cpp.

Side Effects

Values, iterates and output are unchanged: an in-process A/B harness against develop (20 model configurations, 4350 L-BFGS and BFGS runs, 960 Pathfinder runs; I can share the harness) gives the same iterations, gradient evaluations and final log density, bit for bit, in every run, and the same Pathfinder ELBO to the six decimals of its log message. Memory: the work vectors live as long as the optimizer, $O(n)$ doubles, plus $n$ vars in ModelAdaptor.

Documentation

None.

Copyright and Licensing

Please list the copyright holder for the work you are submitting (this will be you or your assignee, such as a university or company):

Jáchym Barvínek jachymb@gmail.com

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the following licenses:

ModelAdaptor computes the gradient with a member std::vector<var> instead
of calling log_prob_grad, which allocated and copied the parameters and
the gradient on every call. s and y of BFGSMinimizer and the two-loop
workspace of LBFGSUpdate are members, and WolfeLineSearch no longer
copies the previous gradient, a copy that reached the caller only as the
gradient output of a failed search.
@stan-buildbot

Copy link
Copy Markdown
Contributor
Name Old Result New Result Ratio Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan 0.04 0.04 1.03 2.53% faster
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan 0.16 0.16 1.0 0.47% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan 0.86 0.86 1.0 0.08% faster
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan 0.01 0.01 1.0 -0.42% slower
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan 0.11 0.11 1.02 1.52% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan 8.1 7.99 1.01 1.4% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan 3.66 3.64 1.0 0.48% faster
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan 33.11 33.5 0.99 -1.17% slower
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan 0.39 0.38 1.01 1.46% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan 6.82 6.81 1.0 0.08% faster
stat_comp_benchmarks/benchmarks/sir/sir.stan 166.73 168.0 0.99 -0.76% slower
stat_comp_benchmarks/benchmarks/arK/arK.stan 3.22 3.21 1.0 0.44% faster
stat_comp_benchmarks/benchmarks/arma/arma.stan 0.6 0.6 1.0 0.18% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan 25.87 25.6 1.01 1.04% faster
performance.compilation 394.6 385.38 1.02 2.34% faster
Mean result: 1.006584545817147

Jenkins Console Log
Jenkins Build Stages
Commit hash: 6a335dd02ea35aac291d2cb571d03092a8b16ad7

Machine information
Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           52 bits physical, 57 bits virtual
CPU(s):                                  192
On-line CPU(s) list:                     0-191
Thread(s) per core:                      2
Core(s) per socket:                      48
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              25
Model:                                   17
Model name:                              AMD EPYC 9474F 48-Core Processor
Stepping:                                1
Frequency boost:                         enabled
CPU MHz:                                 1497.669
CPU max MHz:                             4114.4229
CPU min MHz:                             1500.0000
BogoMIPS:                                7189.39
Virtualization:                          AMD-V
L1d cache:                               3 MiB
L1i cache:                               3 MiB
L2 cache:                                96 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-47,96-143
NUMA node1 CPU(s):                       48-95,144-191
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

@avehtari

Copy link
Copy Markdown
Member

Claude commented something worth checking:

  • LBFGSUpdate::search_direction is const, but it now writes to the mutable member _alphas. Nothing in Stan calls it concurrently on one object, so this is not a bug today. A const method usually promises that concurrent calls are safe, and this one no longer keeps that promise. I suggest removing const from search_direction and mutable from _alphas: every caller in Stan (BFGSMinimizer::step and the two update tests) uses a non-const object, so nothing else changes.

…kspace

The two-loop workspace became a mutable member, so a const call was no
longer safe to make concurrently. No caller uses a const LBFGSUpdate.
@jachymb

jachymb commented Oct 19, 2026

Copy link
Copy Markdown
Author

Thanks, agreed. Adjusted according to the recommendation.

@stan-buildbot

Copy link
Copy Markdown
Contributor
Name Old Result New Result Ratio Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan 0.04 0.04 1.0 -0.07% slower
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan 0.16 0.16 1.0 -0.36% slower
stat_comp_benchmarks/benchmarks/garch/garch.stan 0.87 0.87 1.01 1.12% faster
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan 0.01 0.02 0.92 -8.44% slower
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan 0.11 0.11 1.08 7.61% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan 8.0 7.98 1.0 0.18% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan 3.68 3.68 1.0 -0.04% slower
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan 32.74 32.56 1.01 0.53% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan 0.38 0.38 1.01 1.13% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan 6.82 6.81 1.0 0.07% faster
stat_comp_benchmarks/benchmarks/sir/sir.stan 171.42 168.02 1.02 1.98% faster
stat_comp_benchmarks/benchmarks/arK/arK.stan 3.23 3.2 1.01 0.71% faster
stat_comp_benchmarks/benchmarks/arma/arma.stan 0.61 0.61 1.0 0.49% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan 25.9 25.79 1.0 0.44% faster
performance.compilation 396.22 398.35 0.99 -0.54% slower
Mean result: 1.004116851403355

Jenkins Console Log
Jenkins Build Stages
Commit hash: a2236819d562837e97a9bfe97dc9ebd6d8613729

Machine information
Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           52 bits physical, 57 bits virtual
CPU(s):                                  192
On-line CPU(s) list:                     0-191
Thread(s) per core:                      2
Core(s) per socket:                      48
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              25
Model:                                   17
Model name:                              AMD EPYC 9474F 48-Core Processor
Stepping:                                1
Frequency boost:                         enabled
CPU MHz:                                 1497.419
CPU max MHz:                             4114.4229
CPU min MHz:                             1500.0000
BogoMIPS:                                7189.39
Virtualization:                          AMD-V
L1d cache:                               3 MiB
L1i cache:                               3 MiB
L2 cache:                                96 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-47,96-143
NUMA node1 CPU(s):                       48-95,144-191
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants