تسجيل بيانات أداء عمليات حساب JAX باستخدام XProf

‫XProf هي طريقة رائعة للحصول على بيانات تتبُّع الأداء وملفات التعريف الخاصة ببرنامجك وعرضها بشكل مرئي، بما في ذلك النشاط على وحدة معالجة الرسومات (GPU) ووحدة معالجة Tensor (TPU). يبدو الناتج النهائي على النحو التالي:

مثال على XProf

الالتقاط الآلي

يمكنك تزويد الرمز الخاص بك بأدوات لتسجيل بيانات التتبُّع الخاصة بأداة Profiler لرمز 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.

يمكنك تشغيل واجهة مستخدم أداة Profiler مباشرةً باستخدام أمر 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)

انتقِل إلى عنوان URL المقدَّم (مثل http://localhost:8791/) في المتصفّح للاطّلاع على الملف التجاري.

تظهر عمليات التتبُّع المتاحة في القائمة المنسدلة "الجلسات" على اليمين. اختَر الجلسة التي تهمّك، ثم انقر على "Trace Viewer" (أداة عرض عمليات التتبُّع) من القائمة المنسدلة "الأدوات" (Tools). من المفترض أن يظهر لك الآن مخطط زمني للتنفيذ. يمكنك استخدام مفاتيح WASD للتنقّل في التتبُّع، والنقر أو السحب لاختيار الأحداث للحصول على مزيد من التفاصيل. يمكنك الاطّلاع على مستندات "أداة عرض عمليات التتبُّع" للحصول على مزيد من التفاصيل حول استخدام أداة عرض عمليات التتبُّع.

تسجيل اللقطات يدويًا من خلال XProf

في ما يلي تعليمات لتسجيل عملية تتبُّع لمدة N ثانية يتم تشغيلها يدويًا من برنامج قيد التشغيل.

  1. ابدأ خادم XProf:

    xprof --logdir /tmp/profile-data/

    من المفترض أن تتمكّن من تحميل XProf على <http://localhost:8791/>. يمكنك تحديد منفذ مختلف باستخدام العلامة --port.

  2. في برنامج Python أو العملية التي تريد إنشاء ملف تعريف لها، أضِف ما يلي في مكان ما بالقرب من البداية:

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

    يبدأ هذا الخادم الذي يتصل به XProf. يجب أن يكون خادم أداة تحليل الأداء قيد التشغيل قبل الانتقال إلى الخطوة التالية. عند الانتهاء من استخدام الخادم، يمكنك طلب jax.profiler.stop_server() لإيقافه.

    إذا أردت إنشاء ملف تعريف لجزء من برنامج طويل الأمد (مثل حلقة تدريب طويلة)، يمكنك وضع هذا الرمز في بداية البرنامج وبدء تشغيله كالمعتاد. إذا أردت إنشاء ملف تعريف لبرنامج قصير (مثل اختبار أداء صغير)، يمكنك بدء تشغيل خادم أداة إنشاء الملفات الشخصية في واجهة IPython، وتشغيل البرنامج القصير باستخدام %run بعد بدء عملية التسجيل في الخطوة التالية. هناك خيار آخر وهو بدء تشغيل خادم أداة Profiler في بداية البرنامج واستخدام time.sleep() لمنحك وقتًا كافيًا لبدء عملية التسجيل.

  3. افتح <http://localhost:8791/>، وانقر على الزر "تسجيل الملف الشخصي" (CAPTURE PROFILE) في أعلى يمين الشاشة. أدخِل "localhost:9999" كعنوان URL لخدمة الملفات الشخصية (هذا هو عنوان خادم أداة إنشاء الملفات الشخصية الذي بدأته في الخطوة السابقة). أدخِل عدد المللي ثانية التي تريد إنشاء ملف تعريف لها، ثم انقر على "تسجيل".

  4. إذا لم يكن الرمز الذي تريد إنشاء ملف تعريف له قيد التشغيل (على سبيل المثال، إذا بدأت خادم أداة إنشاء الملفات الشخصية في shell Python)، شغِّله أثناء تشغيل عملية الالتقاط.

  5. بعد انتهاء عملية التسجيل، من المفترض أن يتم إعادة تحميل XProf تلقائيًا. (لا يتم ربط جميع ميزات إنشاء الملفات الشخصية في XProf بـ JAX، لذا قد يبدو في البداية أنّه لم يتم تسجيل أي بيانات). على يمين الصفحة، ضِمن "الأدوات"، انقر على "أداة عرض التتبُّع".

من المفترض أن يظهر لك الآن مخطط زمني للتنفيذ. يمكنك استخدام مفاتيح WASD للتنقّل في التتبُّع، والنقر أو السحب لاختيار الأحداث والاطّلاع على مزيد من التفاصيل في أسفل الصفحة. يمكنك الاطّلاع على مستندات &quot;أداة عرض عمليات التتبُّع&quot; لمزيد من التفاصيل حول استخدام أداة عرض عمليات التتبُّع.

‫XProf وTensorboard

‫XProf هي الأداة الأساسية التي تتيح وظائف إنشاء الملفات الشخصية وتسجيل عمليات التتبُّع في Tensorboard. طالما أنّ xprof مثبّت، ستظهر علامة التبويب "الملف الشخصي" ضمن 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، والتي تتيح التحكّم الدقيق في سلوك أداة تحليل الأداء. يجب أن تكون هذه المَعلمة مثيلاً من jax.profiler.ProfileOptions.

على سبيل المثال، لإيقاف جميع عمليات تتبُّع Python والمضيف:

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: لإيقاف تتبُّع المضيف (وحدة المعالجة المركزية) بالكامل
    • ‫1: يتيح تتبُّع أحداث TraceMe التي نفّذها المستخدم فقط.
    • ‫2: يتضمّن عمليات التتبُّع من المستوى 1 بالإضافة إلى تفاصيل تنفيذ البرنامج على مستوى عالٍ، مثل عمليات XLA المكلفة (الإعداد التلقائي).
    • ‫3: يتضمّن عمليات التتبُّع من المستوى 2 بالإضافة إلى تفاصيل أكثر تفصيلاً حول تنفيذ البرنامج على مستوى منخفض، مثل عمليات XLA غير المكلفة.

  2. ‫device_tracer_level: تتحكّم هذه السياسة في ما إذا كان تتبُّع الجهاز مفعَّلاً.

    القيم المسموح بها:

    • ‫0: لإيقاف تتبُّع الجهاز
    • 1: يتيح تتبُّع الجهاز (الإعداد التلقائي).

  3. ‫python_tracer_level: تتحكّم هذه السمة في ما إذا كان تتبُّع Python مفعَّلاً.

    القيم المسموح بها:

    • ‫0: لإيقاف تتبُّع استدعاء وظائف Python (الإعداد التلقائي)
    • ‫1: تفعيل تتبُّع Python

خيارات الضبط المتقدمة

خيارات وحدة معالجة الموتّرات

  1. ‫tpu_trace_mode: تحدّد هذه السمة وضع تتبُّع وحدة معالجة الموتّرات (TPU).

    القيم المسموح بها:

    • TRACE_ONLY_HOST: يعني ذلك أنّه يتم تتبُّع الأنشطة من جهة المضيف (وحدة المعالجة المركزية) فقط، ولا يتم جمع أي عمليات تتبُّع من الجهاز (وحدة معالجة Tensor أو وحدة معالجة الرسومات).
    • ‫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_max_callback_api_events: تضبط هذه السمة الحد الأقصى لعدد الأحداث التي يتم جمعها من خلال واجهة برمجة التطبيقات CUPTI. القيمة التلقائية هي 2*1024*1024.
  • gpu_max_activity_api_events: يضبط الحد الأقصى لعدد الأحداث التي تجمعها واجهة برمجة التطبيقات الخاصة بنشاط 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 من NVIDIA.
  • gpu_pm_sample_interval_us: يضبط فاصل أخذ العينات بالميكرو ثانية لأخذ عينات CUPTI PM. القيمة التلقائية هي 500.
  • ‫gpu_pm_sample_buffer_size_per_gpu_mb: يضبط حجم مخزن مؤقت لذاكرة النظام لكل جهاز بالميغابايت لأخذ عينات من أداء CUPTI. القيمة التلقائية هي 64 ميغابايت. الحد الأقصى للقيمة المسموح بها هو 4 غيغابايت.
  • gpu_num_chips_to_profile_per_task: تحدّد هذه السمة عدد أجهزة وحدة معالجة الرسومات التي سيتم إنشاء ملف تعريف لها لكل مهمة. في حال عدم تحديد هذه السمة أو ضبطها على 0 أو على قيمة غير صالحة، سيتم إنشاء ملفات تعريف لجميع وحدات معالجة الرسومات المتاحة. ويمكن استخدام ذلك لتقليل حجم مجموعة عمليات التتبُّع.
  • 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 في حال العثور على أي مفاتيح أو قيم خيارات غير معروفة.