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)

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

@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]))

אפשר לראות דוגמה מפורטת יותר לשימוש במערך שיעור המשותף, ובסרטון הדרכה על השימוש ב-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)

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

דרך פשוטה להתחיל להשתמש ב-LA במודלים של 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 לקוד הפעלה (למעבדים (CPU) מסוג 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, אחת לכל אשכול שעבר הידור. מומלץ מאוד לצרף את הפרטים כששולחים דוחות על באגים מסוג XLA.

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

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

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

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

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

הרבה יותר קל ליצור דוח על באג אם הוא כולל קובצי Dump של תוכניות 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
  • ג'וליה: השפה של ג'וליה למחשוב מדעי
  • PyTorch: PyTorch framework
  • Nx: ספריית מחשוב נומרית לשפת התכנות Elixir

הרצאות

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

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