マルチホスト HLO ランナー

このツールを使用すると、1 つ以上の GPU で HLO モジュールを実行できます。また、コードを実行せずに複数の GPU を対象とするコードをコンパイルすることもできます。

マルチ GPU（シャーディング）HLO の実行

これらの HLO は sharding= アノテーションで識別できます。たとえば、sharding={devices=[1,1,2,1]0,1} は、アノテーション付きテンソルを 3 番目のディメンションに沿って 2 つの GPU（GPU0 と GPU1）にシャーディングする必要があることを意味します。

次の手順では、作業ディレクトリが XLA Git リポジトリであり、./configure.py が実行されていることを前提としています。

十分な GPU がある場合は、次のように HLO を再生できます。

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

--run=false を設定することで、実行せずに同じ HLO をコンパイルすることもできます。

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

この場合、自動チューニングキャッシュを使用しない限り、単一の GPU が必要です。

トラブルシューティング

Check failed: result.replicas >= 1 (0 vs. 1) などのエラー:
- 十分な GPU があることを確認する必要があります。
- CUDA_VISIBLE_DEVICES は正しく設定するか、まったく設定しないかのいずれかにする必要があります。
クラッシュ:
- --dynamic_mode=off を使用することもできます。
- CUDA と cuDNN が正しく設定されている必要があります。

例

単一プロセス、複数の GPU の例

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.

XLA マルチホストランナーをビルドする

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

単一プロセスの例: 最適化前のグラフの再生

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

単一プロセスでの例: 最適化グラフの再生後

最適化された HLO を再生するには、--xla_disable_all_hlo_passes または --run_xla_backend_only を使用する必要があります。それ以外の場合、XLA は HLO の再コンパイルを試みますが、これはサポートされていません。そのため、多くの奇妙なエラーが発生します。

完全なコマンド: 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

マルチプロセス、シングルノード

コンテナを起動する

また、不足しているライブラリもインストールします。（ただし、この情報はすぐに古くなる可能性があります。必要に応じて調整します）。

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

元のモデルを実行して HLO をダンプする

この例では、test-pax.sh の 8 GPU PAXML モデルを使用します。（これは、単一プロセスの場合と同じダンプになります。cp -r /tmp/dump /tmp/dump_multi_process がすでにある場合は、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 を実行できます。

HLO ダンプは /tmp/dump_multi_process/ に保存されます。特に PAX の場合、メインモジュールの名前には「pjit__wrapped_step_fn」が含まれます。この例では、/tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt を使用します。

MPI を使用して単一ノードで実行する

run.sh という bash スクリプトを作成します。

#!/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

これで、mpirun を使用して実行できます。

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

SLURM を使用して複数のノードで実行する

SLURM を使用して複数のノードで実行する場合は、SLURM ジョブで次のように SLURM 環境変数を HLO ランナーに転送できます。

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