תודה שהתכווננת ל-Google I/O. צפה בכל ההפעלות לפי דרישה צפה לפי דרישה

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

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

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

לממשק ה-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]))

עיין ב- Colab של המדריך לקבלת דוגמה מפורטת יותר לשימוש, וסרטון הדרכה על jit_compile=True usage.

שימוש עם 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 ללא כל שינוי היא לאפשר clustering אוטומטי , אשר מוצא באופן אוטומטי אשכולות (תתי גרפים מחוברים) בתוך פונקציות 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

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

לאחר ביצוע ההשלכה, תוכל למצוא את הקבצים הבאים ב- /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"

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

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

לקריאה נוספת

XLA Frontends

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

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

שיחות

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

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