پروفایلینگ تماس سفارشی، پروفایلینگ تماس سفارشی

فراخوانی‌های سفارشی XLA به شما امکان می‌دهند هسته‌های سفارشی یا عملیاتی را اجرا کنید که به طور طبیعی توسط XLA پشتیبانی نمی‌شوند. برای مشاهده عملکرد این فراخوانی‌های سفارشی در Trace Viewer ، می‌توانید از پرچم‌های خاص XLA برای فعال کردن ردیابی دقیق و اطلاعات اشکال‌زدایی LLO (بهینه‌ساز سطح پایین) استفاده کنید.

⚠️ ویژگی آزمایشی : تحلیل بهینه‌ساز سطح پایین (LLO) و پروفایلینگ تماس سفارشی آزمایشی هستند. برای دسترسی به این ویژگی‌ها و تمام ابزارهای تحلیل CLI ( get_kernel_stats ، get_llo_analysis ، get_llo_debug_string )، xprof-nightly را نصب کنید . نسخه استاندارد xprof PyPI (2.23.1) فاقد این زیردستورات است.

پیش‌نیازها و الزامات زنجیره ابزار

قبل از ثبت ردپای LLO، بررسی کنید که محیط شما الزامات زیر را برآورده می‌کند:

• پایتون ۳.۱۱+ (پایتون ۳.۱۲ توصیه می‌شود) : ایمیج‌های پیش‌فرض ماشین مجازی ابری TPU (اوبونتو ۲۲.۰۴) با سیستم پایتون ۳.۱۰.۱۲ ارائه می‌شوند که بی‌سروصدا JAX را در نسخه ۰.۶.۲ نگه می‌دارد و libtpu ۰.۰.۱۷ را دریافت می‌کند. نسخه‌های قدیمی‌تر libtpu فاقد تعاریف پرچم LLO هستند و باعث ERROR: Unknown command line flag یا بازگرداندن پروفایل‌های LLO خالی می‌شوند. استفاده از uv برای مدیریت محیط مجازی پایتون ۳.۱۲ اکیداً توصیه می‌شود.

• نصب بسته ( xprof-nightly ) : xprof-nightly را در کنار jax[tpu] نصب کنید:

  # 1. Setup Python 3.12 environment
  pip install uv
  uv python install 3.12
  uv venv --python 3.12 ~/venvs/v312
  
  # 2. Install xprof-nightly and JAX
  uv pip install --python ~/venvs/v312/bin/python \
      'jax[tpu]>=0.11.0' xprof-nightly numpy ml_dtypes absl-py fire
  

• JAX >= 0.11.0 : نسخه پیشنهادی toolchain ( libtpu >= 0.0.44 ). توجه داشته باشید که اطلاعات اشکال‌زدایی LLO در زمان کامپایل ( --xla_xprof_register_llo_debug_info=true ) از jax >= 0.10.2 ( libtpu >= 0.0.42 ) پشتیبانی می‌شود، در حالی که ردیابی فراخوانی سفارشی زمان اجرا ( --xla_xprof_enable_custom_call_tracing=true ) به jax >= 0.11.0 ( libtpu >= 0.0.44 ) نیاز دارد.

• مرتب‌سازی دقیق محیطی : LIBTPU_INIT_ARGS باید قبل از import jax در پوسته export شود یا در os.environ پیکربندی شود. libtpu پرچم‌های مقداردهی اولیه را در اولین import JAX تجزیه می‌کند؛ تنظیم بی‌سروصدای آنها پس از import jax بدون ایجاد استثنا هیچ تاثیری ندارد.

ماتریس سازگاری سخت‌افزار

تشخیص در دسترس بودن پرچم

برای تأیید اینکه فایل باینری libtpu نصب‌شده شما شامل تعاریف پرچم مورد نیاز قبل از راه‌اندازی بارهای کاری است، این قطعه کد تشخیصی را اجرا کنید:

import glob
import os
import libtpu

so_paths = glob.glob(os.path.dirname(libtpu.__file__) + "/*libtpu*.so")
if so_paths:
  blob = open(so_paths[0], "rb").read()
  for flag in (
      b"xla_xprof_register_llo_debug_info",
      b"xla_xprof_enable_custom_call_tracing",
      b"tpu_enable_periodic_counter_sampling",
  ):
    print(flag.decode(), "PRESENT" if flag in blob else "ABSENT")

نحوه فعال کردن ردیابی

پیش‌فرض توصیه‌شده: فقط اطلاعات اشکال‌زدایی LLO

برای تحلیل استاتیک LLO و پروفایل‌سازی HLO/سطح هسته، اطلاعات اشکال‌زدایی LLO را با --xla_xprof_register_llo_debug_info=true ثبت کنید. این کار جریان عملیاتی کامل HLO را دست‌نخورده نگه می‌دارد، بنابراین get_hlo_stats ، get_roofline_model ، get_top_hlo_ops و get_kernel_stats همگی بدون سربار اضافی ردیابی زمان اجرا کار می‌کنند، در حالی که نقشه منبع LLO زمان کامپایل کامل مورد استفاده توسط get_llo_analysis و get_llo_debug_string تولید می‌کنند.

import os

# Flags MUST precede any jax / libtpu import
os.environ["LIBTPU_INIT_ARGS"] = "--xla_xprof_register_llo_debug_info=true"

import jax
# Workload definition and tracing...

• --xla_xprof_register_llo_debug_info=true : اطلاعات اشکال‌زدایی LLO، کدهای عملیاتی و فراداده‌ها را برای مصورسازی XProf ثبت می‌کند.

ردیابی دقیق بسته LLO در زمان اجرا

• --xla_xprof_enable_custom_call_tracing=true : پرچم متعارفی که جزئیات اجرای LLO در زمان اجرا ( Pallas Primitives ، LLO Ops و خطوط دستورالعمل به ازای هر واحد در Trace Viewer) را به صورت دقیق فعال می‌کند و به طور خودکار ابزار دقیق بسته دستورالعمل ( xla_tpu_bundle_instrumentation_options با مقادیر پیش‌فرض trace_best_effort_frequency=10 و trace_guaranteed_frequency=10 ) را فعال می‌کند.

هنگام استفاده از --xla_xprof_enable_custom_call_tracing=true ، اگر افزایش اندازه ردیابی از بافر ردیابی سخت‌افزار سرریز کند، فرکانس vtrace ( trace_best_effort_frequency و trace_guaranteed_frequency در xla_tpu_bundle_instrumentation_options ؛ به نحوه تنظیم در زیر مراجعه کنید) را تنظیم کنید.

نمایشگر ردیابی نمونه

در اینجا مثالی از نحوه نمایش ردپاهای LLO در نمایشگر ردپای Xprof آورده شده است:

پارامترهای پیشرفته (مدیریت حذف رویداد)

اگر در Xprof شاهد افت رویداد یا سرریز بافر هستید، به این معنی است که نقاط ردیابی بیش از حد فعال می‌شوند و بافرهای ردیابی سخت‌افزاری را تحت فشار قرار می‌دهند. می‌توانید فرکانس درج ردیابی LLO را با استفاده از پارامترهای پیشرفته تنظیم کنید.

این پارامترها از طریق xla_tpu_bundle_instrumentation_options پیکربندی می‌شوند. شما می‌توانید کنترل کنید که ردپاها چند وقت یکبار در بسته‌های دستورالعمل بسته‌بندی شوند.

پارامترهای کلیدی

• trace_best_effort_frequency (پیش‌فرض: ۱۰): بازه هدف (به صورت بسته) برای درج ردهای فرصت‌طلبانه بسته‌بندی‌شده در بسته‌های موجود. کامپایلر اغلب سعی می‌کند ردی را درج کند اما بسته‌های جدیدی برای آن ایجاد نمی‌کند .
• trace_guaranteed_frequency (پیش‌فرض: ۱۰): حداکثر تعداد بسته‌های مجاز بین دو ردیابی. این یک تضمین است. هر زمان که نتوانیم با بسته‌بندی ردیابی‌ها در بسته‌های موجود، این مورد را برآورده کنیم، یک بسته جدید ایجاد می‌کنیم و یک ردیابی را (به خودی خود) در آنجا قرار می‌دهیم.

چگونه کوک کنیم

• اگر با خطای Event Drops مواجه شدید : برای ردیابی کمتر ، مقادیر را افزایش دهید (مثلاً روی ۵۰ یا ۱۰۰ تنظیم کنید) تا حجم داده‌های ردیابی تولید شده کاهش یابد.
• اگر به جزئیات دقیق‌تری نیاز دارید : مقادیری را که باید بیشتر ردیابی شوند، کاهش دهید (به قیمت سربار بیشتر و سرریز احتمالی بافر).

نحوه محاسبه تعداد چرخه‌های دستورالعمل

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

کامپایلر هزینه چرخه سخت‌افزاری ذاتی هر دستورالعمل LLO را بر اساس تولید TPU هدف و واحد اجرایی که آن را حل می‌کند، محاسبه می‌کند. این تعداد چرخه نشان دهنده توان عملیاتی اجرا و تأخیرهای تأخیر است.

جریان سطح بالا

1. تجزیه دستورالعمل LLO : شناسایی کد عملیاتی و فراداده.
2. دریافت چرخه‌های سخت‌افزار پایه : چرخه‌ها را بر اساس نسل TPU (v5e/v5p، v6e/v7x و غیره) تعیین کنید.
3. تبدیل به تیک‌های GTC : چرخه‌ها را با استفاده از فرمول زیر به تیک‌های شمارنده تایمر جهانی (GTC) تبدیل کنید: Cycles * (GTC_Freq * 16) / TC_Freq .
4. ایجاد بازه زمانی : رویدادهای میانی را به طور مساوی بین مرزهای ردیابی شناخته شده درون‌یابی کنید.

تخمین چرخه بر اساس واحد و نسل

در زیر نمونه‌هایی از نحوه مدل‌سازی چرخه‌های سخت‌افزار پایه برای واحدهای اجرایی مختلف آمده است:

واحد ضرب ماتریس (MXU)

تعداد چرخه‌های MXU، توان عملیاتی را بر اساس تراکم نوع داده نشان می‌دهد.

واحد انتقال (XLU)

تعداد چرخه‌ها نشان‌دهنده‌ی طرح حافظه‌ی انتقال و تأخیرهای ضربدری است.

مجموعه واحدهای اجرایی (EUP)

دستورالعمل‌های EUP توابع ریاضی برداری (مثلاً tanh ، log ، exp ) را نشان می‌دهند.

مهاجرت و آشتی پرچم

نسخه‌های اولیه مستندات XLA و TPU به پرچم قدیمی --xla_enable_custom_call_region_trace=true ارجاع می‌دادند.

• پرچم متعارف : --xla_xprof_enable_custom_call_tracing (نام متعارف زمانی که طول جدول زمانی بسته درون هسته زمان اجرا مورد نظر است؛ به طور پیش‌فرض از --xla_xprof_register_llo_debug_info=true به تنهایی استفاده کنید). وقتی فعال باشد، ردیابی تماس سفارشی را فعال می‌کند و در عین حال به طور خودکار ابزار دقیق بسته دستورالعمل مورد نیاز و فرکانس‌های ردیابی ( xla_tpu_bundle_instrumentation_options ) را پیکربندی می‌کند.
• پرچم قدیمی : --xla_enable_custom_call_region_trace=true (نام مستعار منسوخ شده). اگرچه هنوز توسط کامپایلرهای قدیمی‌تر پشتیبانی می‌شود، کاربرانی که به بازه‌های زمانی بسته نرم‌افزاری زمان اجرا نیاز دارند باید به --xla_xprof_enable_custom_call_tracing مهاجرت کنند.

مثال ضبط پیش‌فرض (اطلاعات اشکال‌زدایی LLO را ثبت می‌کند در حالی که HLO و آمار هسته را دست‌نخورده نگه می‌دارد):

export LIBTPU_INIT_ARGS="--xla_xprof_register_llo_debug_info=true"
python your_jax_workload.py

وقتی ردیابی تماس سفارشی در کنار اطلاعات اشکال‌زدایی LLO فعال شود، یک خط جدید برای استفاده از LLO در Trace Viewer برای هر هسته TPU یا دستگاهی که تماس سفارشی را اجرا می‌کند، ظاهر می‌شود.

خط استفاده از LLO

خط استفاده از LLO، تصویری از نحوه استفاده از منابع سخت‌افزاری در طول اجرای یک فراخوانی سفارشی ارائه می‌دهد. این امر به ویژه برای شناسایی گلوگاه‌ها در هسته‌های سفارشی (مثلاً آن‌هایی که در پالاس یا موزائیک نوشته شده‌اند) مفید است.

بهترین شیوه‌ها و مشکلات میدانی

• استفاده از xprof-nightly : نسخه استاندارد xprof 2.23.1 فاقد دستورات فرعی get_kernel_stats و LLO CLI (و همچنین اسکریپت مستقل کنسول xparity برای تأیید برابری عددی) است. در محیط‌های غیر Google3، همیشه xprof-nightly را نصب کنید.
• تفسیر معیارها برای هسته‌های پالاس (نقطه کور خط سقف) : XLA هیچ مدل هزینه‌ای برای tpu_custom_call ندارد. بنابراین، get_roofline_model و get_overview حتی زمانی که دستورالعمل‌های LLO به طور کامل ضبط و روی سخت‌افزار اجرا می‌شوند، 0.0 GFLOP/s ، "bound_by": "Unknown" و 0.0% استفاده از MXU را گزارش خواهند کرد.
  • برای مدت زمان و تأخیر هسته، از xprof get_kernel_stats <logdir> استفاده کنید.
  • برای تجزیه و تحلیل اجرای دستورالعمل‌های سطح پایین و تخمین چرخه، از xprof get_llo_analysis <logdir> استفاده کنید.
• بدنه‌های حلقه داخلی Lite Proto : در get_llo_debug_string ، بدنه‌های حلقه داخلی به صورت // Loop body not available in lite proto خلاصه می‌شوند. این بدنه، ساختار ماژول اطراف، تخصیص ثبات و توالی دستورالعمل‌های بیرونی را فراهم می‌کند.
• اکتشاف اعتبارسنجی ردیابی : برای تعیین وجود داده‌های LLO، نام‌های خط ردیابی مانند SALU / VALU / EUP / XLU / VLD / VST / MXU Instructions را بررسی نکنید . ردیابی‌های معتبر LLO از این نام‌های خط استفاده نمی‌کنند. ضبط LLO را با اجرای xprof get_llo_analysis <logdir> و تأیید "success": true اعتبارسنجی کنید.
• اندازه گیری هسته ضبط : اندازه گیری هسته های تست/ضبط بیش از حد بزرگ می تواند خطاهای کامپایلر مانند CompileTimeScopedVmemOom: Scoped allocation with size 32.81M and limit 32.00M exceeded scoped vmem limit . ماتریس های ضبط را با اندازه محافظه کارانه نگه دارید (مثلاً (512, 512, 1024) f32).
• محیط‌های مجازی جداگانه : هنگام انتقال هسته‌ها بین نسخه‌های JAX (مثلاً نسخه‌های منسوخ‌شده JAX 0.11+ مانند pltpu.repeat )، محیط‌های مجازی اختصاصی را حفظ کنید.