多主機 HLO Runner

這個工具可讓您在一或多個 GPU 上執行 HLO 模組。此外，您也可以編譯以多個 GPU 為目標的程式碼，而不需執行。

執行多 GPU (分片) HLO

我們可透過 sharding= 註解識別這些 HLO。舉例來說，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 環境變數轉送至 HLO 執行器，方法是在 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