Pure JAX implementation of ormqr for faster QR solves

[jpbrodrick89]

Ported from jax-ml/jax#36575, this was found to beat cusolver by a factor of 2 for your matrix sizes, so compounded with ormqr outperforming orgqr (raw vs economic) you should see at least a 3x speedup in each QR ormqr (not qr_multiply as a whole as the geqrf call is not sped up) call and probably a lower memory usage too. Crucially, it does not require a jaxlib upgrade or rebuild.

This employs a blocked householder multiplication tuned to an A100, your mileage may vary on other devices.

I tested this on my Mac CPU and observed speedups of 1.6-1.9x on each qr_multiply call.

Let me know if you think the added (hopefully temporary) maintenance burden is worth the performance gain to you?

I wasn't sure where in your codebase to put this so feel free to move this wherever you see fit.

(a) Big Jacobian qr_multiply(J, f) — tall, ratio 6

┌───────┬──────┬───────┬───────┬────────┬──────────┬────────────┐
│   m   │  n   │ orig  │  new  │ native │ orig/new │ native/new │
├───────┼──────┼───────┼───────┼────────┼──────────┼────────────┤
│ 1800  │ 300  │ 18.5  │ 10.1  │ 11.2   │ 1.83×    │ 1.11×      │
├───────┼──────┼───────┼───────┼────────┼──────────┼────────────┤
│ 5136  │ 856  │ 182   │ 116   │ 120    │ 1.57×    │ 1.03×      │
├───────┼──────┼───────┼───────┼────────┼──────────┼────────────┤
│ 20016 │ 3336 │ 6854  │ 4200  │ 4158   │ 1.63×    │ 0.99×      │
├───────┼──────┼───────┼───────┼────────┼──────────┼────────────┤
│ 31128 │ 5188 │ 24277 │ 12647 │ 12206  │ 1.92×    │ 0.97×      │
└───────┴──────┴───────┴───────┴────────┴──────────┴────────────┘

(b) Reduced loop qr_multiply([R; √α·I], zp) — (2n, n), per LM iteration

┌───────┬──────┬──────┬──────┬────────┬──────────┬────────────┐
│   m   │  n   │ orig │ new  │ native │ orig/new │ native/new │
├───────┼──────┼──────┼──────┼────────┼──────────┼────────────┤
│ 600   │ 300  │ 3.67 │ 2.16 │ 1.86   │ 1.70×    │ 0.86×      │
├───────┼──────┼──────┼──────┼────────┼──────────┼────────────┤
│ 1712  │ 856  │ 47.5 │ 25.8 │ 25.3   │ 1.84×    │ 0.98×      │
├───────┼──────┼──────┼──────┼────────┼──────────┼────────────┤
│ 5000  │ 2500 │ 1073 │ 611  │ 664    │ 1.76×    │ 1.09×      │
├───────┼──────┼──────┼──────┼────────┼──────────┼────────────┤
│ 10376 │ 5188 │ 6589 │ 3663 │ 3653   │ 1.80×    │ 1.00×      │
└───────┴──────┴──────┴──────┴────────┴──────────┴────────────┘

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -0.55 %    |     4.046e+03      |     4.024e+03      |    -22.36    |       32.66        |       29.03        |
  test_proximal_jac_w7x_with_eq_update   |   -1.37 %    |     6.623e+03      |     6.532e+03      |    -90.54    |       153.99       |       153.09       |
  test_proximal_freeb_jac                |    0.41 %    |     1.333e+04      |     1.338e+04      |    53.99     |       82.03        |       81.92        |
  test_proximal_freeb_jac_blocked        |    0.90 %    |     7.620e+03      |     7.689e+03      |    68.26     |       69.68        |       68.38        |
  test_proximal_freeb_jac_batched        |   -0.64 %    |     7.649e+03      |     7.600e+03      |    -48.97    |       67.96        |       68.28        |
  test_proximal_jac_ripple               |   -0.60 %    |     3.549e+03      |     3.528e+03      |    -21.27    |       54.05        |       53.25        |
  test_proximal_jac_ripple_bounce1d      |   -0.16 %    |     3.755e+03      |     3.748e+03      |    -6.18     |       67.31        |       69.02        |
  test_eq_solve                          |    1.65 %    |     1.926e+03      |     1.957e+03      |    31.71     |       52.66        |       58.29        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.30 +/- 4.11     | +2.48e-03 +/- 3.35e-02 |  8.16e-01 +/- 2.6e-02  |  8.14e-01 +/- 2.1e-02  |
 test_equilibrium_init_medres            |     -0.07 +/- 2.82     | -4.81e-03 +/- 1.83e-01 |  6.47e+00 +/- 1.3e-01  |  6.48e+00 +/- 1.3e-01  |
 test_equilibrium_init_highres           |     +0.54 +/- 2.11     | +3.90e-02 +/- 1.54e-01 |  7.31e+00 +/- 1.2e-01  |  7.27e+00 +/- 1.0e-01  |
 test_objective_compile_dshape_current   |     +1.65 +/- 2.84     | +6.73e-02 +/- 1.16e-01 |  4.15e+00 +/- 6.8e-02  |  4.08e+00 +/- 9.4e-02  |
 test_objective_compute_dshape_current   |     -6.06 +/- 17.85    | -4.44e-05 +/- 1.31e-04 |  6.88e-04 +/- 8.7e-05  |  7.33e-04 +/- 9.8e-05  |
 test_objective_jac_dshape_current       |     +6.40 +/- 25.72    | +1.43e-03 +/- 5.75e-03 |  2.38e-02 +/- 4.7e-03  |  2.23e-02 +/- 3.3e-03  |
 test_perturb_2                          |     +0.13 +/- 0.65     | +2.45e-02 +/- 1.24e-01 |  1.91e+01 +/- 7.0e-02  |  1.90e+01 +/- 1.0e-01  |
 test_proximal_jac_atf_with_eq_update    |     -0.21 +/- 1.17     | -2.52e-02 +/- 1.38e-01 |  1.18e+01 +/- 7.4e-02  |  1.18e+01 +/- 1.2e-01  |
 test_proximal_freeb_jac                 |     +0.37 +/- 3.72     | +1.71e-02 +/- 1.73e-01 |  4.68e+00 +/- 8.3e-02  |  4.66e+00 +/- 1.5e-01  |
+test_solve_fixed_iter_compiled          |     -5.63 +/- 1.34     | -3.46e-01 +/- 8.22e-02 |  5.80e+00 +/- 7.2e-02  |  6.15e+00 +/- 3.9e-02  |
 test_LinearConstraintProjection_build   |     +0.76 +/- 4.56     | +5.06e-02 +/- 3.03e-01 |  6.69e+00 +/- 1.8e-01  |  6.64e+00 +/- 2.5e-01  |
 test_objective_compute_ripple_bounce1d  |     +0.29 +/- 1.97     | +8.54e-04 +/- 5.83e-03 |  2.97e-01 +/- 4.1e-03  |  2.96e-01 +/- 4.1e-03  |
 test_objective_grad_ripple_bounce1d     |     +1.71 +/- 2.01     | +1.55e-02 +/- 1.83e-02 |  9.22e-01 +/- 1.6e-02  |  9.07e-01 +/- 9.6e-03  |
 test_build_transform_fft_midres         |     +0.52 +/- 4.06     | +4.69e-03 +/- 3.67e-02 |  9.10e-01 +/- 1.3e-02  |  9.05e-01 +/- 3.4e-02  |
 test_build_transform_fft_highres        |     +1.42 +/- 2.69     | +1.69e-02 +/- 3.20e-02 |  1.21e+00 +/- 2.3e-02  |  1.19e+00 +/- 2.2e-02  |
 test_equilibrium_init_lowres            |     +1.33 +/- 3.15     | +8.81e-02 +/- 2.08e-01 |  6.70e+00 +/- 1.5e-01  |  6.61e+00 +/- 1.4e-01  |
 test_objective_compile_atf              |     +2.00 +/- 3.79     | +1.26e-01 +/- 2.39e-01 |  6.43e+00 +/- 1.2e-01  |  6.31e+00 +/- 2.1e-01  |
 test_objective_compute_atf              |    +11.81 +/- 15.89    | +2.51e-04 +/- 3.38e-04 |  2.38e-03 +/- 1.7e-04  |  2.13e-03 +/- 2.9e-04  |
 test_objective_jac_atf                  |     +5.08 +/- 2.44     | +8.10e-02 +/- 3.89e-02 |  1.68e+00 +/- 2.9e-02  |  1.60e+00 +/- 2.6e-02  |
 test_perturb_1                          |     +1.79 +/- 1.34     | +2.89e-01 +/- 2.16e-01 |  1.65e+01 +/- 1.7e-01  |  1.62e+01 +/- 1.4e-01  |
 test_proximal_jac_atf                   |     +0.99 +/- 2.11     | +5.20e-02 +/- 1.10e-01 |  5.29e+00 +/- 8.7e-02  |  5.23e+00 +/- 6.8e-02  |
 test_proximal_freeb_compute             |     +1.09 +/- 2.69     | +1.79e-03 +/- 4.44e-03 |  1.67e-01 +/- 3.7e-03  |  1.65e-01 +/- 2.5e-03  |
 test_solve_fixed_iter                   |     +3.12 +/- 2.02     | +7.54e-01 +/- 4.88e-01 |  2.49e+01 +/- 3.6e-01  |  2.42e+01 +/- 3.3e-01  |
 test_objective_compute_ripple           |     +0.36 +/- 7.46     | +8.46e-04 +/- 1.78e-02 |  2.39e-01 +/- 1.6e-02  |  2.38e-01 +/- 7.4e-03  |
 test_objective_grad_ripple              |     -0.89 +/- 3.13     | -7.94e-03 +/- 2.79e-02 |  8.82e-01 +/- 2.4e-02  |  8.90e-01 +/- 1.5e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[jpbrodrick89]

So I tried to get a handle on what was going on with tests, but I'm not very familiar with the codebase so might be a poor take.

1. The regression test seems to due to an ordering constraint that is hard for the augmented Lagrangian method to obey. I get very close to the 1e-4 limit on CI but don't reach it in 500 iterations. On my Mac I reach it in 320 iterations, but even the current status quo shows oscillation (i.e. once the constraint is met it can be broken again in subsequent iterations). I think the test isn't guaranteed to pass, and that it is flaky but I might have misunderstood.
2. Memory issues, CI reports ~30% peak memory increase, trying to run on my Mac I only measure a 10% increase. This is slightly counterintuitive to me, I tried a few fixed and none moved the dial. I think more detailed profiling is necessary to really understand what is driving the increase in memory. There's a chance this PR is a memory vs compute tradeoff that you might have to make an informed rational decision on.
3. Yes, I think the slight compile time increase is real but hopefully its worth the reduced runtime.

[YigitElma]

1. I think the test isn't guaranteed to pass, and that it is flaky but I might have misunderstood.

That can be true. That test was prone to slight numerical changes that new jax versions introduce. It tends to be more unstable on CI, I am not sure why. But I think the problem here may be solved by changing the method that we solve the trust region problem to svd. The test solves an optimization problem that has bounds, and inside the bounds, the value/derivative of the objective and constraint becomes 0. This may cause the Jacobian to be singular, and in that case, the non-rank-revealing QR doesn't work. I can change the test.

But otherwise, how much numerical difference do you think your method has compared to before?

For other benchmarks, we can rerun the CI since it can be noisy due to GitHub runners. For example, this PR shouldn't have any effect on test_objective_compile_dshape_current or test_build_transform_fft_lowres.

Memory/runtime/compile time is a hard one. Since these can also be problem dependent, we try to decide based on the most common problem types.

[YigitElma]

Also, it might be worth to try jit compiling the new implementation, there seems to be multiple slicing operations that will trigger individual dispatches and compiling. We've previously seen that these individual costs can be overcome by jitting the outer function. For example, see #2153

Adds jit decorator to ObjectiveFunction.unpack_params() (this makes the difference in benchmarks, specifically this line) and to ObjectiveFunction.x(). The substantial speed-up (~1.5s) comes from the compilation overhead, but even after that, each call is 3-4 times faster (but these take couple milliseconds)

[YigitElma]

This may cause the Jacobian to be singular, and in that case, the non-rank-revealing QR doesn't work. I can change the test.

Yes, if I print the shape and the rank of the matrices we solve, the first step gives (1133, 681) 473, so it is singular. The rest of the iterations seem to be fine. Also, the test passed for me locally.

About the memory, this is a crude benchmark on laptop GPU again comparing PR (with qr_multiply jitted), MASTER and PR (without qr_mutliply jitted). Seems like jit solves the problem, let's see what CI benchmarks say.

For clarity, just jit vs not jit,

For clarity, just jit vs master,

For clarity, just no jit vs master,

[codecov]

Codecov Report

❌ Patch coverage is 88.46154% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.32%. Comparing base (7b6ed16) to head (25ce309).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/_qr_multiply.py	88.23%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2247      +/-   ##
==========================================
+ Coverage   94.30%   94.32%   +0.01%     
==========================================
  Files         101      102       +1     
  Lines       28866    28912      +46     
==========================================
+ Hits        27223    27270      +47     
+ Misses       1643     1642       -1     

Files with missing lines	Coverage Δ
desc/backend.py	89.39% <100.00%> (+0.23%)	⬆️
desc/_qr_multiply.py	88.23% <88.23%> (ø)

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jpbrodrick89]

Ah, this is a good spot! I didn't know you were running your loop eagerly (I'm used to optimistix where all our loops are equinox while loops which jits the step function, any reason why don't jit your whole step function?). Looks a lot better with no jit, no memory issues. Would be interested to see how it performs on your production GPU as the laptop GPU may not be optimised enough to benefit from it (it relies on reasonable GPU kernel dispatch (we are launching a lot of kernels whereas cusolver ormqr is theoretically just one), lightning fast GEMM's (which NVIDIA GPU's have) and fairly quick TRSV's). If you just get 5% speedup like on your laptop GPU this is probably not worth it but I'mhoping it will be better than that.

[YigitElma]

any reason why don't jit your whole step function?

There are some objectives and routines that are very hard to make jit-friendly in DESC. Also, we want to keep our optimizers compatible with non-JAX functions. There is still some effort to make it possible, but it is not a priority right now.

I see around 10% improvement in speed and a 30-40% reduction in the max memory of the QR-related region. I think this is good even on my small GPU. We usually use A100's on Princeton clusters and NERSC Perlmutter, so I think your own benchmarks will be reflected here.

In general, computing the Jacobian takes significantly more memory than trust region operations, but this might be useful to minimize the total memory footprint when we compute the Jacobian column by column (then QR becomes the dominant memory bottleneck). Given that, I think we need to discuss/benchmark this further to make sure that it is worth the additional complexity.

I will trigger CI tests with older JAX versions to make sure that the dependency range is still correct.