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

قراءة متعمقة

واجهات XLA

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

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

محادثات

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

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