Pennylane multi-gpu script fails with error even there are enough gpus

Hi @QuantumMan

My colleague Lee has confirmed that the issue seems to be in the environment setup. There is nothing wrong with Cray MPICH and CUDA 12 on Perlmutter with Lightning GPU.

These steps produce a working Lightning GPU install (as of now 0.41.0-rc)

  • Ensure Python 3.10+ is used to create the venv (module load python/3.11)
  • Ensure the appropriate modules are loaded (module load PrgEnv-gnu cray-mpich cudatoolkit craype-accel-nvidia80)
  • Clone pennylane-lightning, install the requirements-dev.txt file and then install Lightning Qubit as python -m pip install -r requirements-dev.txt && CC=$(which cc) CXX=$(which CC) python -m pip install . --verbose) with the CC env-vars used to set the CrayPE compilers, which default to the GNU env
  • Change the package to Lightning GPU and install that using the CrayPE compiler for MPI support (CMAKE_ARGS="-DENABLE_MPI=ON" CC=$(which mpicc) CXX=$(which mpicxx) python -m pip install . --verbose)
  • Install mpi4py using CrayPE MPICH (MPICC="cc -shared" pip install --force-reinstall --no-cache-dir --no-binary=mpi4py mpi4py)
  • Ensure the Cray MPICH libraries are available for dynamic loading by NVIDIA custatevec (export LD_LIBRARY_PATH=$CRAY_LD_LIBRARY_PATH:$LD_LIBRARY_PATH)
  • Allocate a debug session with 4 GPUs (salloc -N 1 -c 32 --qos interactive --time 0:30:00 --constraint gpu --ntasks-per-node=4 --gpus-per-task=1 --gpu-bind=none --account=XYZ)
  • Create an MPI-friendly script and run it with the above allocation (srun -n 4 python myscript.py)

Note: we will release a new version of PennyLane in the next couple of days so if you want to wait until Wednesday to do this you can use PennyLane v0.41