এই টুলটি আপনাকে এক বা একাধিক GPU-তে HLO মডিউল চালাতে দেয়। এটি একাধিক GPU গুলিকে না চালিয়ে কোড কম্পাইল করার অনুমতি দেয়৷

মাল্টি-জিপিইউ (শার্ডেড) এইচএলও চলছে

আমরা sharding= টীকা দেখে এই HLOগুলি সনাক্ত করতে পারি। উদাহরণস্বরূপ sharding={devices=[1,1,2,1]0,1} এর অর্থ হল টীকাযুক্ত টেনসরটি 3য় মাত্রা বরাবর 2টি জিপিইউ (GPU0 এবং GPU1) এ শার্ড করা উচিত।

নিম্নলিখিত নির্দেশাবলী অনুমান করে যে কার্যকারী ডিরেক্টরি হল xla Git সংগ্রহস্থল এবং ./configure.py চালানো হয়েছে।

যদি আমাদের পর্যাপ্ত জিপিইউ থাকে, তাহলে আমরা এই HLOগুলিকে এভাবে রিপ্লে করতে পারি:

bazel run -c opt --config=cuda --dynamic_mode=off \
  //xla/tools/multihost_hlo_runner:hlo_runner_main -- my-hlo.txt

--run=false সেট করে এটি না চালিয়ে একই HLO কম্পাইল করাও সম্ভব

bazel run -c opt --config=cuda --dynamic_mode=off \
  //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 --xla_dump_hlo_as_text"; 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 -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main

একক প্রক্রিয়া উদাহরণ: অপ্টিমাইজেশান গ্রাফ রিপ্লে করার আগে

bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt

একক প্রক্রিয়া উদাহরণ: অপ্টিমাইজেশন গ্রাফ রিপ্লে পরে

একটি অপ্টিমাইজ করা HLO রিপ্লে করতে, আপনাকে অবশ্যই এই দুটি প্যারামিটার ব্যবহার করতে হবে --run_xla_backend_only=true --xla_disable_all_hlo_passes=true । অন্যথায়, এটি HLO পুনরায় কম্পাইল করার চেষ্টা করবে এবং এটি সমর্থিত নয়। তাই এটি আপনাকে অনেক অদ্ভুত ত্রুটি দেবে।

সম্পূর্ণ কমান্ড: bazel run -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run_xla_backend_only=true --xla_disable_all_hlo_passes=true /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 --xla_dump_hlo_as_text" 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 নামে একটি ব্যাশ স্ক্রিপ্ট তৈরি করুন:

#!/bin/bash
export CUDA_VISIBLE_DEVICES=${OMPI_COMM_WORLD_LOCAL_RANK}
bazel run -c opt --config=cuda --dynamic_mode=off //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 env ভেরিয়েবলগুলি hlo রানারে ফরোয়ার্ড করতে পারেন আপনার স্লার্ম কাজের মতো:

bazel run -c opt --config=cuda --dynamic_mode=off //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