Multi-Host HLO Runner

تتيح لك هذه الأداة تشغيل وحدة HLO على وحدة معالجة رسومات واحدة أو أكثر. ويتيح أيضًا تجميع الرمز البرمجي الذي يستهدف وحدات معالجة رسومات متعددة بدون تشغيله.

تشغيل HLO متعددة وحدات معالجة الرسومات (مجزّأة)

يمكننا تحديد هذه الكائنات العالية المستوى من خلال رؤية التعليقات التوضيحية sharding=. على سبيل المثال، يعني sharding={devices=[1,1,2,1]0,1} أنّه يجب تقسيم الموتر المشروح إلى وحدتَي معالجة رسومات (GPU0 وGPU1) على طول البُعد الثالث.

تفترض التعليمات التالية أنّ دليل العمل هو مستودع XLA Git وأنّه تم تنفيذ ./configure.py.

إذا توفّرت لدينا وحدات معالجة الرسومات الكافية، يمكننا إعادة تشغيل هذه الألعاب على النحو التالي:

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

في هذه الحالة، تكون وحدة معالجة رسومات واحدة ضرورية، ما لم يتم استخدام ذاكرة التخزين المؤقت للضبط التلقائي.

تحديد المشاكل وحلّها

  • أخطاء مثل 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 مزوّدًا بـ 8 وحدات معالجة الرسومات من test-pax.sh. (ملاحظة سيكون هذا التفريغ هو نفسه التفريغ في حالة العملية الفردية. لذا يمكنك تنفيذ الخطوة 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