Multi-Host HLO Runner

הכלי הזה מאפשר להריץ מודול HLO במעבד GPU אחד או יותר. בנוסף, אפשר להדר קוד שמיועד למספר מעבדי GPU בלי להריץ אותו.

הרצת HLOs מרובים ב-GPU (חלוקה לרסיסים)

אנחנו יכולים לזהות את ההמלצות האלה לרכיבים מרכזיים באמצעות sharding=הערות. לדוגמה, sharding={devices=[1,1,2,1]0,1} אומר שהטנזור עם ההערות צריך להיות מחולק ל-2 יחידות GPU‏ (GPU0 ו-GPU1) לאורך המימד השלישי.

במסגרת ההוראות הבאות, אנחנו מניחים שספריית העבודה היא מאגר ה-Git של XLA ושפקודת ./configure.py הופעלה.

אם יש לנו מספיק יחידות GPU, אנחנו יכולים להפעיל מחדש את ה-HLOs האלה כך:

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

אפשר גם לקמפל את אותו HLO בלי להריץ אותו, על ידי הגדרת --run=false:

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

במקרה כזה, נדרש GPU יחיד, אלא אם משתמשים בזיכרון מטמון לכוונון אוטומטי.

פתרון בעיות

  • שגיאות כמו Check failed: result.replicas >= 1 (0 vs. 1):
    • אנחנו צריכים לוודא שיש לנו מספיק מעבדים גרפיים.
    • צריך להגדיר את 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

בדוגמה הזו נשתמש במודל 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" 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 //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 למפעיל HLO באופן הבא במשימת SLURM:

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