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

במאמר הזה מופיע מבוא קצר לעבודה עם PyTorch ו-Cloud TPU.

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

לפני שמריצים את הפקודות במסמך הזה, צריך ליצור Google Cloud חשבון, להתקין את Google Cloud CLI ולהגדיר את הפקודה gcloud. מידע נוסף זמין במאמר בנושא הגדרת סביבת 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. מריצים את הפקודה הבאה כדי ליצור מכונת 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 VM ללא כתובת IP חיצונית, פועלים לפי ההוראות במאמר איך מתחברים ל-TPU VM ללא כתובת IP ציבורית.

התקנת PyTorch/XLA ב-TPU VM

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

אימות הגישה של PyTorch ל-TPU

משתמשים בפקודה הבאה כדי לוודא של-PyTorch יש גישה ל-TPU:

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

הפלט של הפקודה אמור להיראות כך:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

ביצוע חישוב בסיסי

  1. יוצרים קובץ בשם tpu-test.py בספרייה הנוכחית ומעתיקים את הסקריפט הבא ומדביקים אותו בקובץ:

    import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

  2. מריצים את הסקריפט:

    (vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

    הפלט מהסקריפט מציג את תוצאת החישוב:

    tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

הסרת המשאבים

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

  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

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

מידע נוסף על מכונות וירטואליות של Cloud TPU