Uruchamiający HLO na wielu hostach

To narzędzie umożliwia uruchamianie modułu HLO na co najmniej 1 procesorze GPU. Umożliwia też kompilowanie kodu przeznaczonego dla wielu procesorów graficznych bez jego uruchamiania.

Uruchamianie HLO z wieloma GPU (podzielonych na fragmenty)

Możemy je rozpoznać po adnotacjach sharding=. Na przykład sharding={devices=[1,1,2,1]0,1} oznacza, że tensor z adnotacjami powinien być podzielony na 2 procesory graficzne (GPU0 i GPU1) wzdłuż 3 wymiaru.

W tych instrukcjach przyjęto założenie, że katalogiem roboczym jest repozytorium Git XLA i że wykonano polecenie ./configure.py.

Jeśli mamy wystarczającą liczbę procesorów GPU, możemy odtworzyć te HLO w ten sposób:

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

Możesz też skompilować ten sam HLO bez uruchamiania go, ustawiając --run=false:

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

W takim przypadku wymagany jest jeden procesor graficzny, chyba że używana jest pamięć podręczna automatycznego dostrajania.

Rozwiązywanie problemów

  • Błędy takie jak Check failed: result.replicas >= 1 (0 vs. 1):
    • Musimy mieć wystarczającą liczbę procesorów graficznych.
    • Wartość CUDA_VISIBLE_DEVICES musi być ustawiona prawidłowo lub nie może być ustawiona wcale.
  • Awarie:
    • Możemy użyć --dynamic_mode=off.
    • CUDA i cuDNN powinny być prawidłowo skonfigurowane.

Przykłady

Przykład pojedynczego procesu z wieloma procesorami graficznymi

Konfigurowanie i uzyskiwanie 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.

Tworzenie narzędzia do uruchamiania XLA na wielu hostach

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

Przykład pojedynczego procesu: przed odtworzeniem wykresu optymalizacji

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

Przykład pojedynczego procesu: odtwarzanie wykresu po optymalizacji

Aby odtworzyć zoptymalizowany HLO, musisz użyć --xla_disable_all_hlo_passes lub --run_xla_backend_only. W przeciwnym razie XLA spróbuje ponownie skompilować HLO, co nie jest obsługiwane. Dlatego może wyświetlać wiele dziwnych błędów.

Pełne polecenie: 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

Wiele procesów, jeden węzeł

Uruchom kontener

Zainstaluj też brakujące biblioteki. (Pamiętaj, że może to szybko stać się nieaktualne. W razie potrzeby dostosuj położenie urządzenia.

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

Uruchamianie oryginalnego modelu i zrzucanie HLO

W tym przykładzie użyjemy modelu PAXML z 8 procesorami GPU z test-pax.sh. (Uwaga: będzie to ten sam zrzut co w przypadku pojedynczego procesu. Możesz to zrobić cp -r /tmp/dump /tmp/dump_multi_process, jeśli masz już taką usługę. 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

Zrzut HLO zostanie zapisany w katalogu /tmp/dump_multi_process/. W przypadku PAX główny moduł będzie miał w nazwie ciąg „pjit__wrapped_step_fn”. W tym przykładzie użyjemy adresu /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Uruchamianie na jednym węźle za pomocą MPI

Utwórz skrypt bash o nazwie 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

Teraz możesz go wykonać za pomocą polecenia mpirun:

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

Uruchamianie na wielu węzłach za pomocą SLURM

Jeśli używasz wielu węzłów z SLURM, możesz przekazać zmienne środowiskowe SLURM do narzędzia uruchamiającego HLO w ten sposób w zadaniu 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