Support for Multiple Devices and Parallelism via MPI

[YigitElma]

Initial support for multi-device optimization.

• Updates set_device and desc.backend for initialization of multiple devices. These are the cause of slightly low coverage (around 90% for patch)
• Adds _device_id and _device to _Objective class (defaults to 0) for making parallelization work with other objectives
• Uses mpi4py for compute and jvp methods of parallel objective functions (there is no separate class but if user passes objectives with different devices and an mpi communicator, parallel strategy will be used.) To make this work a user has to use context manager properly
• Adds rank_per_objective to ObjectiveFunction which is an array that specifies which objective should run on which rank. There are some limitations for the user. For example, 2 different rank cannot use the same device (maybe it can, but it shouldn't in practice). One tricky part about this is, for multiple node case, same device_id can mean different devices. The reason we have rank_per_objective is to be able to distinguish that (MPI.size is equal to total number of devices in all nodes, whereas max(device_id) can be only the number of GPUs connected to a single node)

Things to consider in future PRs,

• Maybe implement a new optimizer that uses distributed matrix operations instead of QR and SVD.
• Find good distributed linear algebra routines for SVD and QR (necessary for speed and low memory usage since the current implementation will use CPU if the Jacobian is big)
• Use CUDA-aware MPI (possibly without mpi4jax)
• Incorporate sharding if possible?
• Implement MPI version of Augmented Lagrangian optimizers. Currently, they are forced to run on GPU0. We can create a flag in current optimizers or write a new optimizer specific to MPI, such that all objective and non-linear constraints are passed inside the main ObjectiveFunction, then which objectives need to be treated as constraints are passed via a flag. Or we can do a bigger structural change to achieve this, but I think putting everything in a single objective and passing to a custom MPI version is a lot easier than some cumbersome logic.

TODOS:

• Add rank to each objective and remove rank_per_objective argument of ObjectiveFunction
• Make sure the API won't change for the support for constraints

Resolves #1071 (but with mpi4py)
Resolves #1601

[codecov]

Codecov Report

❌ Patch coverage is 91.49798% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.44%. Comparing base (24aa7b9) to head (8dc5792).

Files with missing lines	Patch %	Lines
desc/backend.py	51.85%	13 Missing ⚠️
desc/objectives/objective_funs.py	97.19%	5 Missing ⚠️
desc/optimize/_constraint_wrappers.py	92.85%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1495      +/-   ##
==========================================
+ Coverage   94.32%   94.44%   +0.11%     
==========================================
  Files         101      100       -1     
  Lines       28866    29008     +142     
==========================================
+ Hits        27228    27396     +168     
+ Misses       1638     1612      -26     

Files with missing lines	Coverage Δ
desc/objectives/_bootstrap.py	98.68% <ø> (ø)
desc/objectives/_coils.py	99.46% <ø> (ø)
desc/objectives/_equilibrium.py	96.63% <ø> (+0.42%)	⬆️
desc/objectives/_fast_ion.py	95.06% <ø> (ø)
desc/objectives/_free_boundary.py	96.47% <ø> (ø)
desc/objectives/_generic.py	99.58% <ø> (ø)
desc/objectives/_geometry.py	96.96% <ø> (ø)
desc/objectives/_neoclassical.py	94.87% <ø> (ø)
desc/objectives/_omnigenity.py	97.01% <ø> (ø)
desc/objectives/_power_balance.py	91.75% <ø> (ø)
... and 5 more

... 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.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.62 +/- 4.64     | -1.44e-02 +/- 4.11e-02 |  8.72e-01 +/- 1.7e-02  |  8.86e-01 +/- 3.7e-02  |
 test_equilibrium_init_medres            |     -0.48 +/- 3.00     | -3.37e-02 +/- 2.11e-01 |  7.02e+00 +/- 1.5e-01  |  7.06e+00 +/- 1.5e-01  |
 test_equilibrium_init_highres           |     -2.14 +/- 2.60     | -1.71e-01 +/- 2.09e-01 |  7.85e+00 +/- 1.3e-01  |  8.02e+00 +/- 1.6e-01  |
 test_objective_compile_dshape_current   |     -0.85 +/- 2.21     | -3.71e-02 +/- 9.68e-02 |  4.34e+00 +/- 6.1e-02  |  4.38e+00 +/- 7.5e-02  |
 test_objective_compute_dshape_current   |     +4.58 +/- 12.97    | +3.12e-05 +/- 8.85e-05 |  7.14e-04 +/- 6.8e-05  |  6.82e-04 +/- 5.6e-05  |
 test_objective_jac_dshape_current       |     +3.66 +/- 21.06    | +8.82e-04 +/- 5.07e-03 |  2.50e-02 +/- 3.5e-03  |  2.41e-02 +/- 3.7e-03  |
 test_perturb_2                          |     -1.37 +/- 1.57     | -2.83e-01 +/- 3.24e-01 |  2.04e+01 +/- 3.2e-01  |  2.07e+01 +/- 4.6e-02  |
 test_proximal_jac_atf_with_eq_update    |     +0.13 +/- 2.42     | +1.59e-02 +/- 2.94e-01 |  1.21e+01 +/- 2.9e-01  |  1.21e+01 +/- 5.8e-02  |
 test_proximal_freeb_jac                 |     -0.73 +/- 2.52     | -3.52e-02 +/- 1.21e-01 |  4.78e+00 +/- 4.4e-02  |  4.81e+00 +/- 1.1e-01  |
 test_solve_fixed_iter_compiled          |     -1.34 +/- 1.91     | -8.57e-02 +/- 1.23e-01 |  6.33e+00 +/- 1.1e-01  |  6.41e+00 +/- 4.7e-02  |
 test_LinearConstraintProjection_build   |     -1.53 +/- 3.67     | -1.10e-01 +/- 2.64e-01 |  7.09e+00 +/- 2.2e-01  |  7.20e+00 +/- 1.5e-01  |
 test_objective_compute_ripple_bounce1d  |     +0.92 +/- 5.35     | +2.82e-03 +/- 1.63e-02 |  3.08e-01 +/- 1.3e-02  |  3.05e-01 +/- 9.3e-03  |
 test_objective_grad_ripple_bounce1d     |     +1.81 +/- 1.77     | +1.68e-02 +/- 1.64e-02 |  9.44e-01 +/- 1.4e-02  |  9.28e-01 +/- 8.9e-03  |
 test_build_transform_fft_midres         |     +0.30 +/- 3.94     | +2.63e-03 +/- 3.47e-02 |  8.84e-01 +/- 1.5e-02  |  8.81e-01 +/- 3.1e-02  |
 test_build_transform_fft_highres        |     -0.73 +/- 3.20     | -8.63e-03 +/- 3.78e-02 |  1.17e+00 +/- 2.6e-02  |  1.18e+00 +/- 2.7e-02  |
 test_equilibrium_init_lowres            |     +0.96 +/- 3.01     | +6.22e-02 +/- 1.95e-01 |  6.53e+00 +/- 1.5e-01  |  6.47e+00 +/- 1.2e-01  |
 test_objective_compile_atf              |     +1.78 +/- 3.76     | +1.11e-01 +/- 2.34e-01 |  6.33e+00 +/- 2.1e-01  |  6.22e+00 +/- 9.5e-02  |
 test_objective_compute_atf              |     -0.60 +/- 11.24    | -1.26e-05 +/- 2.36e-04 |  2.09e-03 +/- 1.8e-04  |  2.10e-03 +/- 1.5e-04  |
 test_objective_jac_atf                  |     -1.15 +/- 4.80     | -1.77e-02 +/- 7.40e-02 |  1.52e+00 +/- 6.6e-02  |  1.54e+00 +/- 3.3e-02  |
 test_perturb_1                          |     +0.32 +/- 1.83     | +5.13e-02 +/- 2.91e-01 |  1.59e+01 +/- 2.2e-01  |  1.59e+01 +/- 1.9e-01  |
 test_proximal_jac_atf                   |     -0.01 +/- 1.06     | -3.07e-04 +/- 5.50e-02 |  5.18e+00 +/- 4.1e-02  |  5.18e+00 +/- 3.6e-02  |
 test_proximal_freeb_compute             |     +0.94 +/- 2.74     | +1.55e-03 +/- 4.52e-03 |  1.66e-01 +/- 3.2e-03  |  1.65e-01 +/- 3.2e-03  |
 test_solve_fixed_iter                   |     +0.92 +/- 2.46     | +2.23e-01 +/- 5.98e-01 |  2.45e+01 +/- 3.5e-01  |  2.43e+01 +/- 4.8e-01  |
 test_objective_compute_ripple           |     +1.91 +/- 5.21     | +4.40e-03 +/- 1.20e-02 |  2.35e-01 +/- 6.6e-03  |  2.30e-01 +/- 1.0e-02  |
 test_objective_grad_ripple              |     +1.52 +/- 2.33     | +1.31e-02 +/- 2.02e-02 |  8.80e-01 +/- 1.3e-02  |  8.67e-01 +/- 1.5e-02  |

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

[dpanici]

#763
check for overlap with this one

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[ddudt]

I tested this on multiple nodes and it works as expected. It wasn't very slow with 2 nodes on Perlmutter.