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

XLA: מיטוב מהדר למידת מכונה

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

התוצאות הן שיפורים במהירות ושימוש בזיכרון: למשל בהגשת BERT MLPerf באמצעות 8 וולטה V100 GPUs המשתמשים ב- XLA השיגה שיפור ביצועים ~ 7x ושיפור גודל אצווה פי 5:

מבוא

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

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

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

אשכולות אוטומטיים

דרך פשוטה להתחיל להשתמש ב- 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

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

אוסף AOT (Ahead-of-time) עבור מעבד עם 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