PlasmaControl / DESC

Stellarator Equilibrium and Optimization Suite
MIT License
97 stars 26 forks source link

Update the installation instructions for Perlmutter and RAVEN #1324

Open YigitElma opened 4 weeks ago

YigitElma commented 4 weeks ago
          @rahulgaur104 will update the instructions for Perlmutter and RAVEN in a later PR.

Originally posted by @YigitElma in https://github.com/PlasmaControl/DESC/issues/1293#issuecomment-2436078532

Since they use an older JAX version 0.4.23, current instructions require to change scipy version from requirements file for proper installation. If newer versions of JAX are supported for these clusters, the scipy version change is redundant. JAX should be greater than 0.4.24 ( also this is the new requirement for DESC).

dpanici commented 3 weeks ago

1325 @ddudt we should add an example submission script to installation (in FAQS maybe) or as a separate page in docs

daniel-dudt commented 3 weeks ago

Response from NERSC support about installing the latest versions of JAX on Perlmutter:

I would highly recommend that you use containers to run JAX at NERSC, this make your life significantly easier down the line (especially as you start distributing your computation over several nodes): https://docs.nersc.gov/development/languages/python/using-python-perlmutter/#nvidia-containers

Otherwise, to use a nightly release of JAX (versions above 0.4.29) following our documented approach, you can use this page: https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

I have also been told that, with the latest releases, a simple pip install --upgrade "jax[cuda12]" might work on Perlmutter but I have yet to test it (at which point I will update the documentation).

rahulgaur104 commented 2 weeks ago

@daniel-dudt Were you able to install a recent version on jax + DESC on Perlmutter? If so, can you post your instructions here so I can double-check.