تفريغ عمليات حسابية في HLO

تفريغ HLO هو تمثيل نصي لوحدات HLO في مراحل مختلفة من عملية الحساب. وهو مفيد لتصحيح الأخطاء، وغالبًا ما تحتاج إلى تضمينه في تقارير الأخطاء. يكون هذا الملف عادةً ملفًا نصيًا يمكن لشخص عادي قراءته، وهو يتضمّن تعليمات HLO وخصائصها. في بعض الأحيان، يتم تفريغ وحدات HLO النمطية على النحو التالي:

• HloProto: ملفات بروتوكول buffer، وهي تنسيق أكثر تنظيمًا وقابلاً للقراءة آليًا.
• HloSnapshot: وحدة HLO بالإضافة إلى مدخلاتها عند إعادة تشغيل العمليات الحسابية الطويلة الأمد، تحتاج أحيانًا إلى الإدخالات الفعلية التي تم تقديمها إلى عملية حسابية معيّنة بدلاً من البيانات العشوائية.

يمكنك استخدام علامات XLA لتحديد عمليات التفريغ والحصول عليها. في معظم الحالات، يمكنك ضبطها باستخدام متغيّر بيئة. توفّر JAX أيضًا طريقة برمجية لطباعة تفريغ HLO.

التنفيذ المحلي

استخدام متغيرات البيئة

يمكنك ضبط متغيّر البيئة XLA_FLAGS باستخدام العلامات اللازمة للحصول على عمليات تفريغ. يمكن استخدام هذه الميزة مع 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 لعرض أجزاء من الرسم البياني.

تفريغ البطاقات الوسيطة المحدّدة

بالإضافة إلى HLOs العادية التي تم تحسينها مسبقًا أو نهائيًا، يمكنك أيضًا تفريغ حالة HLOs بعد عملية تجميع معيّنة.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

سيتم تفريغ وحدات HLO لعمليات النقل التي تتطابق أسماؤها مع التعبير العادي (regex). على سبيل المثال، يمكنك مراقبة نتائج التحسينات على مستوى الأداء الناتجة عن عمليات التمرير المتعلّقة بتقسيم 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 آليًا باستخدام واجهتَي برمجة التطبيقات lower وcompile في JAX.

يمكنك جلب 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)

Dump jaxprs

jaxprs هي التمثيل الوسيط لعمليات تتبُّع البرنامج في JAX. لتفريغ هذا، اضبط متغيّرات البيئة:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

يمكنك الاطّلاع على مزيد من المعلومات في مستندات JAX حول تصدير العمليات الحسابية التي تم إعدادها وتسلسلها: تصحيح الأخطاء.

Google Colab

متغيرات البيئة

في الخلية الأولى التي يتم تنفيذها في دفتر الملاحظات (لأنّ متغيرات البيئة وعلامات سطر الأوامر تتم معالجتها مرة واحدة فقط عادةً، مثلاً، عند استيراد الوحدة أو عند بدء تشغيل الخلفية XLA)، أضِف XLA_FLAGS الموضّحة أعلاه مع os.environ، مثلاً:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

سيؤدي ذلك إلى إخراج الحساب إلى DIRECTORY_PATH، مثل /tmp. في Colab، انتقِل إلى متصفّح "الملفات" في الشريط الجانبي الأيمن لعرض هذا الدليل والوصول إليه.

يمكنك استخدام جميع العلامات المذكورة في قسم "التنفيذ المحلي".

خيارات خاصة بـ 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)

تجاهل جميع العمليات الحسابية أو بعضها

إذا أردت الاطّلاع على كل شيء في عملية تفريغ، بما في ذلك جميع عمليات التجميع الصغيرة، اضبط متغير بيئة JAX على النحو التالي:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

فسيفساء

‫Mosaic هو برنامج مترجم للواجهة الخلفية لوحدة معالجة Tensor من Pallas، والواجهة الخلفية التجريبية لوحدة معالجة الرسومات من Pallas. لتفريغ عملية حساب فسيفساء، اضبط العلامة التالية:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

أو يمكنك ضبط وسيطات تهيئة TPU كمتغير بيئي:

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

يمكنك الاطّلاع على مستندات JAX حول Pallas وMosaic لمزيد من المعلومات.

المزيد من أعمال HLO Dumps

العثور على الحساب المناسب

عادةً، يتم تجاهل العديد من العمليات الحسابية. يتم تسمية الملفات التي تم تفريغها بشكل صريح باستخدام "اسم العملية الحسابية" في JAX أو TensorFlow أو PyTorch/XLA الذي يتم عرضه في السجلات، ما يسهّل تحديد ملفات HLO ذات الصلة. على سبيل المثال:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

وإلا، يمكنك استخدام ripgrep لتحديد الوحدة التي تتضمّن رموزًا أو عمليات حسابية معيّنة بسرعة.

ملاحظة: أدرِج في تقارير الأخطاء الملفات الثلاث التي تم تفريغها قبل/بعد/أثناء عملية تخصيص المخزن المؤقت والتي تهمك.

HLO Conversion

أداة تُسمى hlo-opt يمكنها الترجمة بين تنسيقات HLOProto والنص. ويكون ذلك مفيدًا في الحالات التي يتوفّر فيها تنسيق واحد، ولكنك تحتاج إلى التنسيق الآخر لتصحيح الأخطاء.

يمكنك الاطّلاع على كيفية استخدامها: مستندات أدوات XLA: hlo-opt.

إعادة التشغيل

يمكنك تشغيل (إعادة تشغيل) العمليات الحسابية التي تم تفريغها على خلفية XLA محدّدة باستخدام بيانات وهمية أو لقطات إدخال. هذه طريقة ملائمة لإعادة إنتاج المشاكل وتكرارها وتصحيحها في XLA.

تستخدم الأوامر التالية بيانات وهمية. إذا كنت قد حفظت "لقطات" HLO، يمكنك تقديمها بدلاً من ذلك، وسيتم استخدام البيانات من اللقطة. لمواصلة استخدام بيانات وهمية أثناء تشغيل اللقطة، مرِّر العلامة --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