멀티 호스트 HLO 러너

이 도구를 사용하면 하나 이상의 GPU에서 HLO 모듈을 실행할 수 있습니다. 또한 실행하지 않고 여러 GPU를 타겟팅하는 코드를 컴파일할 수 있습니다.

멀티 GPU (샤딩된) HLO 실행

sharding= 주석을 확인하여 이러한 HLO를 식별할 수 있습니다. 예를 들어 sharding={devices=[1,1,2,1]0,1}는 주석이 달린 텐서가 세 번째 차원을 따라 GPU 2개 (GPU0 및 GPU1)로 샤딩되어야 함을 의미합니다.

다음 안내에서는 작업 디렉터리가 XLA Git 저장소이고 ./configure.py이 실행되었다고 가정합니다.

GPU가 충분한 경우 다음과 같이 이러한 HLO를 재생할 수 있습니다.

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

--run=false를 설정하여 실행하지 않고 동일한 HLO를 컴파일할 수도 있습니다.

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를 덤프합니다.

이 예에서는 test-pax.sh의 8-GPU PAXML 모델을 사용합니다. (단일 프로세스 사례와 동일한 덤프입니다. 따라서 이미 있는 경우 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를 사용하여 단일 노드에서 실행

run.sh라는 bash 스크립트를 만듭니다.

#!/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 작업에서 다음과 같이 SLURM 환경 변수를 HLO 러너로 전달할 수 있습니다.

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