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

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

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

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

اجرای محلی

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

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

همچنین می‌توانید با تنظیم پرچم --xla_syntax_sugar_async_ops روی true ، HLO Dumps را طوری تنظیم کنید که از بسته‌های قندی نحوی به عنوان نام‌های عملیاتی استفاده کند. این کار می‌تواند حجم dump را حدود 20٪ کاهش دهد. به طور پیش‌فرض، این پرچم روی false تنظیم شده است و از نام‌های عملیاتی واقعی در dump استفاده می‌شود.

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

گذرگاه‌های میانی مخصوص تخلیه

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

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

به جای ارسال پرچم‌ها یا متغیرهای محیطی، می‌توانید HLO را با استفاده از APIهای lower و compile JAX به صورت برنامه‌نویسی شده dump کنید.

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 در مورد Export کردن و سریال‌سازی محاسبات مرحله‌ای مراجعه کنید: اشکال‌زدایی .

گوگل کولب

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

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

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)

حذف همه/محاسبات کوچک

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

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

موزاییک

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

--xla_mosaic_dump_to=/tmp/mosaic_dumps

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

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

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

بیشتر با HLO Dumps

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

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

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

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

نکته: سه فایل مورد نظر قبل/بعد/بافر-انتساب را که در گزارش‌های باگ خود قرار داده‌اید، در گزارش‌های باگ خود قرار دهید.

تبدیل HLO

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

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

پخش مجدد

شما می‌توانید محاسبات ذخیره شده را در یک بک‌اند مشخص XLA با داده‌های جعلی یا اسنپ‌شات‌های ورودی اجرا (بازپخش) کنید. این یک روش مناسب برای بازتولید، تکرار و اشکال‌زدایی مشکلات در XLA است.

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