XLA: تحسين برنامج التحويل للتعلم الآلي

OpenXLA هو مُجمِّع خاص بنطاق معيّن لجبر المعادلات الخطية يمكنه تسريع نماذج TensorFlow بدون أي تغييرات في رمز المصدر.

مقدمة

عند تشغيل برنامج TensorFlow، يتم تنفيذ جميع العمليات بشكل فردي بواسطة أداة تنفيذ TensorFlow. تحتوي كل عملية في TensorFlow على تنفيذ ملف برمجي لوحدة معالجة الرسومات (GPU) تمت تجميعه مسبقًا ويوجّهه المشغّل.

توفّر XLA وضعًا بديلاً لتشغيل النماذج: فهي تُجمِّع مسار TensorFlow إلى تسلسل من نوى الحساب التي يتم إنشاؤها خصيصًا لل نموذج المحدّد. وبما أنّ هذه النوى فريدة للنموذج، يمكنها الاستفادة من المعلومات الخاصة بالنموذج للتحسين. على سبيل المثال، لنلقِ نظرة على أحد عمليات التحسين التي تُجريها XLA في سياق عملية حسابية بسيطة في TensorFlow:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

عند التشغيل بدون XLA، يطلق الرسم البياني ثلاث نوى: نواة للضرب، وناة للجمع وناة للقسمة. ومع ذلك، يمكن أن تحسِّن تقنية XLA الرسم البياني لكي تحسب النتيجة في عملية تشغيل نواة واحدة. ويتم ذلك من خلال "دمج" عمليات الإضافة والضرب والطرح في نواة وحدة معالجة الرسومات واحدة. بالإضافة إلى ذلك، لا تُسجِّل هذه العملية المدمجة القيم الوسيطة التي تنتجها y*z وx+y*z في الذاكرة، بل "تُبثّ" نتائج عمليات الحساب الوسيطة هذه مباشرةً إلى المستخدمين مع إبقائها بالكامل في سجلات وحدة معالجة الرسومات. إنّ Fusion هو أهم ميزة تحسين في XLA. عادةً ما يكون معدل نقل البيانات في الذاكرة هو المورد الأكثر ندرة في مسرعات الأجهزة، لذا فإنّ إزالة عمليات الذاكرة هي إحدى أفضل الطرق لتحسين الأداء.

تفعيل XLA لنماذج TensorFlow

التجميع الصريح باستخدام tf.function(jit_compile=True)

توفّر واجهة برمجة التطبيقات Explicit compilation API عنصر تحكّم دقيقًا لاختيار الدوالّ التي يجب تجميعها. على سبيل المثال، يتم تجميع دالة TensorFlow التالية التي تُجري تدريب MNIST باستخدام XLA:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

تحتوي واجهة برمجة التطبيقات jit_compile على دلالات يجب تجميعها: إما تجميع الدالة بالكامل باستخدام XLA أو طرح استثناء errors.InvalidArgumentError. لا يمكن لواجهة XLA حاليًا تجميع الدوالّ التي لا يمكناستنتاج أبعادها، أي إذا لم يكن من الممكن استنتاج أبعاد كل المتسلسلات بدون تشغيل العملية الحسابية بأكملها. على سبيل المثال، لن يتم تجميع الدالة التالية:

@tf.function
def not_compilable(x):
  return tf.unique(x)

يمكن أن تختلف الأشكال في عمليات التنفيذ:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

اطّلِع على الدليل التعليمي في Colab للحصول على مثال detailed على الاستخدام، وعلى فيديو تعليمي حول استخدام jit_compile=True.

الاستخدام مع Keras

بالنسبة إلى نماذج Keras، يمكن ضبط jit_compile=True كوسيطة لمحاولة model.compile:

model.compile(optimizer="adam", jit_compile=True)

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

XLA:GPU يمكن استخدامها مع استراتيجية TF الموزّعة (MirroredStrategy أو MultiWorkerMirroredStrategy) من خلال التعليق التوضيحي على دالة الخطوة باستخدام jit_compile=True:

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

التجميع التلقائي

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

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

تم تحسين التجميع التلقائي حاليًا لأعمال وحدة معالجة الرسومات، ولكن يمكن أيضًا تفعيله على وحدة المعالجة المركزية باستخدام العلامة --tf_xla_cpu_global_jit:

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

للحصول على مثال مفصّل على الاستخدام، يُرجى الاطّلاع على الدليل التعليمي حول التجميع التلقائي في Colab.

تجميع AOT (قبل وقت التشغيل) لوحدة المعالجة المركزية باستخدام tfcompile

يمكنك أيضًا استخدام أداة tfcompile مستقلة، والتي تحوّل الرسم البياني TensorFlow إلى رمز قابل للتنفيذ (لوحدة المعالجة المركزية x86-64 فقط).

فحص البرامج المجمّعة

توفّر XLA مرافق فحص ذاتي تتيح لك فحص البرامج التي تم إنشاؤها. لتفريغ البرامج التي تم إنشاؤها، استخدِم متغيّر البيئة XLA_FLAGS:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

بعد تنفيذ عملية التفريغ، يمكنك العثور على الملفات التالية في /tmp/generated:

  • module_XXXX.*_optimizations.txt تم إنشاء برامج XLA، واحد لكل مجموعة مجمّعة. إنّ إرفاق هذه الملفات عند إرسال تقارير أخطاء XLA مفيد للغاية.

  • module_XXXX.ir-*.ll الملفات التي تم إنشاؤها في تمثيل وسيط LLVM، مع وظائف برمجية داخلية NVPTX

  • module_XXXX.ptx تم إنشاء ملفات PTX.

يمكنك أيضًا تفريغ الرسم البياني الذي يوضّح عملية تضمين مجموعات XLA داخل رسم بياني TensorFlow باستخدام:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

تقارير الأخطاء القابلة لإعادة الإنتاج

من الأسهل بكثير إعادة إنتاج تقرير الخطأ إذا كان يتضمّن نُسخًا من برامج MATLAB XLA التي تم إنشاؤها وعمليات تضمين التجميع التلقائي المستخدَمة. لإنشاء هذه النماذج لبرنامج TensorFlow الذي يتم تشغيله باستخدام التجميع التلقائي، يمكنك تشغيل:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

عند الإبلاغ عن الأخطاء، يجب إرفاق محتويات الدليل /tmp/generated (المُشار إليه أعلاه).

إذا أمكن، حاوِل عزل الخطأ في برنامج XLA واحد باستخدام run_hlo_module وتشغيله بشكل متكرّر على البرامج التي تم إنشاؤها.

مراجع إضافية

واجهات XLA

بالإضافة إلى TensorFlow، يمكن إنشاء برامج XLA من خلال:

  • JAX: عمليات تحويل قابلة للتجميع لبرامج Python+NumPy
  • Julia: لغة Julia لمعالجة المهام الحسابية العلمية
  • PyTorch: إطار عمل PyTorch
  • Nx: مكتبة الحوسبة الرقمية الخاصة بلغة برمجة Elixir

محادثات

استخدام XLA من TF باستخدام jit_compile=True

نظرة عامة على XLA