Exécuteur HLO multihôte

Cet outil vous permet d'exécuter un module HLO sur un ou plusieurs GPU. Il permet également de compiler du code ciblant plusieurs GPU sans l'exécuter.

Exécuter des HLO multi-GPU (fragmentés)

Nous pouvons identifier ces HLO grâce aux annotations sharding=. Par exemple, sharding={devices=[1,1,2,1]0,1} signifie que le Tensor annoté doit être fragmenté sur deux GPU (GPU0 et GPU1) le long de la troisième dimension.

Les instructions suivantes supposent que le répertoire de travail est le dépôt Git XLA et que ./configure.py a été exécuté.

Si nous disposons de suffisamment de GPU, nous pouvons relire ces HLO comme suit :

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

Il est également possible de compiler le même HLO sans l'exécuter en définissant --run=false :

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

Dans ce cas, un seul GPU est nécessaire, sauf si le cache d'optimisation automatique est utilisé.

Dépannage

Erreurs telles que Check failed: result.replicas >= 1 (0 vs. 1) :
- Nous devons nous assurer de disposer d'un nombre suffisant de GPU.
- CUDA_VISIBLE_DEVICES doit être défini correctement ou pas du tout.
Accidents :
- Nous pouvons utiliser --dynamic_mode=off.
- CUDA et cuDNN doivent être correctement configurés.

Exemples

Exemple de processus unique avec plusieurs GPU

Configurer et obtenir le 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.

Créer un exécuteur multihôte XLA

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

Exemple de processus unique : relecture du graphique avant optimisation

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

Exemple de processus unique : relecture du graphique après optimisation

Pour relire un HLO optimisé, vous devez utiliser --xla_disable_all_hlo_passes ou --run_xla_backend_only. Sinon, XLA tentera de recompiler le HLO, ce qui n'est pas pris en charge. Vous obtiendrez donc de nombreuses erreurs étranges.

Commande complète : 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

Multiprocessus, nœud unique

Lancer le conteneur

Installez également les bibliothèques manquantes. (Notez que ces informations peuvent devenir obsolètes plus rapidement. Ajustez-les si nécessaire.)

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

Exécuter le modèle d'origine et vider HLO

Pour cet exemple, nous utiliserons un modèle PAXML à huit GPU provenant de test-pax.sh. (Notez que ce vidage sera le même que dans le cas d'un processus unique. Vous pouvez donc faire cp -r /tmp/dump /tmp/dump_multi_process si vous l'avez déjà. 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

Le dump HLO sera enregistré dans /tmp/dump_multi_process/. Pour PAX en particulier, le module principal aura "pjit__wrapped_step_fn" dans son nom. Pour cet exemple, nous allons utiliser /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Exécuter sur un seul nœud à l'aide de MPI

Créez un script Bash appelé 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

Vous pouvez maintenant l'exécuter à l'aide de mpirun :

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

Exécuter sur plusieurs nœuds avec SLURM

Lorsque vous exécutez sur plusieurs nœuds à l'aide de SLURM, vous pouvez transférer les variables d'environnement SLURM vers le runner HLO comme suit dans votre 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