Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

マルチホスト 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

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

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

その場合、自動チューニングキャッシュを使用しない限り、単一の GPU が必要になります。

トラブルシューティング

- 十分な 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