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

اجرای HLO های چند GPU (تکه ای).

ما می توانیم این HLO ها را با دیدن sharding= annotations شناسایی کنیم. به عنوان مثال sharding={devices=[1,1,2,1]0,1} به این معنی است که تانسور حاشیه‌نویسی شده باید به ۲ GPU (GPU0 و GPU1) در امتداد بعد سوم تقسیم شود.

دستورالعمل های زیر فرض می کنند که دایرکتوری کاری مخزن xla Git است و ./configure.py اجرا شده است.

اگر پردازنده‌های گرافیکی کافی داشته باشیم، می‌توانیم این HLO‌ها را به صورت زیر پخش کنیم:

bazel run -c opt --config=cuda --dynamic_mode=off \
  //xla/tools/multihost_hlo_runner:hlo_runner_main -- my-hlo.txt

همچنین می توان همان HLO را بدون اجرای آن با تنظیم --run=false کامپایل کرد

bazel run -c opt --config=cuda --dynamic_mode=off \
  //xla/tools/multihost_hlo_runner:hlo_runner_main \
  -- --run=false my-hlo.txt

در این صورت، یک واحد پردازش گرافیکی ضروری است، مگر اینکه از کش خودکار استفاده شود.

عیب یابی

• خطاهایی مانند 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 --xla_dump_hlo_as_text"; 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 -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main

مثال تک فرآیند: قبل از پخش مجدد نمودار بهینه سازی

bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

مثال تک فرآیند: پس از بهینه سازی، پخش مجدد نمودار

برای پخش مجدد یک HLO بهینه شده، باید از آن دو پارامتر --run_xla_backend_only=true --xla_disable_all_hlo_passes=true استفاده کنید. در غیر این صورت، سعی می کند HLO را دوباره کامپایل کند و این پشتیبانی نمی شود. بنابراین به شما خطاهای عجیب و غریب زیادی می دهد.

دستور کامل: bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run_xla_backend_only=true --xla_disable_all_hlo_passes=true /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 8-GPU از test-pax.sh استفاده خواهیم کرد. (توجه داشته باشید که این همان dump مورد تک فرآیندی خواهد بود. بنابراین اگر قبلاً آن را دارید می‌توانید cp -r /tmp/dump /tmp/dump_multi_process انجام دهید. export XLA_FLAGS="--xla_dump_to=/tmp/dump_multi_process --xla_dump_hlo_as_text" 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 روی یک گره اجرا شود

یک اسکریپت bash به نام run.sh ایجاد کنید:

#!/bin/bash
export CUDA_VISIBLE_DEVICES=${OMPI_COMM_WORLD_LOCAL_RANK}
bazel run -c opt --config=cuda --dynamic_mode=off //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 env را به hlo runner ارسال کنید.

bazel run -c opt --config=cuda --dynamic_mode=off //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