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

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

התוצאות הן שיפורי שימוש מהיר זיכרון: למשל ברט MLPerf הגשה באמצעות 8 GPUs וולטה 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)

ה-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 יש סמנטיקת חובת הידור: או הפונקציה כול נערכה עם XLA, או errors.InvalidArgumentError חריג נזרק. XLA לא יכול כרגע לקמפל פונקציות שבו הממדים אינם inferrable: כלומר, אם זה לא אפשרי להסיק את הממדים של כל tensors מבלי להפעיל את החישוב כולו. לדוגמה, הפונקציה הבאה לא תעשה קומפילציה:

@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 שימוש.

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

דרך פשוטה להתחיל להשתמש XLA במודלים TensorFlow ללא כל שינוי הוא לאפשר אשכולות אוטומטי, אשר באופן אוטומטי מוצאת אשכולות (subgraphs קשור) בתוך פונקציות TensorFlow אשר ניתן להדר והוצא להורג באמצעות XLA. Auto-באשכולות על GPU ניתן להפעיל על ידי קביעת TF_XLA_FLAGS משתנה הסביבה:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Auto-באשכולות מותאמים כרגע עבור עומסי עבודת 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 (Ahead of-זמן) הידור עבור מעבד עם 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 Generated תוכניות XLA , אחד לכל אשכול הידור. צירוף אלה בעת הגשת דוחות באגים XLA מועיל ביותר!

  • module_XXXX.ir-*.ll קבצים שנוצרו LLVM ייצוג ביניים, עם NVPTX intrinsics.

  • module_XXXX.ptx Generated PTX קבצים.

אתה יכול גם לזרוק את הגרף הממחיש את ההטמעה של אשכולות XLA בתוך גרף TensorFlow עם:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

דוחות באגים שניתנים לשחזור

דוח באג קל הרבה יותר לשחזר אם הוא כולל dumps עבור תוכניות 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 Frontends

מלבד TensorFlow, ניתן ליצור תוכניות XLA על ידי:

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

שיחות

שימוש XLA מ TF באמצעות jit_compile=True

סקירה כללית של XLA