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
תוכניות XLAmodule_XXXX.*_optimizations.txt
, אחת לכל אשכול מהודר. צירוף אלה בעת הגשת דוחות באגים של XLA מועיל ביותר!module_XXXX.ir-*.ll
קבצים שנוצרו בייצוג ביניים LLVM , עם פנימיות NVPTX .module_XXXX.ptx
קבצי PTXmodule_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
- אדריכלות XLA : סקירה כללית של ארכיטקטורת XLA
- XLA - TensorFlow, הידור : קרא בבלוג Google Developers
- בדוק את מקור ה- XLA ב- Github!
חזית XLA
מלבד TensorFlow, ניתן ליצור תוכניות XLA על ידי: