ডাম্প এইচএলও কম্পিউটেশন

একটি HLO ডাম্প হলো গণনার বিভিন্ন পর্যায়ে HLO মডিউলগুলোর একটি লিখিত উপস্থাপনা। এটি ডিবাগিংয়ের জন্য উপযোগী, এবং বাগ রিপোর্টে প্রায়শই এটি অন্তর্ভুক্ত করার প্রয়োজন হয়। এটি সাধারণত একটি পাঠযোগ্য টেক্সট ফাইল , যেখানে HLO নির্দেশাবলী এবং তাদের বৈশিষ্ট্যগুলো তালিকাভুক্ত থাকে। কখনও কখনও, HLO মডিউলগুলো নিম্নোক্তভাবে ডাম্প করা হয়:

• HloProto: প্রোটোকল বাফার ফাইল, যা একটি আরও সুসংগঠিত ও যন্ত্রপাঠযোগ্য ফরম্যাট।
• HloSnapshot : HLO মডিউল এবং এর ইনপুটসমূহ। HLO রিপ্লে করার জন্য, কখনও কখনও র‍্যান্ডম ডেটার পরিবর্তে কোনো নির্দিষ্ট গণনায় দেওয়া প্রকৃত ইনপুটগুলোর প্রয়োজন হয়।

আপনি ডাম্প নির্দিষ্ট করতে ও পেতে XLA ফ্ল্যাগ ব্যবহার করতে পারেন। বেশিরভাগ ক্ষেত্রে, আপনি এটি একটি এনভায়রনমেন্ট ভেরিয়েবল দিয়ে সেট করতে পারেন। JAX এছাড়াও HLO ডাম্প প্রিন্ট করার জন্য একটি প্রোগ্রাম্যাটিক উপায় প্রদান করে।

স্থানীয় নির্বাহ

এনভায়রনমেন্ট ভেরিয়েবল ব্যবহার করা

ডাম্প পেতে আপনি প্রয়োজনীয় ফ্ল্যাগগুলো দিয়ে XLA_FLAGS এনভায়রনমেন্ট ভেরিয়েবল সেট করতে পারেন। এটি JAX, TensorFlow, এবং PyTorch/XLA-এর জন্য কাজ করে।

HLO মডিউল এবং অন্যান্য ডিবাগিং তথ্য একটি নির্দিষ্ট ডিরেক্টরিতে ডাম্প করতে, আপনার প্রোগ্রামটি --xla_dump_to ফ্ল্যাগ সহ চালান:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

উদাহরণস্বরূপ, আপনি পাথ হিসেবে /tmp অথবা /tmp/xladump ব্যবহার করতে পারেন।

ডিফল্টরূপে, এটি অপ্টিমাইজেশন পাইপলাইনের একেবারে শুরুতে এবং শেষে HLO মডিউলগুলোকে টেক্সট হিসেবে ডাম্প করে।

আপনি ফরম্যাটটি স্পষ্টভাবেও নির্দিষ্ট করে দিতে পারেন:

1. টেক্সট ডাম্প

XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"

1. এইচএলও প্রোটোস

XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"

1. এইচএলও স্ন্যাপশট

XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"

1. গ্রাফভিজ সার্ভার দিয়ে গ্রাফ রেন্ডার (শুধুমাত্র ছোট গ্রাফের জন্য ভালোভাবে কাজ করে)

XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"

1. গ্রাফ HTML ফাইলে রেন্ডার করুন (শুধুমাত্র ছোট গ্রাফের জন্য ভালোভাবে কাজ করে)

XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

বৃহত্তর গ্রাফের ক্ষেত্রে, আপনি গ্রাফের অংশবিশেষ দেখার জন্য interactive_graphviz ব্যবহার করতে পারেন।

আপনি --xla_syntax_sugar_async_ops ফ্ল্যাগটিকে true সেট করে HLO ডাম্পে অপ নেম হিসেবে সিনট্যাকটিক সুগার র‍্যাপার ব্যবহার করার জন্যও সেট করতে পারেন। এটি ডাম্পের আকার প্রায় ২০% কমাতে পারে। ডিফল্টরূপে, এই ফ্ল্যাগটি false সেট করা থাকে এবং ডাম্পে আসল অপ নেম ব্যবহার করা হয়।

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_syntax_sugar_async_ops=true"

নির্দিষ্ট মধ্যবর্তী পাসগুলি ফেলে দিন

সাধারণ প্রি-অপ্টিমাইজড / ফাইনাল-অপ্টিমাইজড HLO-গুলো ছাড়াও, আপনি একটি নির্দিষ্ট কম্পাইলার পাসের পরে HLO-গুলোর অবস্থা ডাম্প করতে পারেন।

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

যেসব পাসের নাম রেগুলার এক্সপ্রেশন (regex)-এর সাথে মেলে, সেগুলোর জন্য HLO মডিউলগুলো ডাম্প করা হবে। উদাহরণস্বরূপ, আপনি SPMD পার্টিশনিং সম্পর্কিত পাসগুলো থেকে প্রাপ্ত HLO-গুলো নিম্নোক্তভাবে পর্যবেক্ষণ করতে পারেন:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

প্রতিটি XLA পাসের পরে ফলাফল ডাম্প করতে (এর ফলে অনেকগুলো ফাইল তৈরি হবে), আপনি সেট করতে পারেন:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX-নির্দিষ্ট বিকল্পগুলি

JAX-এ প্রোগ্রাম্যাটিকভাবে

ফ্ল্যাগ বা এনভায়রনমেন্ট ভেরিয়েবল পাস করার পরিবর্তে, আপনি JAX-এর lower এবং compile API ব্যবহার করে প্রোগ্রাম্যাটিকভাবে HLO ডাম্প করতে পারেন।

স্থানীয়ভাবে অপ্টিমাইজ না করা মূল নিম্ন HLO-টি নিম্নলিখিত উপায়ে আনুন:

jax.jit(f).lower(*args).as_text('hlo')

HLO কম্পাইলেশন পাসের সময় ফাইলে ডাম্প করার জন্য, নির্দিষ্ট করুন:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

ডাম্প জ্যাক্সপিআরএস

jaxpr হলো প্রোগ্রাম ট্রেসের জন্য JAX-এর অন্তর্বর্তীকালীন উপস্থাপনা। এটি ডাম্প করতে, এনভায়রনমেন্ট ভেরিয়েবলগুলো সেট করুন:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

স্টেজড-আউট করা গণনা এক্সপোর্ট এবং সিরিয়ালাইজ করা: ডিবাগিং বিষয়ে JAX ডকুমেন্টেশনে আরও জানুন।

গুগল কোলাব

পরিবেশগত পরিবর্তনশীল

আপনার নোটবুকের প্রথম এক্সিকিউটেড সেলে (কারণ এনভায়রনমেন্ট ভেরিয়েবল এবং কমান্ড-লাইন ফ্ল্যাগ সাধারণত একবারই প্রসেস করা হয়, যেমন, মডিউল-ইম্পোর্ট বা XLA ব্যাকএন্ড ইনিশিয়ালাইজেশনের সময়), os.environ ব্যবহার করে উপরে বর্ণিত XLA_FLAGS টি যোগ করুন, উদাহরণস্বরূপ:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

এটি গণনার ফলাফল DIRECTORY_PATH এ, যেমন /tmp , জমা করবে। Colab-এ, এই ডিরেক্টরিটি দেখতে ও অ্যাক্সেস করতে বাম সাইডবারের "Files" ব্রাউজারে যান।

আপনি লোকাল এক্সিকিউশন বিভাগে উল্লিখিত সমস্ত ফ্ল্যাগ ব্যবহার করতে পারেন।

JAX-নির্দিষ্ট বিকল্পগুলি

স্থানীয় এক্সিকিউশনের মতোই; লাইভ, ইন্টারেক্টিভ ইন্ট্রোস্পেকশনের জন্য আপনি সরাসরি একটি কম্পিউটেশনের প্রি-অপ্টিমাইজড HLO প্রিন্ট করতে পারেন:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

আপনি সরাসরি একটি গণনার অপ্টিমাইজ করা HLO-ও প্রিন্ট করতে পারেন:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

সমস্ত/ছোট গণনা বাদ দেওয়া

যদি আপনি সমস্ত ছোটখাটো কম্পাইলেশন সহ ডাম্পের সবকিছু দেখতে চান, তাহলে JAX এনভায়রনমেন্ট ভেরিয়েবলটি সেট করুন:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

মোজাইক

Mosaic হলো Pallas TPU ব্যাকএন্ড এবং পরীক্ষামূলক Pallas GPU ব্যাকএন্ডের জন্য একটি কম্পাইলার। Mosaic কম্পিউটেশন ডাম্প করতে, নিম্নলিখিত ফ্ল্যাগটি সেট করুন:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

অথবা, TPU ইনিট আর্গুমেন্টগুলোকে একটি এনভায়রনমেন্ট ভেরিয়েবল হিসেবে সেট করুন:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

আরও জানতে Pallas এবং Mosaic সম্পর্কিত JAX ডকুমেন্টেশন দেখুন।

HLO ডাম্পস-এর সাথে আরও তথ্য

সঠিক গণনা খুঁজে বের করা

সাধারণত, অনেক গণনা ডাম্প করা হয়। ডাম্প করা ফাইলগুলোর নাম JAX, Tensorflow, বা PyTorch/XLA "গণনার নাম" দিয়ে স্পষ্টভাবে উল্লেখ করা থাকে, যা লগগুলিতে প্রদর্শিত হয় এবং এর ফলে প্রাসঙ্গিক HLO ফাইলগুলি সনাক্ত করা সহজ হয়। উদাহরণস্বরূপ:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

অন্যথায়, কোন মডিউলে নির্দিষ্ট সিম্বল বা কম্পিউটেশন রয়েছে তা দ্রুত শনাক্ত করতে আপনি ripgrep ব্যবহার করতে পারেন।

পরামর্শ: আপনার বাগ রিপোর্টে ডাম্প করা গুরুত্বপূর্ণ before/after/buffer-assignment ফাইলগুলো অন্তর্ভুক্ত করুন।

এইচএলও রূপান্তর

hlo-opt নামের একটি টুল যা HLOProto এবং টেক্সট ফরম্যাটের মধ্যে অনুবাদ করতে পারে। এটি এমন ক্ষেত্রে উপযোগী যেখানে আপনার কাছে একটি ফরম্যাট আছে, কিন্তু ডিবাগিংয়ের জন্য অন্যটির প্রয়োজন।

এর ব্যবহার শিখুন: XLA টুলিং ডকুমেন্টেশন: hlo-opt ।

পুনরায় চালান

আপনি নকল ডেটা বা ইনপুট স্ন্যাপশট ব্যবহার করে একটি নির্দিষ্ট XLA ব্যাকএন্ডে ডাম্প করা গণনাগুলো চালাতে (রিপ্লে করতে) পারেন। XLA-তে সমস্যাগুলো পুনরুৎপাদন, পুনরাবৃত্তি এবং ডিবাগ করার জন্য এটি একটি সুবিধাজনক উপায়।

নিম্নলিখিত কমান্ডগুলো ফেক ডেটা ব্যবহার করে। আপনার যদি HLO স্ন্যাপশট সেভ করা থাকে, তবে আপনি সেগুলোর পরিবর্তে সেগুলো পাস করতে পারেন এবং স্ন্যাপশটের ডেটা ব্যবহৃত হবে। স্ন্যাপশট চালানোর সময়ও ফেক ডেটা ব্যবহার করতে চাইলে, --force_fake_data ফ্ল্যাগটি পাস করুন।

সিপিইউ ব্যাকএন্ড:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

জিপিইউ ব্যাকএন্ড:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt