הרצת קוד 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}