Многохостовый HLO Runner

Этот инструмент позволяет запускать модуль HLO на одном или нескольких графических процессорах. Он также позволяет компилировать код, предназначенный для нескольких графических процессоров, без его запуска.

Запуск многопроцессорных (шардированных) HLO

Мы можем идентифицировать эти HLO по аннотациям sharding= . Например, sharding={devices=[1,1,2,1]0,1} означает, что аннотированный тензор должен быть разбит на 2 графических процессора (GPU0 и GPU1) по третьему измерению.

Следующие инструкции предполагают, что рабочая директория — это репозиторий XLA Git и что файл ./configure.py был запущен.

Если у нас достаточно графических процессоров, мы можем воспроизвести эти HLO следующим образом:

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- my-hlo.txt

Также можно скомпилировать тот же HLO без его запуска, установив параметр --compile_only=true :

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- --compile_only=true my-hlo.txt

В этом случае потребуется всего один графический процессор, если только не используется кэш автонастройки .

Поиск неисправностей

  • Ошибки, такие как Check failed: result.replicas >= 1 (0 vs. 1) :
    • Мы должны убедиться, что у нас достаточно графических процессоров.
    • CUDA_VISIBLE_DEVICES должен быть установлен правильно или не должен быть установлен вообще.
  • Сбои:
    • Возможно, нам стоит использовать --dynamic_mode=off .
    • Необходимо правильно настроить CUDA и cuDNN.

Примеры

Пример работы в однопроцессной среде с несколькими графическими процессорами.

Настройте и получите HLO. 

You can use a container with the following instructions:

  docker run -it --shm-size=1g --gpus all ghcr.io/nvidia/jax:pax-2024-06-03
  cd /opt/xla/

Note, those instructions can be outdated more quickly. Adjust as needed.

# The 8 below is the number of GPUs you have.
# test-pax.sh --help for more details on the parallelization options
(export XLA_FLAGS="--xla_dump_to=/tmp/dump"; test-pax.sh --fsdp 8 --batch-per-gpu 1)

ls -lSh /tmp/dump/*before_optimizations.txt
# The biggest file one is normally the one you care about.
# I picked one, for the rest of the scripts, but the name could change when you change the JAX or XLA version.

Создание многохостового исполнителя XLA 

cd /opt/xla/
./configure.py --backend CUDA --nccl
bazel build //xla/tools/multihost_hlo_runner:hlo_runner_main

Пример отдельного процесса: Воспроизведение графа до оптимизации 

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \
  /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

Пример отдельного процесса: Воспроизведение графа после оптимизации

Для воспроизведения оптимизированного HLO необходимо использовать либо --xla_disable_all_hlo_passes , либо --run_xla_backend_only . В противном случае XLA попытается перекомпилировать HLO, что не поддерживается. В результате вы получите множество странных ошибок.

Полная команда: bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run_xla_backend_only /tmp/dump/module_0023.pjit__wrapped_step_fn.sm_8.0_gpu_after_optimizations.txt

Многопроцессная, одноузловая архитектура

Запустить контейнер

Также установите некоторые недостающие библиотеки. (Обратите внимание, что они могут быстро устаревать. Внесите необходимые корректировки.)

docker run -it --shm-size=1g --gpus all ghcr.io/nvidia/jax:pax-2024-06-03
apt-get update && apt-get install -y openmpi-bin openmpi-common libopenmpi-dev

Запустите исходную модель и сделайте дамп HLO.

В этом примере мы будем использовать модель PAXML с 8 графическими процессорами из test-pax.sh . (Обратите внимание, что это будет тот же дамп, что и в случае с одним процессом. Поэтому вы можете выполнить cp -r /tmp/dump /tmp/dump_multi_process , если она у вас уже есть. export XLA_FLAGS="--xla_dump_to=/tmp/dump_multi_process" mpirun --allow-run-as-root -np 8 test-pax.sh --fsdp 8 --batch-per-gpu 1 -o /tmp/checkpoint --multiprocess

Дамп HLO будет сохранен в /tmp/dump_multi_process/ . В частности, для PAX основной модуль будет содержать в своем имени "pjit__wrapped_step_fn". В этом примере мы будем использовать /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt .

Запуск на одном узле с использованием MPI.

Создайте bash-скрипт с именем run.sh :

#!/bin/bash
export CUDA_VISIBLE_DEVICES=${OMPI_COMM_WORLD_LOCAL_RANK}
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \
  --task_id=${OMPI_COMM_WORLD_RANK} \
  --num_nodes=${OMPI_COMM_WORLD_SIZE} \
  --address=127.0.0.1:12345 \
  /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

Теперь вы можете запустить его с помощью mpirun:

chmod a+x run.sh
mpirun --allow-run-as-root -np 8 run.sh

Запуск на нескольких узлах с использованием SLURM.

При запуске SLURM на нескольких узлах вы можете передавать переменные окружения SLURM исполнителю HLO следующим образом в вашем задании SLURM: 

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \
  --task_id=${SLURM_PROCID} \
  --num_nodes=${SLURM_NTASKS} \
  --address="${SLURM_LAUNCH_NODE_IPADDR}:12345" \
  /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt