多主机 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，可以执行 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