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 في الذاكرة ؛ وبدلاً من ذلك ، تقوم "بدفق" نتائج هذه الحسابات الوسيطة مباشرةً إلى مستخدميها مع الاحتفاظ بها بالكامل في سجلات وحدة معالجة الرسومات. 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 على دلالات يجب تجميعها : إما أن يتم ترجمة الوظيفة بالكامل باستخدام 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

تم تحسين التجميع التلقائي حاليًا لأحمال عمل GPU ، ولكن يمكن أيضًا تمكينه على وحدة المعالجة المركزية عن طريق استخدام العلامة --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"

تقارير الشوائب القابلة لإعادة الإنتاج

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

محادثات

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

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