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

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

• HloProto: প্রোটোকল বাফার ফাইল, যা আরও কাঠামোগত, মেশিন-পাঠযোগ্য বিন্যাস।
• HloSnapshot : 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. HLO প্রোটোস

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

1. HLO স্ন্যাপশট

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. এইচটিএমএল ফাইলে গ্রাফ রেন্ডার (শুধুমাত্র ছোট গ্রাফের জন্য ভাল কাজ করে)

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

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

স্পেসিফিক ইন্টারমিডিয়েট পাস ডাম্প করুন

স্ট্যান্ডার্ড প্রাক-অপ্টিমাইজড/ফাইনাল-অপ্টিমাইজড এইচএলও ছাড়াও, আপনি একটি নির্দিষ্ট কম্পাইলার পাসের পরেও এইচএলও-এর অবস্থা ডাম্প করতে পারেন।

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

এইচএলও মডিউলগুলি সেই পাসগুলির জন্য ডাম্প করা হবে যার নাম রেগুলার এক্সপ্রেশন (রেজেক্স) এর সাথে মেলে৷ উদাহরণস্বরূপ, আপনি 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)

ডাম্প jaxprs

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

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

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

Google Colab

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

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

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

এটি কম্পিউটেশনটিকে DIRECTORY_PATH এ ডাম্প করবে, উদাহরণস্বরূপ /tmp । Colab-এ, এই ডিরেক্টরিটি দেখতে এবং অ্যাক্সেস করতে বাম সাইডবারে "ফাইলস" ব্রাউজারে নেভিগেট করুন।

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

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

মোজাইক

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

--xla_mosaic_dump_to=/tmp/mosaic_dumps

অথবা, পরিবেশ পরিবর্তনশীল হিসাবে TPU init আর্গুমেন্ট সেট করুন:

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

আরও জানতে প্যালাস এবং মোজাইকের JAX ডকুমেন্টেশন দেখুন।

এইচএলও ডাম্পের সাথে আরও

সঠিক গণনা খোঁজা

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

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

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

টিপ: আপনার বাগ রিপোর্টে আগ্রহের আগে/পরে/বাফার-অ্যাসাইনমেন্ট ফাইল ডাম্প করা 3টি অন্তর্ভুক্ত করুন।

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

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

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

রিপ্লে

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

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

CPU ব্যাকএন্ড:

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

GPU ব্যাকএন্ড:

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