احفظ التاريخ! يعود مؤتمر Google I / O من 18 إلى 20 مايو. سجل الآن
ترجمت واجهة Cloud Translation API‏ هذه الصفحة.
Switch to English

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

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

النتائج عبارة عن تحسينات في السرعة واستخدام الذاكرة: على سبيل المثال ، في تقديم BERT MLPerf باستخدام 8 وحدات معالجة رسومات Volta 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 واجهة برمجة التطبيقات jit_compile على دلالات يجب تجميعها : إما أن يتم ترجمة الوظيفة بالكامل باستخدام 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 للحصول على مثال استخدام أكثر تفصيلاً.

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

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

يمكنك أيضًا استخدام أداة 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 واحد باستخدام replay_computation متكرر على البرامج التي تم إنشاؤها.

قراءة متعمقة

واجهات XLA

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

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