Questo strumento consente di eseguire un modulo HLO su una o più GPU. Consente inoltre di compilare il codice destinato a più GPU senza eseguirlo.

Esecuzione di HLO multi-GPU (con partizionamento)

Possiamo identificare questi HLO dalle annotazioni sharding=. Ad esempio, sharding={devices=[1,1,2,1]0,1} significa che il tensore annotato deve essere partizionato su 2 GPU (GPU0 e GPU1) lungo la terza dimensione.

Le seguenti istruzioni presuppongono che la directory di lavoro sia il repository Git XLA e che sia stato eseguito ./configure.py.

Se abbiamo un numero sufficiente di GPU, possiamo riprodurre questi HLO nel seguente modo:

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

È anche possibile compilare lo stesso HLO senza eseguirlo impostando --compile_only=true:

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

In questo caso, è necessaria una singola GPU, a meno che non venga utilizzata la cache di ottimizzazione automatica.

Risoluzione dei problemi

• Errori come Check failed: result.replicas >= 1 (0 vs. 1):
  • Dobbiamo assicurarci di avere un numero sufficiente di GPU.
  • CUDA_VISIBLE_DEVICES deve essere impostato correttamente o non impostato.
• Arresti anomali:
  • Potrebbe essere necessario utilizzare --dynamic_mode=off.
  • CUDA e cuDNN devono essere configurati correttamente.

Esempi

Esempio di più GPU in un singolo processo

Configurazione e recupero dell'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.

Creazione dell'esecutore multi-host XLA

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

Esempio di un singolo processo: prima della riproduzione del grafico di ottimizzazione

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

Esempio di un singolo processo: dopo la riproduzione del grafico di ottimizzazione

Per riprodurre un HLO ottimizzato, devi utilizzare --xla_disable_all_hlo_passes o --run_xla_backend_only. In caso contrario, XLA tenterà di ricompilare l'HLO, ma questa operazione non è supportata. Di conseguenza, si verificheranno molti errori strani.

Comando completo: 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

Più processi, un singolo nodo

Esegui il deployment del container

Installa anche alcune librerie mancanti. (Tieni presente che questa operazione può diventare obsoleta più rapidamente. Apporta le modifiche necessarie.)

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

Esegui il modello originale e scarica l'HLO

Per questo esempio, utilizzeremo un modello PAXML a 8 GPU da test-pax.sh. (Tieni presente che questo dump sarà lo stesso del caso di un singolo processo. Quindi, se lo hai già, puoi eseguire 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

Il dump HLO verrà salvato in /tmp/dump_multi_process/. Per PAX in particolare, il modulo principale avrà "pjit__wrapped_step_fn" nel nome. Per questo esempio, utilizzeremo /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Esegui su un singolo nodo utilizzando MPI

Crea uno script bash denominato 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

Ora puoi eseguirlo utilizzando mpirun:

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

Esegui su più nodi con SLURM

Quando esegui su più nodi utilizzando SLURM, puoi inoltrare le variabili di ambiente SLURM all'esecutore HLO nel job 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