使用 GPU 的 Dask-MPI
目录
使用 GPU 的 Dask-MPI¶
在支持 GPU 的系统上运行 dask-mpi 时，每个 MPI 秩将提供一个或多个 GPU。
目前 Dask 假设每个 GPU 对应一个 worker 进程，并且 workers 正确地绑定到每个 GPU。为了实现这一点，存在 dask-cuda 包，其中包含设计用于正确配置 GPU 环境的 cluster 和 worker 类。
conda install -c rapidsai -c nvidia -c conda-forge dask-cuda
# or
python -m pip install dask-cuda
通过将 worker 类设置为使用 dask_cuda.CUDAWorker，可以结合使用 dask-cuda 和 dask-mpi。
mpirun -np 4 dask-mpi --worker-class dask_cuda.CUDAWorker
from dask_mpi import initialize
initialize(worker_class="dask_cuda.CUDAWorker")
提示
如果您的集群配置为每个秩代表一个节点，那么每个节点可能有多个 GPU。Workers 是按 GPU 创建的，而不是按秩创建的，因此 CUDAWorker 将为每个 GPU 创建一个 worker，其名称遵循模式 {rank}-{gpu_index}。因此，如果您设置了 -np 4，但每个节点有四个 GPU，最终您的集群中将有十六个 workers。
附加配置¶
除了 dask-mpi 支持的选项外，您可能还希望向 dask_cuda.CUDAWorker 传递附加配置选项。通常会为 GPU workers 配置内存管理和网络协议等事项。
您可以使用 worker 选项参数传递 dask_cuda.CUDAWorker 接受的任何附加选项。
在 CLI 上，这预计是一个 JSON 序列化的值字典。
mpirun -np 4 dask-mpi --worker-class dask_cuda.CUDAWorker --worker-options '{"rmm_managed_memory": true}'
在 Python 中，它只是一个字典。
from dask_mpi import initialize
initialize(worker_class="dask_cuda.CUDAWorker", worker_options={"rmm_managed_memory": True})
提示
有关将 GPU 与 Dask 结合使用的更多信息，请查阅 dask-cuda 文档。