הפעלת חישוב במכונה וירטואלית של 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 VM

מתחברים למכונת ה-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