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

Cómo ejecutar HLO con varias GPU (fragmentados)

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

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

Si tenemos suficientes GPUs, podemos volver a ejecutar estos HLO 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 configurando --run=false:

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run=false 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):
  • Tenemos que asegurarnos de tener suficientes GPUs.
  • CUDA_VISIBLE_DEVICES debe configurarse correctamente o no establecerse en absoluto.
• Accidentes automovilísticos:
  • Es posible que queramos usar --dynamic_mode=off.
  • CUDA y cuDNN deben estar configurados correctamente.

Ejemplos

Ejemplo de varios procesos 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 se admite. Por lo tanto, te mostrará 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

Multiprocesos, un solo nodo

Inicia el contenedor

También instala algunas bibliotecas faltantes. (Ten en cuenta que puede quedar desactualizada más rápido. Realiza los ajustes necesarios.

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

En este ejemplo, usaremos un modelo PAXML de 8 GPUs de test-pax.sh. (Ten en cuenta que este será el mismo volcado que en 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 ejecutarlo 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