zoomy_jax.gpurank

zoomy_jax.gpurank#

One GPU per MPI rank, bound before anything loads CUDA.

mpiexec -n 4 python -m zoomy_jax.gpurank run.py [args ...]

Sets CUDA_VISIBLE_DEVICES from the node-local rank (round robin over the node’s GPUs), then runs the script. Every CUDA user in the process – JAX, AmgX, mpi4jax’s copies – sees only that GPU. A launcher that binds GPUs itself (srun --gpus-per-task=1, an existing CUDA_VISIBLE_DEVICES) is left alone. Binding later, inside the run, is too late for libraries that use the CUDA runtime directly: import jax already loads CUDA.

Module Contents#

Functions#

bind

Set CUDA_VISIBLE_DEVICES for this rank; True if it did.

main

API#

zoomy_jax.gpurank.bind()#

Set CUDA_VISIBLE_DEVICES for this rank; True if it did.

zoomy_jax.gpurank.main()#
Contents