Multi-Host HLO Runner

Alat ini memungkinkan Anda menjalankan modul HLO di satu atau beberapa GPU. Hal ini juga memungkinkan mengompilasi kode yang menargetkan beberapa GPU tanpa menjalankannya.

Menjalankan HLO multi-GPU (sharded)

Kita dapat mengidentifikasi HLO ini dengan melihat anotasi sharding=. Misalnya, sharding={devices=[1,1,2,1]0,1} berarti tensor yang dianotasi harus di-shard ke 2 GPU (GPU0 dan GPU1) di sepanjang dimensi ke-3.

Petunjuk berikut mengasumsikan direktori kerja adalah repositori XLA Git dan ./configure.py telah dijalankan.

Jika memiliki GPU yang memadai, kita dapat memutar ulang HLO ini seperti berikut:

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

Anda juga dapat mengompilasi HLO yang sama tanpa menjalankannya dengan menyetel --run=false:

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

Dalam hal ini, satu GPU diperlukan, kecuali jika cache penyetelan otomatis digunakan.

Pemecahan masalah

  • Error seperti Check failed: result.replicas >= 1 (0 vs. 1):
    • Kita harus memastikan bahwa kita memiliki GPU yang cukup.
    • CUDA_VISIBLE_DEVICES harus ditetapkan dengan benar atau tidak ditetapkan sama sekali.
  • Error:
    • Kita mungkin ingin menggunakan --dynamic_mode=off.
    • CUDA dan cuDNN harus disiapkan dengan benar.

Contoh

Contoh satu proses, beberapa GPU

Siapkan dan dapatkan 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.

Membangun runner multihost XLA

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

Contoh proses tunggal: Sebelum pemutaran ulang grafik pengoptimalan

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

Contoh proses tunggal: Setelah pemutaran ulang grafik pengoptimalan

Untuk memutar ulang HLO yang dioptimalkan, Anda harus menggunakan --xla_disable_all_hlo_passes atau --run_xla_backend_only. Jika tidak, XLA akan mencoba mengompilasi ulang HLO dan hal ini tidak didukung. Jadi, akan muncul banyak error aneh.

Perintah lengkap: 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

Multi-proses, satu node

Luncurkan penampung

Instal juga beberapa library yang belum ada. (Perhatikan, informasi tersebut dapat menjadi tidak berlaku lebih cepat. Sesuaikan seperlunya.)

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

Jalankan model asli dan dump HLO

Untuk contoh ini, kita akan menggunakan model PAXML 8 GPU dari test-pax.sh. (Catatan: dump ini akan sama dengan dump kasus proses tunggal. Jadi, Anda dapat melakukan cp -r /tmp/dump /tmp/dump_multi_process jika sudah memilikinya. 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

Dump HLO akan disimpan ke /tmp/dump_multi_process/. Khusus untuk PAX, modul utama akan memiliki "pjit__wrapped_step_fn" dalam namanya. Untuk contoh ini, kita akan menggunakan /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Menjalankan di satu node menggunakan MPI

Buat skrip bash bernama 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

Sekarang, Anda dapat menjalankannya menggunakan mpirun:

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

Menjalankan di beberapa node dengan SLURM

Saat berjalan di beberapa node menggunakan SLURM, Anda dapat meneruskan variabel lingkungan SLURM ke runner HLO seperti ini di tugas SLURM Anda:

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