Multi-Host-HLO-Runner

Mit diesem Tool können Sie ein HLO-Modul auf einer oder mehreren GPUs ausführen. Außerdem kann Code für mehrere GPUs kompiliert werden, ohne dass er ausgeführt wird.

Ausführen von HLOs mit mehreren GPUs (sharded)

Diese HLOs erkennen Sie an den sharding=-Anmerkungen. sharding={devices=[1,1,2,1]0,1} bedeutet beispielsweise, dass der annotierte Tensor entlang der dritten Dimension auf zwei GPUs (GPU0 und GPU1) aufgeteilt werden soll.

Bei der folgenden Anleitung wird davon ausgegangen, dass das Arbeitsverzeichnis das XLA-Git-Repository ist und dass ./configure.py ausgeführt wurde.

Wenn wir genügend GPUs haben, können wir diese HLOs so wiedergeben:

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

Es ist auch möglich, denselben HLO zu kompilieren, ohne ihn auszuführen. Dazu müssen Sie --run=false festlegen:

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

In diesem Fall ist eine einzelne GPU erforderlich, sofern nicht der Cache für automatisches Tuning verwendet wird.

Fehlerbehebung

Fehler wie Check failed: result.replicas >= 1 (0 vs. 1):
- Wir müssen sicherstellen, dass wir genügend GPUs haben.
- CUDA_VISIBLE_DEVICES muss richtig festgelegt oder gar nicht festgelegt werden.
Abstürze:
- Möglicherweise möchten wir --dynamic_mode=off verwenden.
- CUDA und cuDNN sollten richtig eingerichtet sein.

Beispiele

Beispiel für einen einzelnen Prozess mit mehreren GPUs

HLO einrichten und abrufen

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-Multihost-Runner erstellen

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

Beispiel für einen einzelnen Prozess: Wiedergabe des Graphen vor der Optimierung

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

Beispiel für einen einzelnen Prozess: Wiedergabe des Optimierungsgraphen

Wenn Sie ein optimiertes HLO noch einmal ausführen möchten, müssen Sie entweder --xla_disable_all_hlo_passes oder --run_xla_backend_only verwenden. Andernfalls versucht XLA, das HLO neu zu kompilieren, was nicht unterstützt wird. Daher werden viele seltsame Fehler ausgegeben.

Vollständiger Befehl: 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

Mehrere Prozesse, ein Knoten

Container starten

Installieren Sie auch einige fehlende Bibliotheken. Hinweis: Diese Informationen können schnell veralten. Passen Sie die Position nach Bedarf an.)

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

Originalmodell ausführen und HLO-Datei erstellen

In diesem Beispiel verwenden wir ein PAXML-Modell mit 8 GPUs aus test-pax.sh. Hinweis: Dies ist derselbe Dump wie im Fall mit einem einzelnen Prozess. Wenn Sie cp -r /tmp/dump /tmp/dump_multi_process bereits haben, können Sie 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 nutzen.

Der HLO-Dump wird in /tmp/dump_multi_process/ gespeichert. Speziell für PAX enthält das Hauptmodul „pjit__wrapped_step_fn“ im Namen. In diesem Beispiel verwenden wir /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Auf einem einzelnen Knoten mit MPI ausführen

Erstellen Sie ein Bash-Script mit dem Namen 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

Jetzt können Sie sie mit „mpirun“ ausführen:

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

Mit SLURM auf mehreren Knoten ausführen

Wenn Sie SLURM auf mehreren Knoten ausführen, können Sie die SLURM-Umgebungsvariablen so an den HLO-Runner weiterleiten:

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