הרצת קוד PyTorch בפרוסות TPU

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

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

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

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

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

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

יצירת פרוסת Cloud TPU

  1. הגדרת משתני סביבה לפרמטרים של פקודות ב-Google Cloud CLI.

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central1-b
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e

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

    • ‫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 \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --accelerator-type=$ACCELERATOR_TYPE \
     --version=$RUNTIME_VERSION

התקנת PyTorch/XLA בפרוסה

אחרי שיוצרים את פרוסת ה-TPU, צריך להתקין את PyTorch בכל המארחים בפרוסת ה-TPU. כדי לעשות את זה, משתמשים בפקודה gcloud compute tpus tpu-vm ssh עם הפרמטרים --worker=all ו---command.

  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
torchvision
ray[default]==2.40.0

  2. מעתיקים את requirements.txt לכל מכונה וירטואלית בפרוסה:

    gcloud compute tpus tpu-vm tpu-vm scp ./requirements.txt \
     $TPU_NAME:~/ \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --worker=all

  3. מתקינים את יחסי התלות בכל VM בפרוסה:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --worker=all \
     --command="pip3 install -r requirements.txt"

  4. משכפלים את XLA בכל ה-workers של ה-TPU VM:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --worker=all \
     --command="git clone https://github.com/pytorch/xla.git"

הרצת סקריפט אימון בפרוסת ה-TPU

מריצים את סקריפט האימון בכל ה-workers. סקריפט האימון משתמש באסטרטגיית חלוקה של תוכנית אחת עם נתונים מרובים (SPMD). מידע נוסף על SPMD זמין במדריך למשתמש של PyTorch/XLA SPMD.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --project=$PROJECT_ID \
   --worker=all \
   --command="PJRT_DEVICE=TPU python3 ~/xla/test/spmd/test_train_spmd_imagenet.py  \
   --fake_data \
   --model=resnet50  \
   --num_epochs=1 2>&1 | tee ~/logs.txt"

ההדרכה נמשכת כ-15 דקות. בסיום התהליך אמורה להופיע הודעה דומה לזו:

   Epoch 1 test end 23:49:15, Accuracy=100.00
   10.164.0.11 [0] Max Accuracy: 100.00%

הסרת המשאבים

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

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

    exit

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

  2. מוחקים את המשאבים של Cloud TPU.

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

  3. כדי לוודא שהמשאבים נמחקו, מריצים את הפקודה gcloud compute tpus tpu-vm list. המחיקה עשויה להימשך כמה דקות. הפלט של הפקודה הבאה לא אמור לכלול אף אחד מהמשאבים שנוצרו במדריך הזה:

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