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

Esecuzione di HLO multi-GPU (frammentati)

Possiamo identificare questi HLO grazie alle annotazioni sharding=. Ad esempio, sharding={devices=[1,1,2,1]0,1} significa che il tensore annotato deve essere suddiviso in 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 abbastanza GPU, possiamo riprodurre questi HLO in questo modo:

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

È anche possibile compilare lo stesso HLO senza eseguirlo impostando --run=false:

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

In questo caso, è necessaria una sola 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 GPU sufficienti.
  • CUDA_VISIBLE_DEVICES deve essere impostato correttamente o non impostato affatto.
• Incidenti:
  • Potremmo voler utilizzare --dynamic_mode=off.
  • CUDA e cuDNN devono essere configurati correttamente.

Esempi

Esempio di più GPU in un singolo processo

Configurare e ottenere l'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.

Crea il runner multihost XLA

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

Esempio di singolo processo: riproduzione del grafico prima dell'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: riproduzione del grafico dopo l'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 e questa operazione non è supportata. Pertanto, riceverai 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, singolo nodo

Avvia container

Installa anche alcune librerie mancanti. Tieni presente che queste informazioni potrebbero diventare obsolete più rapidamente. Regola in base alle esigenze.)

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 HLO

Per questo esempio, utilizzeremo un modello PAXML a 8 GPU di test-pax.sh. (Nota: questo sarà lo stesso dump del caso di un singolo processo. Quindi puoi fare cp -r /tmp/dump /tmp/dump_multi_process se ce l'hai già. 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 viene eseguito su più nodi utilizzando SLURM, puoi inoltrare le variabili di ambiente SLURM al runner HLO nel seguente modo nel tuo 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