Esta herramienta te permite ejecutar un módulo HLO en una o más GPUs. También permite compilar código destinado a varias GPUs sin ejecutarlo.

Ejecuta HLOs de varias GPUs (fragmentados)

Podemos identificar estos HLOs con las anotaciones sharding=. Por ejemplo, sharding={devices=[1,1,2,1]0,1} significa que el tensor anotado se debe fragmentar en 2 GPUs (GPU0 y GPU1) a lo largo de la tercera dimensión.

En las siguientes instrucciones, se da por hecho que el directorio de trabajo es el repositorio de Git de XLA y que se ejecutó ./configure.py.

Si tenemos suficientes GPUs, podemos reproducir estos HLOs de la siguiente manera:

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

También es posible compilar el mismo HLO sin ejecutarlo si se configura --compile_only=true:

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

En ese caso, se necesita una sola GPU, a menos que se use la caché de ajuste automático.

Solución de problemas

• Errores como Check failed: result.replicas >= 1 (0 vs. 1):
  • Debemos asegurarnos de tener suficientes GPUs.
  • CUDA_VISIBLE_DEVICES se debe configurar correctamente o no configurarse en absoluto.
• Bloqueos:
  • Es posible que queramos usar --dynamic_mode=off.
  • CUDA y cuDNN deben configurarse correctamente.

Ejemplos

Ejemplo de un solo proceso y varias GPUs

Configura y obtén el 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.

Compila el ejecutor de varios hosts de XLA

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

Ejemplo de un solo proceso: Antes de la reproducción del gráfico de optimización

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

Ejemplo de un solo proceso: Después de la reproducción del gráfico de optimización

Para reproducir un HLO optimizado, debes usar --xla_disable_all_hlo_passes o --run_xla_backend_only. De lo contrario, XLA intentará volver a compilar el HLO, lo que no es compatible. Por lo tanto, te dará muchos errores extraños.

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

Varios procesos, un solo nodo

Inicia el contenedor

También instala algunas bibliotecas faltantes. (Ten en cuenta que se puede desactualizar más rápido. Realiza ajustes según sea necesario).

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

Ejecuta el modelo original y vuelca el HLO

Para este ejemplo, usaremos un modelo PAXML de 8 GPUs de test-pax.sh. (Ten en cuenta que este será el mismo volcado que el caso de un solo proceso. Por lo tanto, puedes hacer cp -r /tmp/dump /tmp/dump_multi_process si ya lo tienes. 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

El volcado de HLO se guardará en /tmp/dump_multi_process/. En el caso de PAX específicamente, el módulo principal tendrá "pjit__wrapped_step_fn" en el nombre. Para este ejemplo, usaremos /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Ejecuta en un solo nodo con MPI

Crea una secuencia de comandos de Bash llamada 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

Ahora, puedes ejecutarla con mpirun:

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

Ejecuta en varios nodos con SLURM

Cuando se ejecuta en varios nodos con SLURM, puedes reenviar las variables de entorno de SLURM al ejecutor de HLO de la siguiente manera en tu trabajo de 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