این ابزار به شما امکان می‌دهد یک ماژول HLO را روی یک یا چند پردازنده گرافیکی اجرا کنید. همچنین امکان کامپایل کدی را فراهم می‌کند که چندین پردازنده گرافیکی را بدون اجرای آن هدف قرار می‌دهد.

اجرای HLO های چند پردازنده گرافیکی (sharded)

ما می‌توانیم این HLOها را با دیدن حاشیه‌نویسی‌های sharding= شناسایی کنیم. برای مثال sharding={devices=[1,1,2,1]0,1} به این معنی است که تانسور حاشیه‌نویسی شده باید در امتداد بعد سوم به ۲ پردازنده گرافیکی (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

در این حالت، یک پردازنده گرافیکی واحد لازم است، مگر اینکه از حافظه پنهان اتوتیونینگ استفاده شود.

عیب‌یابی

• خطاهایی مانند Check failed: result.replicas >= 1 (0 vs. 1) :
  • ما باید مطمئن شویم که به اندازه کافی پردازنده گرافیکی (GPU) داریم.
  • CUDA_VISIBLE_DEVICES باید به درستی تنظیم شود یا اصلاً تنظیم نشود.
• خرابی‌ها:
  • شاید بخواهیم از --dynamic_mode=off استفاده کنیم.
  • CUDA و cuDNN باید به درستی تنظیم شده باشند.

مثال‌ها

مثال یک فرآیند، چندین پردازنده گرافیکی

راه‌اندازی و دریافت 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 را تخلیه کنید

برای این مثال، ما از یک مدل PAXML با ۸ پردازنده گرافیکی از test-pax.sh استفاده خواهیم کرد. (توجه داشته باشید که این همان dump مربوط به حالت تک فرآیندی خواهد بود. بنابراین اگر از قبل آن را دارید، می‌توانید دستور زیر را اجرا کنید: 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 dump در مسیر /tmp/dump_multi_process/ ذخیره خواهد شد. به طور خاص برای PAX، ماژول اصلی دارای عبارت "pjit__wrapped_step_fn" در نام خود خواهد بود. برای این مثال ما از /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt استفاده خواهیم کرد.

اجرا روی یک گره واحد با استفاده از MPI

یک اسکریپت bash به نام 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

حالا می‌توانید آن را با استفاده از mpirun اجرا کنید:

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

اجرا روی چندین گره با SLURM

هنگام اجرا روی چندین گره با استفاده از SLURM، می‌توانید متغیرهای محیط SLURM را مانند job 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