הפעלת חישוב במכונה וירטואלית של Cloud TPU באמצעות JAX

במאמר הזה מוסבר בקצרה איך עובדים עם JAX ו-Cloud TPU.

לפני שמתחילים

לפני שמריצים את הפקודות במסמך הזה, צריך ליצור חשבון, להתקין את Google Cloud CLI ולהגדיר את הפקודה gcloud. Google Cloudמידע נוסף זמין במאמר בנושא הגדרת סביבת Cloud TPU.

התפקידים הנדרשים

כדי לקבל את ההרשאות שדרושות ליצירת TPU ולהתחבר אליו באמצעות SSH, צריך לבקש מהאדמין להקצות לכם בפרויקט את תפקידי ה-IAM הבאים:

להסבר על מתן תפקידים, ראו איך מנהלים את הגישה ברמת הפרויקט, התיקייה והארגון.

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

יצירת מכונה וירטואלית של Cloud TPU באמצעות gcloud

  1. כדי להקל על השימוש בפקודות, מגדירים כמה משתני סביבה.

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east5-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

    תיאורים של משתני סביבה

    • PROJECT_ID: מזהה הפרויקט ב- Google Cloud . משתמשים בפרויקט קיים או יוצרים פרויקט חדש.
    • ‫TPU_NAME: השם של ה-TPU.
    • ‫ZONE: האזור שבו יוצרים את מכונת ה-TPU. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותחומים של TPU.
    • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
    • ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.

  2. מריצים את הפקודה הבאה מ-Cloud Shell או ממסוף המחשב שבו מותקן Google Cloud CLI כדי ליצור מכונת TPU.

    $ gcloud compute tpus tpu-vm create $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --accelerator-type=$ACCELERATOR_TYPE \
    --version=$RUNTIME_VERSION

התחברות למכונה וירטואלית של Cloud TPU

מתחברים למכונת ה-TPU הווירטואלית באמצעות SSH באמצעות הפקודה הבאה:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

אם לא הצלחתם להתחבר למכונת TPU וירטואלית באמצעות SSH, יכול להיות שהמכונה הווירטואלית של TPU לא כוללת כתובת IP חיצונית. כדי לגשת למכונת TPU וירטואלית ללא כתובת IP חיצונית, פועלים לפי ההוראות במאמר התחברות למכונת TPU וירטואלית ללא כתובת IP ציבורית.

התקנת JAX במכונת ה-VM של Cloud TPU

(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

בדיקת מערכת

מוודאים של-JAX יש גישה ל-TPU והיא יכולה להריץ פעולות בסיסיות:

  1. מפעילים את מתורגמן Python 3:

    (vm)$ python3
     
    >>> import jax

  2. הצגת מספר ליבות ה-TPU הזמינות:

    >>> jax.device_count()

מוצג מספר ליבות ה-TPU. מספר ליבות ה-TPU שמוצג תלוי בגרסת ה-TPU שבה אתם משתמשים. מידע נוסף זמין במאמר בנושא גרסאות TPU.

ביצוע חישוב

>>> jax.numpy.add(1, 1)

התוצאה של פעולת החיבור ב-numpy מוצגת:

פלט מהפקודה: 

Array(2, dtype=int32, weak_type=True)

יציאה ממתורגמן Python

>>> exit()

הרצת קוד JAX במכונה וירטואלית של TPU

עכשיו אפשר להריץ כל קוד JAX שרוצים. הדוגמאות של Flax הן מקום מצוין להתחיל בו כדי להריץ מודלים רגילים של למידת מכונה ב-JAX. לדוגמה, כדי לאמן רשת בסיסית של MNIST:

  1. מתקינים את יחסי התלות של דוגמאות Flax:

    (vm)$ pip install --upgrade clu
(vm)$ pip install tensorflow
(vm)$ pip install tensorflow_datasets
(vm)$ pip install tensorboard

  2. מתקינים את Flax:

    (vm)$ git clone https://github.com/google/flax.git
(vm)$ pip install --user flax

  3. מריצים את סקריפט האימון של Flax MNIST:

    (vm)$ cd flax/examples/mnist
(vm)$ python3 main.py --workdir=/tmp/mnist \
    --config=configs/default.py \
    --config.learning_rate=0.05 \
    --config.num_epochs=5

הסקריפט מוריד את מערך הנתונים ומתחיל באימון. פלט הסקריפט אמור להיראות כך:

I0214 18:00:50.660087 140369022753856 train.py:146] epoch:  1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88
I0214 18:00:52.015867 140369022753856 train.py:146] epoch:  2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72
I0214 18:00:53.377511 140369022753856 train.py:146] epoch:  3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04
I0214 18:00:54.727168 140369022753856 train.py:146] epoch:  4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15
I0214 18:00:56.082807 140369022753856 train.py:146] epoch:  5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18

הסרת המשאבים

כדי לא לצבור חיובים לחשבון Google Cloud על המשאבים שבהם השתמשתם בדף הזה, פועלים לפי השלבים הבאים:

כשמסיימים להשתמש במכונת ה-TPU הווירטואלית, פועלים לפי השלבים הבאים כדי לנקות את המשאבים.

  1. אם עדיין לא עשיתם זאת, מתנתקים ממכונת ה-Cloud TPU:

    (vm)$ exit

    ההנחיה אמורה להיות עכשיו username@projectname, כדי להראות שאתם ב-Cloud Shell.

  2. כדי למחוק את Cloud TPU:

    $ gcloud compute tpus tpu-vm delete $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

  3. מריצים את הפקודה הבאה כדי לוודא שהמשאבים נמחקו. מוודאים ש-TPU לא מופיע יותר ברשימה. המחיקה עשויה להימשך כמה דקות.

    $ gcloud compute tpus tpu-vm list \
    --zone=$ZONE

הערות לגבי ביצועים

אלה כמה פרטים חשובים שרלוונטיים במיוחד לשימוש ב-TPU ב-JAX.

מרווח פנימי

אחת הסיבות הנפוצות ביותר לביצועים איטיים ב-TPU היא הוספה של padding לא מכוון:

  • מערכים ב-Cloud TPU מחולקים למקטעים. המשמעות היא שצריך להוסיף לאחד מהמאפיינים ערך כדי להגיע למספר שהוא כפולה של 8, ולמאפיין אחר ערך כדי להגיע למספר שהוא כפולה של 128.
  • הביצועים של יחידת הכפל המטריציוני הם הכי טובים עם זוגות של מטריצות גדולות שבהן הצורך בריפוד הוא מינימלי.

סוג הנתונים bfloat16

כברירת מחדל, הכפלת מטריצות ב-JAX ב-TPU משתמשת ב-bfloat16 עם צבירה של float32. אפשר לשלוט בזה באמצעות ארגומנט הדיוק בקריאות רלוונטיות לפונקציה jax.numpy (matmul,‏ dot,‏ einsum וכו'). הקפידו במיוחד על הדברים הבאים:

  • ‫precision=jax.lax.Precision.DEFAULT: שימוש בדיוק מעורב של bfloat16 (המהיר ביותר)
  • ‫precision=jax.lax.Precision.HIGH: משתמש בכמה מעברים של MXU כדי להשיג רמת דיוק גבוהה יותר
  • ‫precision=jax.lax.Precision.HIGHEST: משתמש בעוד יותר מעברים של MXU כדי להשיג דיוק מלא של float32

בנוסף, JAX מוסיף את סוג הנתונים bfloat16, שאפשר להשתמש בו כדי להמיר מערכים באופן מפורש ל-bfloat16. לדוגמה, jax.numpy.array(x, dtype=jax.numpy.bfloat16).

המאמרים הבאים

מידע נוסף על Cloud TPU: