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

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

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

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

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

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

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

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

  1. כדי להקל על השימוש בפקודות, מגדירים כמה משתני סביבה.

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-32
export RUNTIME_VERSION=v2-alpha-tpuv5

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

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

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

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

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

  1. התקנת PyTorch/XLA בכל העובדים של TPU VM:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --worker=all \
    --command="pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

  2. משכפלים את XLA בכל העובדים של 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 slice

מריצים את סקריפט האימון על כל העובדים. סקריפט האימון משתמש בשיטת שרדינג של Single Program Multiple Data (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 הווירטואלית, פועלים לפי השלבים הבאים כדי לנקות את המשאבים.

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

    (vm)$ exit

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

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

    $ gcloud compute tpus tpu-vm delete  \
    --zone=${ZONE}

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

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