To narzędzie umożliwia uruchamianie modułu HLO na co najmniej 1 procesorze graficznym. Pozwala też kompilować kod przeznaczony na wiele procesorów graficznych bez jego uruchamiania.

Uruchamianie HLO na wielu procesorach graficznych (z podziałem)

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

W tych instrukcjach przyjęto, że katalog roboczy to repozytorium Git XLA i że uruchomiono skrypt ./configure.py.

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

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

Można też skompilować ten sam HLO bez jego uruchamiania, ustawiając flagę --compile_only=true:

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

W takim przypadku wymagany jest 1 procesor graficzny, chyba że używany jest pamięć podręczna autotuningu.

Rozwiązywanie problemów

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

Przykłady

Przykład z jednym procesem i wieloma procesorami graficznymi

Konfiguracja i pobieranie 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 XLA multihost runner

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

Przykład z jednym procesem: odtwarzanie wykresu przed optymalizacją

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

Przykład z jednym procesem: odtwarzanie wykresu po optymalizacji

Aby odtworzyć zoptymalizowany HLO, musisz użyć flagi --xla_disable_all_hlo_passes lub --run_xla_backend_only. W przeciwnym razie XLA spróbuje ponownie skompilować HLO, co nie jest obsługiwane. Spowoduje to 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ł

Uruchamianie kontenera

Zainstaluj też brakujące biblioteki. (Pamiętaj, że mogą one 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 zrzutu HLO

W tym przykładzie użyjemy modelu PAXML z 8 procesorami graficznymi z pliku test-pax.sh. (Pamiętaj, że będzie to ten sam zrzut co w przypadku jednego procesu. Jeśli masz już zrzut, możesz użyć polecenia 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

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 pliku /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Uruchamianie w 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

Możesz go teraz uruchomić 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 uruchamiasz zadanie na wielu węzłach za pomocą SLURM, możesz przekazać zmienne środowiskowe SLURM do narzędzia HLO runner w ten sposób:

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