Parallelization via sharding (multi-device) composed with vmap (across memory in single device)

[unalmis]

Robust sharding for factor of N improvement in compute and memory

This is the continuation of #1869. Broadly, there are two methods of parallelizing computation.

1. The first is to split the computation and parallelize across the memory of a single device. This is ideal if the problem can fit on a single device and if the device has sufficient processing power to complete the computation in parallel.

2. The second is to split the computation and parallelize across multiple devices, which is useful when the problem size becomes too large for the first strategy.

In JAX, option 1 is done with vmap. When the problem size becomes larger, we mimic the benefits of option 2 with vmap_chunked. With vmap_chunked we choose the chunk_size such that the italic part of option 1 is best satisfied and then compute on each chunk in sequence.

This PR composes multiple device parallelization with vmap_chunked to improve compute ability.

With the current JAX sharding API, it is most practical to first split the data S across n devices, then on each subset of data of size S/n we may chunk with size m across the memory of a single device with vmap_chunked. So each device is computing on a subset of size S/n/m simultaneously. These API require S to be divisible by n and m, but it is just some simple extra logic to compute the remainder of S/n/m on each device after they finish the divisible portion, and likewise compute the remainder of S/n on a single device after that.

The net result will be reducing the compute time. This also ensures that if there was a memory bottlneck in constructing the Jacobian, then this bottleneck will be pushed away toward storing the Jacobian. Then that can later be avoided by using iterative subspace methods so that only a subset of the Jacobian needs to be stored at a time.

If the use case is not clear, then explicitly, this would enable parallelization without rebuilding an objective for each flux surface/duplcating computation and transforms.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    4.18 %    |     3.934e+03      |     4.099e+03      |    164.50    |       32.24        |       29.44        |
  test_proximal_jac_w7x_with_eq_update   |   -1.11 %    |     6.569e+03      |     6.496e+03      |    -72.68    |       149.88       |       151.21       |
  test_proximal_freeb_jac                |   -0.06 %    |     1.338e+04      |     1.338e+04      |    -8.18     |       86.31        |       85.73        |
  test_proximal_freeb_jac_blocked        |   -0.40 %    |     7.731e+03      |     7.700e+03      |    -31.23    |       76.19        |       74.90        |
  test_proximal_freeb_jac_batched        |    0.56 %    |     7.623e+03      |     7.666e+03      |    42.99     |       75.00        |       75.59        |
  test_proximal_jac_ripple               |   -0.34 %    |     3.560e+03      |     3.548e+03      |    -12.26    |       53.76        |       52.93        |
+ test_proximal_jac_ripple_bounce1d      |   -20.66 %   |     2.056e+02      |     1.631e+02      |    -42.48    |        4.11        |        4.11        |
  test_eq_solve                          |   -0.60 %    |     2.081e+03      |     2.069e+03      |    -12.48    |       82.11        |       81.71        |

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         |     -1.03 +/- 2.90     | -5.93e-03 +/- 1.66e-02 |  5.68e-01 +/- 1.2e-02  |  5.74e-01 +/- 1.2e-02  |
 test_equilibrium_init_medres            |     -0.82 +/- 2.17     | -4.17e-02 +/- 1.10e-01 |  5.02e+00 +/- 7.0e-02  |  5.06e+00 +/- 8.5e-02  |
 test_equilibrium_init_highres           |     -0.34 +/- 1.52     | -1.94e-02 +/- 8.54e-02 |  5.61e+00 +/- 6.0e-02  |  5.63e+00 +/- 6.1e-02  |
 test_objective_compile_dshape_current   |     +0.59 +/- 3.17     | +1.99e-02 +/- 1.07e-01 |  3.40e+00 +/- 9.6e-02  |  3.38e+00 +/- 4.7e-02  |
 test_objective_compute_dshape_current   |     -4.53 +/- 6.85     | -3.51e-05 +/- 5.30e-05 |  7.39e-04 +/- 3.5e-05  |  7.74e-04 +/- 3.9e-05  |
 test_objective_jac_dshape_current       |     -1.16 +/- 18.84    | -3.81e-04 +/- 6.21e-03 |  3.26e-02 +/- 4.1e-03  |  3.30e-02 +/- 4.7e-03  |
 test_perturb_2                          |     +0.60 +/- 1.67     | +1.02e-01 +/- 2.85e-01 |  1.71e+01 +/- 1.1e-01  |  1.70e+01 +/- 2.6e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.70 +/- 0.89     | +9.45e-02 +/- 1.21e-01 |  1.37e+01 +/- 9.8e-02  |  1.36e+01 +/- 7.1e-02  |
 test_proximal_freeb_jac                 |     +0.46 +/- 1.81     | +2.32e-02 +/- 9.02e-02 |  5.02e+00 +/- 6.6e-02  |  5.00e+00 +/- 6.2e-02  |
 test_solve_fixed_iter_compiled          |     +0.31 +/- 1.80     | +5.34e-02 +/- 3.08e-01 |  1.72e+01 +/- 2.7e-01  |  1.71e+01 +/- 1.5e-01  |
 test_LinearConstraintProjection_build   |     -1.08 +/- 2.23     | -9.17e-02 +/- 1.90e-01 |  8.42e+00 +/- 1.2e-01  |  8.51e+00 +/- 1.5e-01  |
 test_objective_compute_ripple_spline    |     +1.20 +/- 3.65     | +3.53e-03 +/- 1.07e-02 |  2.98e-01 +/- 8.6e-03  |  2.94e-01 +/- 6.5e-03  |
 test_objective_grad_ripple_spline       |     -0.41 +/- 3.84     | -4.59e-03 +/- 4.29e-02 |  1.11e+00 +/- 3.8e-02  |  1.12e+00 +/- 1.9e-02  |

[codecov]

Codecov Report

❌ Patch coverage is 30.43478% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.44%. Comparing base (58768e7) to head (72293f1).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/batching.py	30.43%	16 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1773      +/-   ##
==========================================
- Coverage   94.48%   94.44%   -0.05%     
==========================================
  Files         102      102              
  Lines       28721    28741      +20     
==========================================
+ Hits        27137    27143       +6     
- Misses       1584     1598      +14     

Files with missing lines	Coverage Δ
desc/batching.py	75.93% <30.43%> (-6.11%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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