Essa ferramenta permite executar um módulo HLO em uma ou mais GPUs. Também é possível compilar código para várias GPUs sem executá-lo.

Execução de HLOS com várias GPUs (fragmentados)

Podemos identificar esses HLOs com as anotações sharding=. Por exemplo, sharding={devices=[1,1,2,1]0,1} significa que o tensor anotado precisa ser fragmentado em duas GPUs (GPU0 e GPU1) ao longo da terceira dimensão.

As instruções a seguir pressupõem que o diretório de trabalho seja o repositório Git do XLA e que ./configure.py tenha sido executado.

Se tivermos GPUs suficientes, poderemos reproduzir esses HLOs assim:

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

Também é possível compilar o mesmo HLO sem executá-lo definindo --run=false:

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

Nesse caso, uma única GPU é necessária, a menos que o cache de ajuste automático seja usado.

Solução de problemas

• Erros como Check failed: result.replicas >= 1 (0 vs. 1):
  • Precisamos garantir que temos GPUs suficientes.
  • O campo CUDA_VISIBLE_DEVICES precisa ser definido corretamente ou não ser definido.
• Acidentes:
  • Talvez seja necessário usar --dynamic_mode=off.
  • O CUDA e o cuDNN precisam estar configurados corretamente.

Exemplos

Exemplo de processo único com várias GPUs

Configurar e receber o 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.

Criar um executor multihost do XLA

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

Exemplo de processo único: antes da repetição do gráfico de otimização

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

Exemplo de processo único: após a repetição do gráfico de otimização

Para reproduzir um HLO otimizado, use --xla_disable_all_hlo_passes ou --run_xla_backend_only. Caso contrário, o XLA vai tentar recompilar o HLO, o que não é compatível. Isso vai gerar muitos erros estranhos.

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

Vários processos, um único nó

Iniciar contêiner

Instale também algumas bibliotecas ausentes. Essa informação pode ficar desatualizada rapidamente. Ajuste conforme necessário.

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

Executar o modelo original e despejar o HLO

Neste exemplo, vamos usar um modelo PAXML de 8 GPUs do test-pax.sh. Observação: esse será o mesmo despejo do caso de processo único. Então, você pode fazer cp -r /tmp/dump /tmp/dump_multi_process se já tiver. 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

O despejo de HLO será salvo em /tmp/dump_multi_process/. Para o PAX especificamente, o módulo principal terá "pjit__wrapped_step_fn" no nome. Neste exemplo, vamos usar /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Executar em um único nó usando MPI

Crie um script bash chamado 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

Agora, execute usando mpirun:

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

Executar em vários nós com o SLURM

Ao executar em vários nós usando o SLURM, você pode encaminhar as variáveis de ambiente do SLURM para o executor de HLO da seguinte maneira no seu job do 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