ترجمت واجهة Cloud Translation API‏ هذه الصفحة.
Switch to English

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

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

النتائج هي تحسينات في السرعة واستخدام الذاكرة: تعمل معظم المعايير الداخلية بشكل أسرع ~ 1.15 مرة بعد تمكين XLA. يتم تقييم مجموعة البيانات أدناه على وحدة معالجة رسومات NVidia V100 واحدة:

المقدمة

عند تشغيل برنامج 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. الانصهار هو أهم تحسين لـ XLA. عادةً ما يكون عرض النطاق الترددي للذاكرة هو المورد الأكثر ندرة على مسرعات الأجهزة ، لذا تعد إزالة عمليات الذاكرة إحدى أفضل الطرق لتحسين الأداء.

تمكين XLA لنماذج TensorFlow

التكتل التلقائي

تتمثل أبسط طريقة لبدء استخدام XLA في نماذج TensorFlow في تمكين التجميع التلقائي ، الذي يعثر تلقائيًا على مجموعات (مخططات فرعية متصلة) ضمن الرسم البياني TensorFlow الذي يمكن تجميعه وتنفيذه باستخدام XLA. يمكن تمكين التكتل التلقائي على GPU من خلال تعيين متغير البيئة 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
 

للحصول على مثال مفصل عن الاستخدام ، انظر إلى البرنامج التعليمي للتجمع التلقائي .

تجميع صريح مع وظيفة tf.

يعد التجميع التلقائي أداة رائعة لجعل النموذج أسرع بدون أي تغييرات على التعليمات البرمجية ، ولكن قد يكون من الصعب فهم التغييرات التي تم إجراؤها.

تقدم واجهة برمجة التطبيقات صريحة تحكمًا أكثر دقة لاختيار الوظائف التي يجب ترجمتها. على سبيل المثال ، يتم تجميع وظيفة TensorFlow التالية التي تجري تدريب MNIST باستخدام XLA:

 @tf.function(experimental_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))
 

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

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

يمكن أن تختلف الأشكال عبر الجري بالرغم من ذلك:

 @tf.function(experimental_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 التعليمي للحصول على مثال استخدام أكثر تفصيلاً.

تجميع tfcompile (قبل الوقت) لوحدة المعالجة المركزية مع 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 التي تم module_XXXX.*_optimizations.txt ، واحدة لكل مجموعة مجمعة. يعد إرفاقها عند إرسال تقارير أخطاء XLA أمرًا مفيدًا للغاية!

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

  • module_XXXX.ptx ملفات PTX التي تم module_XXXX.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 : التحولات المركبة لبرامج Python + NumPy
  • جوليا : لغة جوليا للحوسبة العلمية
  • PyTorch : إطار PyTorch

قراءة متعمقة