پروفایل‌بندی محاسبات JAX با XProf

XProf روشی عالی برای به دست آوردن و مصورسازی ردپاها و پروفایل‌های عملکرد برنامه شما، از جمله فعالیت روی GPU و TPU، است. نتیجه نهایی چیزی شبیه به این خواهد بود:

ضبط برنامه‌ای

شما می‌توانید کد خود را برای ثبت ردیابی پروفایلر برای کد JAX از طریق متدهای jax.profiler.start_trace و jax.profiler.stop_trace ابزارک‌سازی کنید. jax.profiler.start_trace را با دایرکتوری که فایل‌های ردیابی در آن نوشته می‌شوند، فراخوانی کنید. این باید همان دایرکتوری --logdir باشد که برای شروع XProf استفاده می‌شود. سپس، می‌توانید XProf را برای مشاهده ردیابی‌ها اجرا کنید.

برای مثال، برای گرفتن رد یک پروفایلر:

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

به فراخوانی jax.block_until_ready توجه کنید. ما از این برای اطمینان از ثبت اجرای روی دستگاه توسط ردیابی استفاده می‌کنیم. برای جزئیات بیشتر در مورد لزوم این کار، به اعزام ناهمزمان مراجعه کنید.

همچنین می‌توانید از مدیر زمینه‌ی jax.profiler.trace به عنوان جایگزینی برای start_trace و stop_trace استفاده کنید:

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

مشاهده‌ی ردپا

پس از ثبت یک مسیر، می‌توانید آن را با استفاده از رابط کاربری XProf مشاهده کنید.

شما می‌توانید رابط کاربری پروفایلر را مستقیماً با استفاده از دستور مستقل XProf و با ارجاع آن به دایرکتوری لاگ خود، اجرا کنید:

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

برای مشاهده پروفایل، به آدرس اینترنتی ارائه شده (مثلاً http://localhost:8791/ ) در مرورگر خود بروید.

ردپاهای موجود در منوی کشویی "Sessions" در سمت چپ ظاهر می‌شوند. جلسه مورد نظر خود را انتخاب کنید و سپس در زیر منوی کشویی "Tools"، "Trace Viewer" را انتخاب کنید. اکنون باید یک جدول زمانی از اجرا را مشاهده کنید. می‌توانید از کلیدهای WASD برای پیمایش ردپا استفاده کنید و برای جزئیات بیشتر، برای انتخاب رویدادها کلیک یا بکشید. برای جزئیات بیشتر در مورد استفاده از نمایشگر ردپا، به مستندات ابزار Trace Viewer مراجعه کنید.

ضبط دستی از طریق XProf

دستورالعمل‌های زیر برای ثبت یک مسیر N ثانیه‌ایِ فعال‌شده به صورت دستی از یک برنامه‌ی در حال اجرا ارائه شده است.

1. یک سرور XProf راه‌اندازی کنید:

   xprof --logdir /tmp/profile-data/
   

   شما باید بتوانید XProf را در <http://localhost:8791/> بارگذاری کنید. می‌توانید با استفاده از پرچم --port پورت دیگری را مشخص کنید.

2. در برنامه یا فرآیند پایتونی که می‌خواهید پروفایل آن را مشخص کنید، کد زیر را در جایی نزدیک به ابتدا اضافه کنید:

   import jax.profiler
   jax.profiler.start_server(9999)
   

   این کار سرور پروفایلر را که XProf به آن متصل است، آغاز می‌کند. سرور پروفایلر باید قبل از رفتن به مرحله بعدی در حال اجرا باشد. وقتی استفاده از سرور تمام شد، می‌توانید تابع jax.profiler.stop_server() را برای خاموش کردن آن فراخوانی کنید.

   اگر می‌خواهید قطعه‌ای از یک برنامه‌ی طولانی‌مدت (مثلاً یک حلقه‌ی آموزشی طولانی) را پروفایل کنید، می‌توانید این را در ابتدای برنامه قرار دهید و برنامه‌ی خود را طبق معمول شروع کنید. اگر می‌خواهید یک برنامه‌ی کوتاه‌مدت (مثلاً یک میکروبنچمارک) را پروفایل کنید، یک گزینه این است که سرور پروفایلر را در یک پوسته‌ی IPython شروع کنید و برنامه‌ی کوتاه را با %run پس از شروع ضبط در مرحله‌ی بعدی اجرا کنید. گزینه‌ی دیگر این است که سرور پروفایلر را در ابتدای برنامه شروع کنید و از time.sleep() استفاده کنید تا زمان کافی برای شروع ضبط داشته باشید.

3. <http://localhost:8791/> را باز کنید و روی دکمه "CAPTURE PROFILE" در بالا سمت چپ کلیک کنید. "localhost:9999" را به عنوان URL سرویس پروفایل وارد کنید (این آدرس سرور پروفایلری است که در مرحله قبل شروع کردید). تعداد میلی ثانیه‌هایی را که می‌خواهید پروفایل کنید وارد کنید و روی "CAPTURE" کلیک کنید.

4. اگر کدی که می‌خواهید پروفایل کنید، هنوز اجرا نشده است (مثلاً اگر سرور پروفایلر را در یک پوسته پایتون راه‌اندازی کرده‌اید)، آن را هنگام اجرای ضبط اجرا کنید.

5. پس از اتمام ضبط، XProf باید به طور خودکار رفرش شود. (همه ویژگی‌های پروفایل XProf با JAX مرتبط نیستند، بنابراین ممکن است در ابتدا به نظر برسد که چیزی ضبط نشده است.) در سمت چپ، در زیر «ابزارها»، «نمایشگر ردیابی» را انتخاب کنید.

اکنون باید یک جدول زمانی از اجرا را مشاهده کنید. می‌توانید از کلیدهای WASD برای پیمایش در مسیر استفاده کنید و برای مشاهده جزئیات بیشتر در پایین، روی رویدادها کلیک یا بکشید. برای جزئیات بیشتر در مورد استفاده از نمایشگر مسیر، به مستندات ابزار نمایشگر مسیر مراجعه کنید.

XProf و Tensorboard

XProf ابزار اساسی است که قابلیت پروفایلینگ و ردیابی را در Tensorboard فراهم می‌کند. تا زمانی که xprof نصب باشد، یک تب "Profile" در Tensorboard وجود خواهد داشت. استفاده از این تب مانند اجرای مستقل XProf است، البته تا زمانی که با اشاره به همان دایرکتوری لاگ اجرا شود. این شامل ضبط پروفایل، تجزیه و تحلیل و مشاهده قابلیت‌ها می‌شود. XProf جایگزین قابلیت tensorboard_plugin_profile می‌شود که قبلاً توصیه می‌شد.

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

افزودن رویدادهای ردیابی سفارشی

به طور پیش‌فرض، رویدادهای موجود در نمایشگر ردیابی، عمدتاً توابع داخلی JAX سطح پایین هستند. می‌توانید رویدادها و توابع خودتان را با استفاده از jax.profiler.TraceAnnotation و jax.profiler.annotate_function در کد خود اضافه کنید.

پیکربندی گزینه‌های پروفایلر

متد start_trace یک پارامتر اختیاری profiler_options را می‌پذیرد که امکان کنترل دقیق بر رفتار profiler را فراهم می‌کند. این پارامتر باید نمونه‌ای از jax.profiler.ProfileOptions باشد.

برای مثال، برای غیرفعال کردن تمام ردپاهای پایتون و میزبان:

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

گزینه‌های عمومی

1. host_tracer_level : سطح ردیابی را برای فعالیت‌های سمت میزبان تنظیم می‌کند.

   مقادیر پشتیبانی شده:

   • 0 : ردیابی میزبان (CPU) را به طور کامل غیرفعال می‌کند.
   • 1 : ردیابی رویدادهای TraceMe که فقط توسط کاربر تنظیم شده‌اند را فعال می‌کند.
   • 2 : شامل ردیابی‌های سطح ۱ به علاوه جزئیات اجرای برنامه سطح بالا مانند عملیات XLA پرهزینه (پیش‌فرض).
   • 3 : شامل ردیابی‌های سطح ۲ به علاوه جزئیات اجرای برنامه سطح پایین و مفصل‌تر مانند عملیات XLA ارزان است.

2. device_tracer_level : کنترل می‌کند که آیا ردیابی دستگاه فعال باشد یا خیر.

   مقادیر پشتیبانی شده:

   • 0 : ردیابی دستگاه را غیرفعال می‌کند.
   • 1 : ردیابی دستگاه را فعال می‌کند (پیش‌فرض).

3. python_tracer_level : کنترل می‌کند که آیا ردیابی پایتون فعال باشد یا خیر.

   مقادیر پشتیبانی شده:

   • 0 : ردیابی فراخوانی تابع پایتون را غیرفعال می‌کند (پیش‌فرض).
   • 1 : ردیابی پایتون را فعال می‌کند.

گزینه‌های پیکربندی پیشرفته

گزینه‌های TPU

1. tpu_trace_mode : حالت ردیابی TPU را مشخص می‌کند.

   مقادیر پشتیبانی شده:

   • TRACE_ONLY_HOST : این به این معنی است که فقط فعالیت‌های سمت میزبان (CPU) ردیابی می‌شوند و هیچ اثری از دستگاه (TPU/GPU) جمع‌آوری نمی‌شود.
   • TRACE_ONLY_XLA : این به این معنی است که فقط عملیات سطح XLA روی دستگاه ردیابی می‌شوند.
   • TRACE_COMPUTE : این تابع عملیات محاسباتی روی دستگاه را ردیابی می‌کند.
   • TRACE_COMPUTE_AND_SYNC : این تابع هم عملیات محاسباتی و هم رویدادهای همگام‌سازی را روی دستگاه ردیابی می‌کند.

   اگر "tpu_trace_mode" ارائه نشده باشد، trace_mode به طور پیش‌فرض روی TRACE_ONLY_XLA تنظیم می‌شود.

2. tpu_num_sparse_cores_to_trace : تعداد هسته‌های پراکنده برای ردیابی روی TPU را مشخص می‌کند.

3. tpu_num_sparse_core_tiles_to_trace : تعداد کاشی‌های درون هر هسته پراکنده برای ردیابی روی TPU را مشخص می‌کند.

4. tpu_num_chips_to_profile_per_task : تعداد تراشه‌های TPU برای پروفایل‌بندی در هر وظیفه را مشخص می‌کند.

گزینه‌های پردازنده گرافیکی

گزینه‌های زیر برای پروفایلینگ GPU در دسترس هستند:

• gpu_max_callback_api_events : حداکثر تعداد رویدادهای جمع‌آوری‌شده توسط API فراخوانی CUPTI را تنظیم می‌کند. مقدار پیش‌فرض آن 2*1024*1024 است.
• gpu_max_activity_api_events : حداکثر تعداد رویدادهای جمع‌آوری‌شده توسط API فعالیت CUPTI را تنظیم می‌کند. مقدار پیش‌فرض آن 2*1024*1024 است.
• gpu_max_annotation_strings : حداکثر تعداد رشته‌های حاشیه‌نویسی قابل جمع‌آوری را تنظیم می‌کند. مقدار پیش‌فرض آن 1024*1024 است.
• gpu_enable_nvtx_tracking : ردیابی NVTX را در CUPTI فعال می‌کند. مقدار پیش‌فرض آن False .
• gpu_enable_cupti_activity_graph_trace : ردیابی نمودار فعالیت CUPTI را برای نمودارهای CUDA فعال می‌کند. مقدار پیش‌فرض False است.
• gpu_pm_sample_counters : رشته‌ای از معیارهای نظارت بر عملکرد GPU که با کاما از هم جدا شده‌اند و با استفاده از ویژگی نمونه‌گیری PM در CUPTI جمع‌آوری می‌شوند (مثلاً "sm__cycles_active.avg.pct_of_peak_sustained_elapsed" ). نمونه‌گیری PM به طور پیش‌فرض غیرفعال است. برای معیارهای موجود، به مستندات CUPTI انویدیا مراجعه کنید.
• gpu_pm_sample_interval_us : فاصله نمونه‌برداری را بر حسب میکروثانیه برای نمونه‌برداری CUPTI PM تنظیم می‌کند. مقدار پیش‌فرض 500 است.
• gpu_pm_sample_buffer_size_per_gpu_mb : اندازه بافر حافظه سیستم را برای هر دستگاه بر حسب مگابایت برای نمونه‌برداری CUPTI PM تنظیم می‌کند. مقدار پیش‌فرض ۶۴ مگابایت است. حداکثر مقدار پشتیبانی‌شده ۴ گیگابایت است.
• gpu_num_chips_to_profile_per_task : تعداد دستگاه‌های GPU را برای پروفایل‌بندی در هر وظیفه مشخص می‌کند. اگر مشخص نشده باشد، روی ۰ تنظیم شود یا روی مقدار نامعتبر تنظیم شود، تمام GPUهای موجود پروفایل‌بندی می‌شوند. این می‌تواند برای کاهش اندازه جمع‌آوری ردیابی استفاده شود.
• gpu_dump_graph_node_mapping : در صورت فعال بودن، اطلاعات نگاشت گره گراف CUDA را در مسیر ردیابی ذخیره می‌کند. مقدار پیش‌فرض False .

برای مثال:

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

اگر کلیدها یا مقادیر گزینه‌ی ناشناخته‌ای یافت شوند، InvalidArgumentError را برمی‌گرداند.