Trình chạy HLO nhiều máy chủ

Công cụ này cho phép bạn chạy một mô-đun HLO trên một hoặc nhiều GPU. Thao tác này cũng cho phép biên dịch mã nhắm đến nhiều GPU mà không cần chạy mã đó.

Chạy HLO đa GPU (phân mảnh)

Chúng ta có thể xác định những HLO này bằng cách xem chú thích sharding=. Ví dụ: sharding={devices=[1,1,2,1]0,1} có nghĩa là tenxơ được chú thích sẽ được phân mảnh thành 2 GPU (GPU0 và GPU1) dọc theo chiều thứ 3.

Hướng dẫn sau đây giả định rằng thư mục đang hoạt động là kho lưu trữ XLA Git và ./configure.py đã được chạy.

Nếu có đủ GPU, chúng ta có thể phát lại các HLO này như sau:

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

Bạn cũng có thể biên dịch cùng một HLO mà không cần chạy bằng cách thiết lập --run=false:

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

Trong trường hợp đó, bạn cần có một GPU, trừ phi bạn sử dụng bộ nhớ đệm tự động điều chỉnh.

Khắc phục sự cố

  • Các lỗi như Check failed: result.replicas >= 1 (0 vs. 1):
    • Chúng ta phải đảm bảo có đủ GPU.
    • Bạn phải đặt CUDA_VISIBLE_DEVICES một cách chính xác hoặc hoàn toàn không đặt.
  • Vấn đề trục trặc:
    • Chúng ta có thể muốn sử dụng --dynamic_mode=off.
    • Bạn phải thiết lập CUDA và cuDNN đúng cách.

Ví dụ

Ví dụ về một quy trình, nhiều GPU

Thiết lập và nhận 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.

Tạo trình chạy nhiều máy chủ XLA

cd /opt/xla/
./configure.py --backend CUDA --nccl
bazel build //xla/tools/multihost_hlo_runner:hlo_runner_main

Ví dụ về quy trình đơn: Trước khi phát lại biểu đồ tối ưu hoá

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \
  /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

Ví dụ về một quy trình: Sau khi phát lại biểu đồ tối ưu hoá

Để phát lại một HLO được tối ưu hoá, bạn phải sử dụng --xla_disable_all_hlo_passes hoặc --run_xla_backend_only. Nếu không, XLA sẽ cố gắng biên dịch lại HLO và điều này không được hỗ trợ. Vì vậy, nó sẽ gây ra nhiều lỗi lạ.

Câu lệnh đầy đủ: 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

Nhiều quy trình, một nút

Khởi chạy vùng chứa

Đồng thời cài đặt một số thư viện còn thiếu. (Xin lưu ý rằng thông tin này có thể nhanh chóng trở nên lỗi thời. Điều chỉnh nếu cần.)

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

Chạy mô hình gốc và kết xuất HLO

Trong ví dụ này, chúng ta sẽ sử dụng mô hình PAXML 8 GPU từ test-pax.sh. (Lưu ý: đây sẽ là cùng một tệp kết xuất như trường hợp một quy trình. Bạn có thể làm cp -r /tmp/dump /tmp/dump_multi_process nếu đã có. 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

Tệp kết xuất HLO sẽ được lưu vào /tmp/dump_multi_process/. Cụ thể đối với PAX, mô-đun chính sẽ có "pjit__wrapped_step_fn" trong tên. Trong ví dụ này, chúng ta sẽ sử dụng /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt.

Chạy trên một nút bằng MPI

Tạo một tập lệnh bash có tên là 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

Giờ đây, bạn có thể thực thi bằng mpirun:

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

Chạy trên nhiều nút bằng SLURM

Khi chạy trên nhiều nút bằng SLURM, bạn có thể chuyển tiếp các biến môi trường SLURM đến trình chạy HLO như sau trong công việc 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