مؤتمر Google I / O هو التفاف! تابع جلسات TensorFlow اعرض الجلسات

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

يعد XLA (الجبر الخطي المتسارع) مترجمًا خاصًا بالمجال للجبر الخطي يمكنه تسريع نماذج TensorFlow مع عدم وجود تغييرات محتملة في كود المصدر.

وكانت النتائج تحسينات في السرعة واستخدام الذاكرة: على سبيل المثال في بيرت MLPerf تقديم باستخدام 8 فولتا V100 وحدات معالجة الرسومات باستخدام XLA حققت ~ 7X تحسين الأداء و~ 5X دفعة تحسن حجم:

مقدمة

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

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

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

تشغيل بدون XLA ، يطلق الرسم البياني ثلاث نواة: واحدة للضرب ، وواحدة للإضافة وواحدة للتصغير. ومع ذلك ، يمكن لـ XLA تحسين الرسم البياني بحيث يحسب النتيجة في إطلاق نواة واحدة. يقوم بذلك عن طريق "دمج" عمليات الجمع والضرب والاختزال في نواة GPU واحدة. وعلاوة على ذلك، فإن هذه العملية تنصهر لا يكتب خارج القيم الوسيطة التي تنتجها y*z و x+y*z إلى الذاكرة. بدلاً من ذلك ، يقوم "بدفق" نتائج هذه الحسابات الوسيطة مباشرةً إلى مستخدميها مع الاحتفاظ بها بالكامل في سجلات GPU. 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 لا يمكن ترجمة حاليا وظائف حيث أبعاد ليست inferrable: هذا هو، إذا كان من غير الممكن أن نستنتج أبعاد كل التنسورات دون تشغيل الحساب بأكمله. على سبيل المثال ، لن يتم ترجمة الوظيفة التالية:

@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 الاستخدام.

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

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

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

هو الأمثل لصناعة السيارات في تجميع حاليا لأعباء العمل GPU، ولكن يمكن أيضا أن يدعم على وحدة المعالجة المركزية من قبل بالإضافة إلى ذلك باستخدام العلامة --tf_xla_cpu_global_jit :

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

للمثال استخدام مفصل، انظر لصناعة السيارات في تجمع تعليمي colab .

عثمان (قبل من الوقت) تجميع وحدة المعالجة المركزية مع tfcompile

يمكنك أيضا استخدام مستقل tfcompile أداة، والذي يحول TensorFlow الرسم البياني إلى رمز قابل للتنفيذ (وحدة المعالجة المركزية إكس86-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 intrinsics.

  • 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 واحد باستخدام replay_computation وتشغيله على برامج ولدت بشكل متكرر.

قراءة متعمقة

واجهات XLA

بصرف النظر عن TensorFlow ، يمكن إنشاء برامج XLA من خلال:

  • JAX التحولات Composable برامج بيثون + نمباي:
  • جوليا : اللغة جوليا عن الحوسبة العلمية
  • PyTorch إطار PyTorch:
  • NX : العددي الحوسبة مكتبة للغة البرمجة إكسير

محادثات

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

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