דף זה תורגם על ידי Cloud Translation API.
Switch to English

XLA: מהדר אופטימיזציה ללמידה של מכונות

XLA (אלגברה לינארית מואצת) הוא מהדר ספציפי לתחום לאלגברה לינארית שיכול להאיץ דגמי TensorFlow ללא כל שינוי בקוד המקור.

התוצאות הן שיפורים בשימוש במהירות ובזיכרון: מרבית המוצגים הפנימיים פועלים במהירות ~ 1.15x לאחר הפעלת XLA. מערך הנתונים שלמטה מוערך על גבי NVidia V100 GPU יחיד:

מבוא

כאשר מופעלת תוכנית 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

האשכול האוטומטי מותאם כרגע לעומסי עבודה של GPU, אך ניתן לאפשר אותו גם במעבד על ידי שימוש נוסף בדגל - --tf_xla_cpu_global_jit : 

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

לקבלת דוגמא לשימוש מפורט ראה במדריך ההדרכה לאשכול אוטומטי .

אוסף מפורש עם פונקצית tf.

אשכול אוטומטי הוא כלי נהדר להפיכת הדגם למהיר יותר ללא שינויים בקוד, אך יתכן שקשה להבין אילו שינויים בוצעו.

ממשק API של קומפילציה מפורשת מציע שליטה עדינה יותר לבחירת הפונקציות שצריך לערוך. לדוגמה, הפונקציה הבאה של 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))

ממשק ה- API experimental_compile מכיל סמנטיקה חובה : כל הפונקציה מורכבת עם XLA, או שגיאה. חריגה של 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]))

ראה דוגמת שימוש במדריך לקבלת דוגמא מפורטת יותר לשימוש.

אוסף 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 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 (לעיל).

במידת האפשר, נסה לבודד באג replay_computation XLA יחידה באמצעות replay_computation והפעל אותו באופן איטרטיבי בתוכניות שנוצרו.

חזיתות XLA

מלבד TensorFlow, ניתן להפיק תוכניות XLA על ידי:

  • JAX : טרנספורמציות מחוברות של תוכניות Python + NumPy
  • ג'וליה : שפת הג'וליה למחשוב מדעי
  • PyTorch : מסגרת PyTorch

לקריאה נוספת