Implement multi-gpu functionality for separate objectives with MPI4Jax

PlasmaControl / DESC

Stellarator Equilibrium and Optimization Suite

MIT License

87 stars 22 forks source link

Open kianorr opened 2 months ago

kianorr commented 2 months ago

The general idea is

Physics objectives are put onto separate GPUs
Constraints are taken out
Then combine jacobians on a single GPU to create the $A$ matrix, and perform SVD on this $A$ matrix
From this a new eq is created, which is then fed back to each GPU

dpanici commented 1 month ago

Add view-only link to hackathon doc

kianorr commented 3 weeks ago

instructions to install on della are under "mpi4jax installation instructions on della-gpu" in this google doc: https://docs.google.com/document/d/1x6nGZEiZnAiWBDf20Mcbwbob9a6GMbx3ZpT_huUVZF4/edit?usp=sharing