Multi-Host HLO Runner

เครื่องมือนี้ช่วยให้คุณเรียกใช้โมดูล HLO บน GPU อย่างน้อย 1 ตัวได้ นอกจากนี้ยังช่วยให้ คอมไพล์โค้ดที่กำหนดเป้าหมายไปยัง GPU หลายตัวได้โดยไม่ต้องเรียกใช้

การเรียกใช้ HLO แบบหลาย GPU (แบบ Shard)

เราสามารถระบุ HLO เหล่านี้ได้โดยดูsharding=คำอธิบายประกอบ เช่น sharding={devices=[1,1,2,1]0,1} หมายความว่าควรแบ่งพาร์ติชันเทนเซอร์ที่ใส่คำอธิบายประกอบเป็น 2 GPU (GPU0 และ GPU1) ตามมิติที่ 3

วิธีการต่อไปนี้จะถือว่าไดเรกทอรีการทำงานคือที่เก็บ Git ของ XLA และได้เรียกใช้ ./configure.py แล้ว

หากมี GPU เพียงพอ เราจะเล่น HLO เหล่านี้ซ้ำได้ดังนี้

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):
    • เราต้องตรวจสอบว่ามี 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"; 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 (หมายเหตุ การทิ้งข้อมูลนี้จะเหมือนกับกรณีที่มีกระบวนการเดียว คุณจึงทำ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