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

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

مقدمة

عند تشغيل برنامج TensorFlow، يتم تنفيذ جميع العمليات بشكل فردي من خلال تنفيذ TensorFlow. تحتوي كل عملية TensorFlow على تنفيذ نواة وحدة معالجة رسومات مجمعة مسبقًا والتي يرسِل إليها المسؤول.

توفر 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)

توفّر واجهة برمجة التطبيقات للتجميع الصريح إمكانية تحكم دقيقة لاختيار الدوال التي يجب تجميعها. على سبيل المثال، يتم تجميع دالة 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 API على دلالات يجب تجميعها، أي أنّه يتم تجميع الدالة بأكملها باستخدام 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 التعليمي للاطّلاع على مثال أكثر تفصيلاً على الاستخدام، وفيديو تعليمي حول استخدام 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

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

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

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

مجموعة مجمّعة من AOT (مسبقة) لوحدة المعالجة المركزية (CPU) من خلال 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"

تقارير الأخطاء القابلة للتكرار

يسهل إعادة إنتاج تقرير الأخطاء إذا كان يتضمن تفريغ لبرامج 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
  • جوليا: لغة جوليا للحوسبة العلمية
  • PyTorch: إطار عمل PyTorch
  • Nx: مكتبة الحوسبة الرقمية للغة برمجة Elxir

محادثات

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

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