এই টুলটি আপনাকে এক বা একাধিক জিপিইউ-তে একটি এইচএলও মডিউল চালাতে দেয়। এটি কোড না চালিয়েই একাধিক জিপিইউ-কে লক্ষ্য করে কম্পাইল করার সুযোগও দেয়।

একাধিক জিপিইউ (শার্ডেড) এইচএলও চালানো হচ্ছে

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

নিম্নলিখিত নির্দেশাবলীতে ধরে নেওয়া হয়েছে যে ওয়ার্কিং ডিরেক্টরি হলো XLA গিট রিপোজিটরি এবং ./configure.py চালানো হয়েছে।

আমাদের কাছে পর্যাপ্ত GPU থাকলে, আমরা এই HLO-গুলো এইভাবে পুনরায় খেলতে পারি:

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

--compile_only=true সেট করে একই HLO রান না করেও কম্পাইল করা সম্ভব:

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

সেক্ষেত্রে একটিমাত্র জিপিইউ প্রয়োজন, যদি না অটোটিউনিং ক্যাশে ব্যবহার করা হয়।

সমস্যা সমাধান

• এরর যেমন 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

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

একটি অপ্টিমাইজ করা 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 নামে একটি ব্যাশ স্ক্রিপ্ট তৈরি করুন:

#!/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