Эта страница переведена с помощью Cloud Translation API.

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

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

Запуск многочиповых (сегментированных) HLO

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

В следующих инструкциях предполагается, что рабочим каталогом является репозиторий xla Git и запущен файл ./configure.py .

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

bazel run -c opt --config=cuda --dynamic_mode=off \
  //xla/tools/multihost_hlo_runner:hlo_runner_main -- my-hlo.txt

Также можно скомпилировать тот же HLO, не запуская его, установив --run=false

bazel run -c opt --config=cuda --dynamic_mode=off \
  //xla/tools/multihost_hlo_runner:hlo_runner_main \
  -- --run=false 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 --xla_dump_hlo_as_text"; 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 -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main

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

bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

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

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

Полная команда: bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run_xla_backend_only=true --xla_disable_all_hlo_passes=true /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 --xla_dump_hlo_as_text" mpirun --allow-run-as-root -np 8 test-pax.sh --fsdp 8 --batch-per-gpu 1 -o /tmp/checkpoint --multiprocess он у вас уже есть. export XLA_FLAGS="--xla_dump_to=/tmp/dump_multi_process --xla_dump_hlo_as_text" 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 -c opt --config=cuda --dynamic_mode=off //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 вы можете перенаправить переменные env SLURM в программу hlo, например, в своем задании slurm:

bazel run -c opt --config=cuda --dynamic_mode=off //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