فراخوانیهای سفارشی XLA به شما امکان میدهند هستههای سفارشی یا عملیاتی را اجرا کنید که به طور طبیعی توسط XLA پشتیبانی نمیشوند. برای مشاهده عملکرد این فراخوانیهای سفارشی در Trace Viewer ، میتوانید از پرچمهای خاص XLA برای فعال کردن ردیابی دقیق و اطلاعات اشکالزدایی LLO (بهینهساز سطح پایین) استفاده کنید.
نحوه فعال کردن ردیابی
برای کامپایل مجدد حجم کار خود با ابزار دقیق، پرچم XLA زیر را تنظیم کنید:
--xla_xprof_enable_custom_call_tracing=true
اگر از طریق LIBTPU_INIT_ARGS روی Cloud TPU در حال اجرا هستید، میتوانید آن را به صورت زیر ارسال کنید:
LIBTPU_INIT_ARGS="--xla_xprof_enable_custom_call_tracing=true --xla_xprof_register_llo_debug_info=true" python your_jax_workload.py
نمایشگر ردیابی نمونه
در اینجا مثالی از نحوه نمایش ردپاهای LLO در نمایشگر ردپای Xprof آورده شده است:
پارامترهای پیشرفته (مدیریت حذف رویداد)
اگر در Xprof شاهد افت رویداد یا سرریز بافر هستید، به این معنی است که نقاط ردیابی بیش از حد فعال میشوند و بافرهای ردیابی سختافزاری را تحت فشار قرار میدهند. میتوانید فرکانس درج ردیابی LLO را با استفاده از پارامترهای پیشرفته تنظیم کنید.
این پارامترها از طریق xla_tpu_bundle_instrumentation_options پیکربندی میشوند. شما میتوانید کنترل کنید که ردپاها چند وقت یکبار در بستههای دستورالعمل بستهبندی شوند.
پارامترهای کلیدی:
-
trace_best_effort_frequency(پیشفرض: ۱۰): بازه هدف (به صورت بسته) برای درج ردهای فرصتطلبانه بستهبندیشده در بستههای موجود. کامپایلر اغلب سعی میکند ردی را درج کند اما بستههای جدیدی برای آن ایجاد نمیکند . -
trace_guaranteed_frequency(پیشفرض: ۱۰): حداکثر تعداد بستههای مجاز بین دو ردیابی. این یک تضمین است. هر زمان که نتوانیم با بستهبندی ردیابیها در بستههای موجود، این مورد را برآورده کنیم، یک بسته جدید ایجاد میکنیم و یک ردیابی را (به خودی خود) در آنجا قرار میدهیم.
نحوه تنظیم:
- اگر با خطای Event Drops مواجه شدید : برای ردیابی کمتر ، مقادیر را افزایش دهید (مثلاً روی ۵۰ یا ۱۰۰ تنظیم کنید) تا حجم دادههای ردیابی تولید شده کاهش یابد.
- اگر به جزئیات دقیقتری نیاز دارید : مقادیری را که باید بیشتر ردیابی شوند، کاهش دهید (به قیمت سربار بیشتر و سرریز احتمالی بافر).
نحوه محاسبه تعداد چرخههای دستورالعمل
از آنجا که نقاط ردیابی به صورت فرصتطلبانه تزریق میشوند و نه در هر دستورالعمل واحد، مهرهای زمانی میانی بر اساس هزینههای چرخه سختافزاری تخمینی درونیابی میشوند.
کامپایلر هزینه چرخه سختافزاری ذاتی هر دستورالعمل LLO را بر اساس تولید TPU هدف و واحد اجرایی که آن را حل میکند، محاسبه میکند. این تعداد چرخه نشان دهنده توان عملیاتی اجرا و تأخیرهای تأخیر است.
جریان سطح بالا
- تجزیه دستورالعمل LLO : شناسایی کد عملیاتی و فراداده.
- دریافت چرخههای سختافزار پایه : چرخهها را بر اساس نسل TPU (v5e/v5p، v6e/v7x و غیره) تعیین کنید.
- تبدیل به تیکهای GTC : چرخهها را با استفاده از فرمول زیر به تیکهای شمارنده تایمر جهانی (GTC) تبدیل کنید:
Cycles * (GTC_Freq * 16) / TC_Freq. - ایجاد بازه زمانی : رویدادهای میانی را به طور مساوی بین مرزهای ردیابی شناخته شده درونیابی کنید.
تخمین چرخه بر اساس واحد و نسل
در زیر نمونههایی از نحوه مدلسازی چرخههای سختافزار پایه برای واحدهای اجرایی مختلف آمده است:
واحد ضرب ماتریس (MXU)
تعداد چرخههای MXU، توان عملیاتی را بر اساس تراکم نوع داده نشان میدهد.
| دسته بندی دستورالعمل | زیرنوع / قالب | (نسخه ۵e/نسخه ۵p) | (نسخه ۶/۷) |
|---|---|---|---|
| وکتور ماتمول | اف۳۲ | ۸ | ۸ |
| پیشپردازش Matmul (F8 تا BF16) | ۴ | ۴ | |
| بسته بندی شده BF16 | ۲ | ۲ | |
| فرمت های عدد صحیح (U8، S8، U4، S4) | ۱ | ۱ | |
| لچهای برداری | F32 جابجا شده | ۴ | ۴ |
| BF16 جابجا شده | ۸ | ۸ | |
| F32 بدون جابجایی | ۲ | ۲ | |
| BF16 بدون جابجایی | ۴ | ۴ | |
| آماده سازی تشک / Dwg | همه | ۱ | ۱ |
واحد انتقال (XLU)
تعداد چرخهها نشاندهندهی طرح حافظهی انتقال و تأخیرهای ضربدری است.
| دسته بندی دستورالعمل | زیرنوع / قالب | (نسخه ۵e/نسخه ۵p) | (نسخه ۶/۷) |
|---|---|---|---|
| جابجایی بستهبندیشده | همه | ۱۷ | ۴ |
| ترانهاده استاندارد | B32 ترانسپوز | ۹ | ۴ |
| B16 انتقال (قطعه قطعه/فشرده) | ۱۷ | ۴ |
مجموعه واحدهای اجرایی (EUP)
دستورالعملهای EUP توابع ریاضی برداری (مثلاً tanh ، log ، exp ) را نشان میدهند.
| دسته بندی دستورالعمل | (نسخه ۵e/نسخه ۵p) | (نسخه ۶/۷) |
|---|---|---|
ریاضی برداری ( tanh ، exp و غیره) | ۲ | ۱ |
(قدیمی) فعال کردن قابلیت مشاهده تماس سفارشی
برای فعال کردن پروفایلینگ تماس سفارشی، باید هنگام اجرای بار کاری خود، پرچمهای XLA زیر را تنظیم کنید:
-
--xla_enable_custom_call_region_trace=true: این پرچم ردیابی مناطقی را که حاوی فراخوانیهای سفارشی هستند، فعال میکند. -
--xla_xprof_register_llo_debug_info=true: این پرچم اطلاعات اشکالزدایی LLO را ثبت میکند، که به XProf اجازه میدهد آمار دقیق استفاده را برای فراخوانی سفارشی نمایش دهد.
مثال:
LIBTPU_INIT_ARGS="--xla_enable_custom_call_region_trace=true --xla_xprof_register_llo_debug_info=true" python your_jax_workload.py
وقتی این پرچمها فعال شوند، یک خط جدید استفاده از LLO در Trace Viewer برای هر هسته یا دستگاه TPU که فراخوانی سفارشی را اجرا میکند، ظاهر میشود.
خط استفاده از LLO
خط استفاده از LLO، تصویری از نحوه استفاده از منابع سختافزاری در طول اجرای یک فراخوانی سفارشی ارائه میدهد. این امر به ویژه برای شناسایی گلوگاهها در هستههای سفارشی (مثلاً آنهایی که در پالاس یا موزائیک نوشته شدهاند) مفید است.
بهترین شیوهها
- فقط در صورت نیاز فعال کنید : این پرچمها میتوانند اندازه پروفایل ضبط شده را افزایش دهند و ممکن است کمی بر عملکرد در طول جمعآوری تأثیر بگذارند. از آنها در درجه اول برای اشکالزدایی و بهینهسازی فراخوانیهای سفارشی استفاده کنید.
- بررسی اطلاعات LLO : اگر این پرچمها را فعال کردهاید اما خط مربوط به استفاده از LLO را نمیبینید، مطمئن شوید که کامپایلر شما از ثبت اطلاعات اشکالزدایی LLO برای پیادهسازی فراخوانی سفارشی خاص شما پشتیبانی میکند.