הפעלת חישוב במכונה וירטואלית של 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 VM. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותחומים של TPU.
    • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
    • ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.

  2. מריצים את הפקודה הבאה כדי ליצור TPU VM:

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

התחברות ל-Cloud TPU VM

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

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

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

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

  1. שומרים את הטקסט הבא בקובץ בשם requirements.txt:

    --find-links https://storage.googleapis.com/libtpu-releases/index.html
--find-links https://storage.googleapis.com/libtpu-wheels/index.html
torch~=2.6.0
torch_xla[tpu]~=2.6.0
ray[default]==2.40.0

  2. מתקינים את יחסי התלות:

    pip install -r requirements.txt

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

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

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. מריצים את הסקריפט:

    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:

    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