Вопрос о запуске MPI и Horovod в ядре Jupyter Python

Question

Я хочу поэкспериментировать с ноутбуком, на котором работает хоровод, распределенный по трем узлам HPC, каждый с одним графическим процессором. Я загружаю эти модули в определение своего ядра:

"module load shared slurm jupyter-eg-kernel-wlm-py39 horovod-tensorflow2-py39-cuda11.2-gcc9/0.22.1 nccl2-cuda11.2-gcc9/2.14.3 tensorflow2-py39-cuda11.2-gcc9/2.7.0 openmpi4-cuda11.2-ofed51-gcc9"

Затем ядро ​​запускается с этой отправкой SLURM:

      "submit_script": [
          "#!/bin/bash",
          "#SBATCH --nodes=3",
          "#SBATCH --ntasks=3",
          "#SBATCH --ntasks-per-node=1",
          "#SBATCH --gres=gpu:1",
          "#SBATCH --mem=18G",
          "#SBATCH --export=NONE",
          "#SBATCH --job-name=jupyter-eg-kernel-slurm-py39-{kernel_id}",

...и т. д

Итак, все запускается нормально, и я вижу три узла, выделенных для трех моих задач, но когда я запускаюhvd.init()а потомhvd.size()в ячейке блокнота только печатает1. Разве я не должен видеть размер «3»? Различные онлайн-примеры загрузки набора данных MNIST и обучения модели работают нормально, но только на одном из графических процессоров одного из узлов. Я всегда получаю это сообщение, когда запускаю образец (с помощью nvidia-smi я вижу графические процессоры на всех узлах):

I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1631 MB memory: -> device: 0, name: Quadro P600, pci bus id: 0000:07:00.0, compute capability: 6.1