इस टूल की मदद से, एक या उससे ज़्यादा जीपीयू पर एचएलओ मॉड्यूल चलाया जा सकता है. इससे कोड को चलाए बिना, कई जीपीयू को टारगेट करने वाले कोड को कंपाइल भी किया जा सकता है.

मल्टी-जीपीयू (शार्ड किए गए) 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

--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) जैसी गड़बड़ियां:
  • हमें यह पक्का करना होगा कि हमारे पास ज़रूरत के मुताबिक़ जीपीयू हों.
  • 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

सिंगल प्रोसेस का उदाहरण: ऑप्टिमाइज़ेशन ग्राफ़ को फिर से चलाने के बाद

ऑप्टिमाइज़ किए गए एचएलओ को फिर से चलाने के लिए, आपको --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

ओरिजनल मॉडल चलाएं और एचएलओ डंप करें

इस उदाहरण के लिए, हम test-pax.sh से 8-जीपीयू वाला 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

एचएलओ डंप को /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 नाम की एक बैश स्क्रिप्ट बनाएं:

#!/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 रनर पर इस तरह से फ़ॉरवर्ड किया जा सकता है:

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