perf: use jax.linearize when chunking jacobians in forward-mode

[jpbrodrick89]

This should provide a runtime and compile time improvement at the possible cost of a very slight peak memory increase. I think if this pushes you over the limit you are probably in the region where you should be switching to iterative rather than direct solvers.

jax.linearize caches the residuals from the primal pass just like jax.vjp does which means sequentially evaluating (chunks of) jvp's it will be a lot faster. For example, if Proximal was fully jittable and went through this route this would have allowed the reuse of the SVD decomposition automatically. As it stands Proximal does not directly benefit but the improvement @YigitElma is making in #2239 should chain nicely with this (it would be interesting to see the benchmarks for these two PR's combined).

@unalmis this should probably be upstreamed in adv-jax-math, lmk if you'd like me to make a PR there.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.03 %    |     3.985e+03      |     3.986e+03      |     1.00     |       31.07        |       29.63        |
  test_proximal_jac_w7x_with_eq_update   |    3.47 %    |     6.534e+03      |     6.761e+03      |    227.04    |       155.08       |       152.65       |
  test_proximal_freeb_jac                |   -0.35 %    |     1.341e+04      |     1.336e+04      |    -46.28    |       80.32        |       82.21        |
  test_proximal_freeb_jac_blocked        |   -8.77 %    |     7.727e+03      |     7.049e+03      |   -677.75    |       73.27        |       69.00        |
  test_proximal_freeb_jac_batched        |   -7.53 %    |     7.634e+03      |     7.059e+03      |   -574.68    |       72.57        |       69.56        |
  test_proximal_jac_ripple               |   -1.43 %    |     3.597e+03      |     3.546e+03      |    -51.41    |       54.33        |       54.83        |
  test_proximal_jac_ripple_bounce1d      |    1.59 %    |     3.695e+03      |     3.754e+03      |    58.73     |       68.20        |       67.38        |
  test_eq_solve                          |   -6.87 %    |     2.125e+03      |     1.979e+03      |   -145.95    |       54.14        |       53.14        |
  test_objective_quadratic_flux_jac      |   -0.83 %    |     2.537e+03      |     2.516e+03      |    -21.16    |       53.95        |       53.90        |

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.51 +/- 4.95     | +4.23e-03 +/- 4.13e-02 |  8.39e-01 +/- 3.1e-02  |  8.35e-01 +/- 2.7e-02  |
 test_equilibrium_init_lowres            |     -1.98 +/- 3.20     | -1.30e-01 +/- 2.10e-01 |  6.42e+00 +/- 1.6e-01  |  6.55e+00 +/- 1.4e-01  |
 test_objective_compile_atf              |     -2.65 +/- 3.52     | -1.55e-01 +/- 2.06e-01 |  5.69e+00 +/- 1.4e-01  |  5.85e+00 +/- 1.5e-01  |
 test_objective_compute_atf              |     -7.21 +/- 9.50     | -2.00e-04 +/- 2.64e-04 |  2.58e-03 +/- 1.3e-04  |  2.78e-03 +/- 2.3e-04  |
 test_objective_jac_atf                  |     -3.32 +/- 4.93     | -6.02e-02 +/- 8.94e-02 |  1.75e+00 +/- 7.0e-02  |  1.81e+00 +/- 5.6e-02  |
 test_perturb_1                          |     -3.62 +/- 2.09     | -4.31e-01 +/- 2.49e-01 |  1.15e+01 +/- 1.8e-01  |  1.19e+01 +/- 1.7e-01  |
 test_proximal_jac_atf                   |     -3.77 +/- 2.15     | -2.06e-01 +/- 1.18e-01 |  5.27e+00 +/- 4.8e-02  |  5.47e+00 +/- 1.1e-01  |
 test_proximal_freeb_compute             |     -3.49 +/- 4.20     | -5.08e-03 +/- 6.11e-03 |  1.41e-01 +/- 4.4e-03  |  1.46e-01 +/- 4.3e-03  |
 test_solve_fixed_iter                   |     -2.06 +/- 2.92     | -5.28e-01 +/- 7.48e-01 |  2.51e+01 +/- 5.2e-01  |  2.56e+01 +/- 5.4e-01  |
 test_LinearConstraintProjection_build   |     -2.60 +/- 3.25     | -1.72e-01 +/- 2.15e-01 |  6.43e+00 +/- 1.0e-01  |  6.60e+00 +/- 1.9e-01  |
 test_objective_compute_ripple           |     +2.65 +/- 6.70     | +5.64e-03 +/- 1.43e-02 |  2.19e-01 +/- 1.3e-02  |  2.13e-01 +/- 6.7e-03  |
 test_objective_grad_ripple              |     -1.43 +/- 3.64     | -1.41e-02 +/- 3.60e-02 |  9.75e-01 +/- 3.0e-02  |  9.89e-01 +/- 2.0e-02  |
 test_objective_quadratic_flux_compute   |     -5.54 +/- 10.47    | -3.24e-03 +/- 6.12e-03 |  5.52e-02 +/- 2.5e-03  |  5.85e-02 +/- 5.6e-03  |
 test_build_transform_fft_midres         |     -0.72 +/- 3.31     | -6.23e-03 +/- 2.86e-02 |  8.59e-01 +/- 1.6e-02  |  8.65e-01 +/- 2.4e-02  |
 test_build_transform_fft_highres        |     +0.54 +/- 3.33     | +6.14e-03 +/- 3.81e-02 |  1.15e+00 +/- 3.1e-02  |  1.14e+00 +/- 2.3e-02  |
 test_equilibrium_init_medres            |     +0.14 +/- 3.80     | +9.79e-03 +/- 2.59e-01 |  6.81e+00 +/- 1.9e-01  |  6.80e+00 +/- 1.7e-01  |
 test_objective_compile_dshape_current   |     -0.30 +/- 3.24     | -1.18e-02 +/- 1.25e-01 |  3.86e+00 +/- 1.2e-01  |  3.87e+00 +/- 4.6e-02  |
 test_objective_compute_dshape_current   |    +10.81 +/- 14.07    | +6.47e-05 +/- 8.42e-05 |  6.63e-04 +/- 7.8e-05  |  5.99e-04 +/- 3.2e-05  |
 test_objective_jac_dshape_current       |    +14.35 +/- 28.13    | +3.29e-03 +/- 6.46e-03 |  2.63e-02 +/- 4.8e-03  |  2.30e-02 +/- 4.3e-03  |
 test_perturb_2                          |     +0.11 +/- 1.39     | +1.69e-02 +/- 2.06e-01 |  1.48e+01 +/- 1.5e-01  |  1.48e+01 +/- 1.5e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.56 +/- 0.97     | +6.93e-02 +/- 1.21e-01 |  1.26e+01 +/- 8.9e-02  |  1.25e+01 +/- 8.2e-02  |
 test_proximal_freeb_jac                 |     +0.26 +/- 4.12     | +1.23e-02 +/- 1.93e-01 |  4.70e+00 +/- 1.5e-01  |  4.69e+00 +/- 1.2e-01  |
 test_solve_fixed_iter_compiled          |     +0.96 +/- 1.85     | +5.91e-02 +/- 1.14e-01 |  6.21e+00 +/- 9.7e-02  |  6.15e+00 +/- 6.0e-02  |
 test_objective_compute_ripple_bounce1d  |     +1.05 +/- 5.50     | +2.91e-03 +/- 1.52e-02 |  2.80e-01 +/- 1.0e-02  |  2.77e-01 +/- 1.1e-02  |
 test_objective_grad_ripple_bounce1d     |     -0.89 +/- 2.14     | -8.73e-03 +/- 2.09e-02 |  9.71e-01 +/- 1.9e-02  |  9.79e-01 +/- 7.7e-03  |
 test_objective_quadratic_flux_jac       |     +0.30 +/- 0.86     | +2.56e-02 +/- 7.34e-02 |  8.58e+00 +/- 2.0e-02  |  8.55e+00 +/- 7.1e-02  |

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

[YigitElma]

The failing unit test is related to jax-finufft version. It is similar to the one we faced in #2264. This commit aimed to solve the issue but it didn't and the later commit deleted a unit test which in return moved the failing test to a different Python version that uses a newer jax-finufft (the test passed as if the original commit solved the issue). Anyway, I can fix this in a separate PR, the changes in this PR shouldn't affect Gamma_c at all.

[YigitElma]

Ok, #2287 will resolve the test issue.

I ran this script with different jac_chunk_size values for the objective. The issue in #2239 only happens if ForceBalance constraint is chunked. BoundaryError is one of the expensive objectives, and we usually need to chunk it to fit to the memory, so I guess this is a good representative case for what this PR can achieve.

The above plots show the memory usage for chunk sizes of [1, 10, 30, 50, 100, 150, None] and the total batch size is 291. I call the same function 6 times and each time the first shorter peak is constraint evaluation and the second bigger peak is objective evaluation. Note: no chunking could barely fit to my GPU with rematerialization, so the speed comparison is not fair there.

The memory usage for very low chunk sizes increases quite a bit, whereas the evaluation is 2 times faster. They converge to the same speed and the same memory as the chunk size decreases. I checked these on my laptop RTX4080 12GB, having similar thing on cluster A100 would help. I think we need to test this a bit more. We usually accept the slowdown for less memory usage.

[jpbrodrick89]

This is exactly what I'd expect to see, there is essentially a constant memory overhead equal to keeping the residuals in memory. For a chunk size of one this will be noticeable and its proportional controbution will decrease proportional to chunk size. I think it all depends what chunk size you need to run in production. If its typically O(1) then maybe we should gate linearising to skip when chunk size is small (i think even 1-2 will be reasonable). As i said though if you are running with chunk sizes of 1-2 already my gut feel is youd be better off with iterative Krylov solvers in those cases.

[YigitElma]

Yeah, one weird thing tho (which I realised before but worth mentioning here), sometimes doing the operation in chunks is faster than no chunking or bigger chunks. I guess this depends on the hardware you have and the vectorization capabilities. I mean, sometimes I choose a lower chunk size not for memory purposes but for speed too. I will try to check different types of problems on A100. DESC has many objectives and many options, so it is hard to say whether this is the production run. In the meantime, people are welcome to test it using my repo; it should be quite straightforward when you have a script.

[jpbrodrick89]

Im surprised that some medium chunk sizes seemed to show a slowdown but I'm guessing this is noise. Also note that for smaller chunk sizes you need more chunks so the speedup might actually be more worth it. I think the interesting thing to check is if you are running at the largest chunk size that fits on an A100 and this PR forces you to shrink that chunk size slightly whether the linearisation speedup compensates for the extra chunks required.