Bu araç, bir HLO modülünü bir veya daha fazla GPU'da çalıştırmanıza olanak tanır. Ayrıca, kodu çalıştırmadan birden fazla GPU'yu hedefleyen kod derlemeye de olanak tanır.

Çoklu GPU (parçalanmış) HLO'ları çalıştırma

Bu HLO'ları sharding= ek açıklamalarından tanıyabiliriz. Örneğin, sharding={devices=[1,1,2,1]0,1}, açıklama eklenmiş tensörün 3. boyut boyunca 2 GPU'ya (GPU0 ve GPU1) parçalanması gerektiğini gösterir.

Aşağıdaki talimatlarda, çalışma dizininin XLA Git deposu olduğu ve ./configure.py komutunun çalıştırıldığı varsayılmaktadır.

Yeterli sayıda GPU'muz varsa bu HLO'ları şu şekilde yeniden oynatabiliriz:

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

Aynı HLO'yu çalıştırmadan derlemek de mümkündür. Bunun için --run=false ayarlanmalıdır:

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

Bu durumda, otomatik ayarlama önbelleği kullanılmıyorsa tek bir GPU gerekir.

Sorun giderme

• Check failed: result.replicas >= 1 (0 vs. 1) gibi hatalar:
  • Yeterli sayıda GPU'ya sahip olduğumuzdan emin olmamız gerekir.
  • CUDA_VISIBLE_DEVICES doğru şekilde ayarlanmalı veya hiç ayarlanmamalıdır.
• Kazalar:
  • --dynamic_mode=off kullanmak isteyebiliriz.
  • CUDA ve cuDNN doğru şekilde ayarlanmalıdır.

Örnekler

Tek işlem, çoklu GPU örneği

Kurulum ve HLO'yu alma

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 çoklu ana makine çalıştırıcısı oluşturma

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

Tek süreç örneği: Optimizasyon grafiği yeniden oynatılmadan önce

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

Tek süreç örneği: Optimizasyon grafiği tekrarından sonra

Optimize edilmiş bir HLO'yu yeniden oynatmak için --xla_disable_all_hlo_passes veya --run_xla_backend_only kullanmanız gerekir. Aksi takdirde, XLA, HLO'yu yeniden derlemeye çalışır ve bu desteklenmez. Bu nedenle, birçok tuhaf hata alırsınız.

Tam komut: 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

Çoklu işlemler, tek düğüm

Kapsayıcıyı başlatma

Ayrıca eksik kitaplıkları da yükleyin. (Bu bilgilerin daha hızlı eski hale gelebileceğini unutmayın. Gerektiği şekilde ayarlayın.)

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

Orijinal modeli çalıştırma ve HLO'yu dökme

Bu örnekte, test-pax.sh adresinden 8 GPU'lu bir PAXML modeli kullanacağız. (Not: Bu, tek süreçli durumla aynı döküm olacaktır. Bu nedenle, cp -r /tmp/dump /tmp/dump_multi_process varsa 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 yapabilirsiniz.

HLO dökümü /tmp/dump_multi_process/ konumuna kaydedilir. PAX için özellikle, ana modülün adında "pjit__wrapped_step_fn" bulunur. Bu örnekte /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt kullanılacaktır.

MPI kullanarak tek bir düğümde çalıştırma

run.sh adlı bir bash komut dosyası oluşturun:

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

Artık mpirun kullanarak yürütebilirsiniz:

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

SLURM ile birden fazla düğümde çalıştırma

SLURM kullanarak birden fazla düğümde çalıştırırken SLURM ortam değişkenlerini SLURM işinizde aşağıdaki gibi HLO çalıştırıcısına iletebilirsiniz:

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