הרצת חישוב ב-Cloud TPU VM באמצעות 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: