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
ולהפעיל אותו באופן איטרטיבי בתוכניות שנוצרו.
קריאה נוספת
- מסמכי תיעוד של OpenXLA מסמכי OpenXLA
- בעיות מוכרות רשימת בעיות ידועות ב-XA+TF
- XLA – TensorFlow, הידור,: לקריאת הבלוג של Google Developers
- עיינו במקור של XLA ב-GitHub!
ממשקי XLA
מלבד TensorFlow, ניתן ליצור תוכניות XLA בדרכים הבאות:
- JAX: טרנספורמציות קומפוזביליות של תוכניות Python+NumPy
- ג'וליה: השפה של ג'וליה למחשוב מדעי
- PyTorch: PyTorch framework
- Nx: ספריית מחשוב נומרית לשפת התכנות Elixir