Using multiple GPUs

Using multiple GPUs#

Learning how to make use of the multi-GPU capabilities in PtyPy.

All GPU-accelerated examples so far were using only a single GPU working on the reconstructions, but PtyPy supports multi-GPU execution for all its engines via MPI. Unfortunately, this is not compatible with Jupyter notebooks, but we can instead use the shell terminal provided within JupyterHub.

Tip

When running PtyPy with a Python script in a compute environment that uses OpenMPI, e.g. as part of a cluster job, the multi-GPU capabilities of engines such as DM, RAAR or ML can be utilised with mpirun -n 4 python ptypy_example.py

Making a Python run script for PtyPy#

Revisiting the earlier example on using JSON/YAML Config Files we have saved the parameter tree from the previous example on Working With Large Data into a config file named ./config/dls_i08_nanogold_spiral.yaml making a few changes, i.e. switching to verbose_level=info and providing p.io.rfile for saving the output of the reconstruction. We can then generate a simple Python script

with open("./ptypy_run_dls_i08_nanogold.py", "w") as f:
    f.write("""\
import ptypy
ptypy.load_ptyscan_module("hdf5_loader")
ptypy.load_gpu_engines("cuda")  
p = ptypy.utils.param_from_yaml("config/dls_i08_nanogold_spiral.yaml")
P = ptypy.core.Ptycho(p,level=5)
""")

which imports the relevant PtyScan and GPU acceleration modules, loads the parameter tree and runs the reconstruction.

with open("./ptypy_run_dls_i08_nanogold.py", "w") as f:
    f.write("""\
import ptypy
ptypy.load_ptyscan_module("hdf5_loader")
ptypy.load_gpu_engines("cupy")  
p = ptypy.utils.param_from_yaml("config/dls_i08_nanogold_spiral.yaml")
P = ptypy.core.Ptycho(p,level=5)
""")

Executing script in terminal on JupyterHub#

Now, we need to open a new terminal (click blue “+” in top left corner of JupyterHub and then “Terminal”) which should look something like this

(PtyPy) benedikt@nid001848:/global/u2/b/benedikt>

with the (PtyPy) environment already loaded. We can then navigate to the folder for this example

cd $HOME/tutorials/notebooks/experimental_xray_data

and execute a multi-GPU reconstruction via srun using all 4 GPUs that is available on a Perlmutter node.

srun -n 4 -c 2 --gpus-per-task=1 --gpu-bin=None python ptypy_run_dls_i08_nanogold.py

Challenge

Run the I08 nanogold example in a terminal using 4 GPUs and compare the speed against the previous example from the Jupyter notebook which was using just a single GPU.

Display reconstruction from ptyr file#

Once the multi-GPU reconstruction has finished, we can switch back to this notebook and display the final output using

import ptypy.utils.plot_client as pc
fig = pc.figure_from_ptyr("./recons/dls_i08_nanogold_DM_cupy_0200.ptyr")

import ptypy.utils.plot_client as pc
fig = pc.figure_from_ptyr("./recons/dls_i08_nanogold_DM_cupy_0200.ptyr")
Contents