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

OpenXLA הוא מהדר ספציפי לדומיין של אלגברה לינארית, שיכול לזרז מודלים של TensorFlow ללא שינויים בקוד המקור.

מבוא

כשתוכנית 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)

באמצעות Explicit compilation 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))

ל-API של jit_compile יש סמנטיקה של must-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. כדי להפעיל את הקיבוץ האוטומטי ב-GPU, מגדירים את משתנה הסביבה TF_XLA_FLAGS:

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

בשלב הזה, האשכולות האוטומטיים מותאמים לעומסי עבודה של GPU, אבל אפשר להפעיל אותם גם ב-CPU באמצעות הדגל --tf_xla_cpu_global_jit:

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

דוגמה מפורטת לשימוש זמינה במדריך ב-Colab בנושא אשכולות אוטומטיים.

הידור AOT (מראש) ל-CPU באמצעות 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

אחרי ביצוע ה-dump, הקבצים הבאים יופיעו ב-/tmp/generated:

  • module_XXXX.*_optimizations.txt תוכניות XLA שנוצרו, אחת לכל אשכול שעבר הידור. צירוף הקבצים האלה כששולחים דוחות על באגים ב-XLA עוזר מאוד.

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

  • module_XXXX.ptx קובצי PTX שנוצרו.

אפשר גם ליצור גרסת dump של התרשים שמציג את הטמעת האשכולות של 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 אחת באמצעות run_hlo_module ולהריץ אותה באופן איטרטיבי בתוכניות שנוצרו.

קריאה נוספת

ממשקי קצה של XLA

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

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

הרצאות

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

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