محاسبات HLO را تخلیه کنید

یک HLO dump یک نمایش متنی از ماژول های HLO در مراحل مختلف محاسبات است. برای اشکال‌زدایی مفید است و اغلب لازم است آن را در گزارش‌های باگ بگنجانید. این معمولاً یک فایل متنی قابل خواندن توسط انسان است که دستورالعمل های HLO و ویژگی های آنها را فهرست می کند. گاهی اوقات، ماژول های HLO به صورت زیر ریخته می شوند:

• HloProto: فایل‌های بافر پروتکل، که قالبی ساختارمندتر و قابل خواندن توسط ماشین هستند.
• HloSnapshot : ماژول HLO به اضافه ورودی های آن. برای پخش مجدد HLO ها، گاهی اوقات به ورودی های واقعی به جای داده های تصادفی نیاز به یک محاسبات معین دارید.

شما می توانید از پرچم های XLA برای مشخص کردن و دریافت dump استفاده کنید. در بیشتر موارد، می توانید آن را با یک متغیر محیطی تنظیم کنید. JAX همچنین یک روش برنامه ریزی شده برای چاپ HLO dump ارائه می دهد.

اجرای محلی

استفاده از متغیرهای محیطی

می‌توانید متغیر محیطی XLA_FLAGS را با پرچم‌های لازم برای دریافت dumps تنظیم کنید. این برای 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. رندر نمودار با سرور graphviz (فقط برای نمودارهای کوچک خوب کار می کند)

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 برای تجسم بخشی از نمودار استفاده کنید.

گذرهای میانی خاص را تخلیه کنید

علاوه بر HLOهای استاندارد از پیش بهینه شده / نهایی بهینه شده، همچنین می توانید وضعیت HLO ها را پس از عبور کامپایلر خاص تخلیه کنید.

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

ماژول های HLO برای پاس هایی که نام آنها با عبارت منظم (regex) مطابقت دارد، ریخته می شود. به عنوان مثال، می توانید HLO های حاصل از پاس های مربوط به پارتیشن بندی SPMD را با موارد زیر مشاهده کنید:

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

به‌جای ارسال پرچم‌ها یا متغیرهای محیطی، می‌توانید با استفاده از APIهای lower JAX و compile 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 درباره صادرات و سریال‌سازی محاسبات مرحله‌ای: اشکال‌زدایی بیشتر بیاموزید.

گوگل کولب

متغیرهای محیطی

در اولین سلول اجرا شده نوت بوک (چون متغیرهای محیطی و پرچم‌های خط فرمان معمولاً فقط یک بار پردازش می‌شوند، به عنوان مثال، در زمان وارد کردن ماژول یا زمان اولیه‌سازی باطن XLA)، XLA_FLAGS که در بالا توضیح داده شد با os.environ اضافه کنید، برای مثال:

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

این محاسبات را به DIRECTORY_PATH ، برای مثال /tmp ، منتقل می‌کند. در Colab، برای مشاهده و دسترسی به این فهرست، به مرورگر "Files" در نوار کناری سمت چپ بروید.

می توانید از تمام پرچم های ذکر شده در قسمت Local Execution استفاده کنید.

گزینه های مخصوص 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)

تخلیه همه / محاسبات کوچک

اگر می خواهید همه چیز را در یک dump مشاهده کنید، از جمله تمام کامپایل های کوچک، متغیر محیطی JAX را تنظیم کنید:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

موزاییک

Mosaic یک کامپایلر برای باطن TPU Pallas و باطن GPU آزمایشی Pallas است. برای حذف محاسبات موزاییک، پرچم زیر را تنظیم کنید:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

یا، آرگومان های اولیه TPU را به عنوان یک متغیر محیطی تنظیم کنید:

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

برای کسب اطلاعات بیشتر ، مستندات JAX در مورد Pallas and Mosaic را بررسی کنید.

بیشتر با HLO Dumps

پیدا کردن محاسبات مناسب

معمولاً بسیاری از محاسبات حذف می شوند. فایل‌های ریخته‌شده به صراحت با «نام محاسباتی» JAX، Tensorflow، یا PyTorch/XLA نام‌گذاری می‌شوند که در گزارش‌ها فراخوانی می‌شوند و شناسایی فایل‌های HLO مربوطه را آسان می‌کنند. به عنوان مثال:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

در غیر این صورت، می توانید از ripgrep برای تشخیص سریع اینکه کدام ماژول دارای نمادها یا محاسبات خاص است استفاده کنید.

نکته: 3 فایل مورد علاقه قبل/پس از/تخصیص بافر را در گزارش های اشکال خود قرار دهید.

تبدیل HLO

ابزاری به نام hlo-opt که می تواند بین فرمت های HLOProto و متن ترجمه کند. در مواردی که یک فرمت دارید، اما برای اشکال زدایی به دیگری نیاز دارید، مفید است.

استفاده از آن را بیاموزید: اسناد XLA Tooling: hlo-opt .

پخش مجدد

می‌توانید محاسبات تخلیه‌شده را روی یک باطن خاص XLA با داده‌های جعلی یا عکس‌های فوری ورودی اجرا کنید. این یک راه راحت برای بازتولید، تکرار، و اشکال زدایی مسائل در XLA است.

دستورات زیر از داده های جعلی استفاده می کنند. اگر عکس‌های فوری HLO را ذخیره کرده‌اید، می‌توانید آن‌ها را به جای آن ارسال کنید و از داده‌های عکس فوری استفاده می‌شود. برای اینکه همچنان از داده‌های جعلی هنگام اجرای عکس فوری استفاده کنید، پرچم --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