יצירת פלח TPU עם כמה מארחים

יצירת פלח TPU עם כמה מארחים

במאמר הזה מוסבר איך ליצור פלח TPU עם כמה מארחים באמצעות קבוצת מופעי מכונה מנוהלים (MIG), איך להתחבר לפלח ואיך להריץ חישוב. במדריך הזה נשתמש באפשרות של צריכה על פי דרישה. מריצים את הפקודות במדריך למתחילים הזה בטרמינל המקומי או ב-Cloud Shell.

לפני שמתחילים

נכנסים לחשבון Google Cloud . אנחנו ממליצים למשתמשים חדשים ב- Google Cloud ליצור חשבון כדי שיוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים Google Cloud פרויקט.

תפקידים שנדרשים כדי לבחור או ליצור פרויקט

• Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

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

• יוצרים Google Cloud פרויקט:

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

אם משתמשים בפרויקט קיים, מוודאים שיש את ההרשאות הנדרשות כדי להשלים את ההדרכה. אם משתמשים בפרויקט חדש, לא צריך לוודא כי כבר יש את ההרשאות הנדרשות.

מוודאים שהחיוב מופעל בפרויקט Google Cloud .

מפעילים את Compute Engine API, אם הוא עדיין לא מופעל:

תפקידים שנדרשים להפעלת ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable compute.googleapis.com

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים Google Cloud פרויקט.

תפקידים שנדרשים כדי לבחור או ליצור פרויקט

• Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

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

• יוצרים Google Cloud פרויקט:

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

אם משתמשים בפרויקט קיים, מוודאים שיש את ההרשאות הנדרשות כדי להשלים את ההדרכה. אם משתמשים בפרויקט חדש, לא צריך לוודא כי כבר יש את ההרשאות הנדרשות.

מוודאים שהחיוב מופעל בפרויקט Google Cloud .

מפעילים את Compute Engine API, אם הוא עדיין לא מופעל:

תפקידים שנדרשים להפעלת ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable compute.googleapis.com

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

כדי לקבל את ההרשאות שדרושות ליצירת קבוצת מכונות מנוהלת (MIG) שיוצרת פרוסת TPU מרובת מארחים, צריך להתחבר לכל מכונה וירטואלית בקבוצת המכונות המנוהלת באמצעות SSH ולהריץ פקודות. לשם כך, צריך לבקש מהאדמין להקצות לכם את תפקידי ה-IAM הבאים בפרויקט:

• Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1)
• משתמש בחשבון שירות (roles/iam.serviceAccountUser)
• אדמין Service Usage (roles/serviceusage.serviceUsageAdmin)

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

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

יצירת תבנית של הגדרות מכונה

כדי ליצור תבנית של הגדרות מכונה ל-TPU v6e, משתמשים בפקודה gcloud compute instance-templates create:

הערה: אם אתם משתמשים בסוג מכונה אחר, אתם צריכים לוודא שאתם משתמשים בתמונת מערכת הפעלה נתמכת. רשימה של תמונות מערכת הפעלה שמנוהלות על ידי Google לכל גרסת TPU נתמכת זמינה במאמר תמונות מערכת הפעלה של TPU.

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

יצירת מדיניות של עומס עבודה

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

כדי ליצור מדיניות עומס עבודה עבור פרוסת TPU מרובת-מארחים, משתמשים בפקודה gcloud compute resource-policies create workload-policy עם הדגל --accelerator-topology. הפקודה הבאה יוצרת מדיניות של עומס עבודה עם טופולוגיה של 2x4:

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

יצירת קבוצת מופעים מנוהלת (MIG)

מריצים את הפקודות הבאות כדי ליצור קבוצת MIG שיוצרת פלח TPU עם כמה מארחים.

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

1. כדי ליצור קבוצת מופעים מנוהלת (MIG) שיוצרת פרוסת TPU מרובת-מארחים, משתמשים בפקודה gcloud compute instance-groups managed create:

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   מחליפים את PROJECT_ID במזהה הפרויקט ב- Google Cloud .

2. כדי לוודא שהמכונות המנוהלות פועלות, משתמשים בפקודות הבאות:

   • כדי לראות את הסטטוס הכולל של ה-MIG, משתמשים בפקודה gcloud compute instance-groups managed describe:

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • כדי לראות את הסטטוס של כל מכונה, משתמשים בפקודה gcloud compute instance-groups managed list-instances:

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

   זהירות: לפני שממשיכים, חשוב לוודא שכל המופעים במצב RUNNING. כשמנסים להתחבר למכונות באמצעות SSH בזמן שהן עדיין בהקצאה, מתקבלת שגיאה.

התקנה של JAX

מתקינים יחסי תלות ואת מסגרת JAX בסביבה וירטואלית בכל מכונות ה-TPU VM ב-MIG. אם במכונות הווירטואליות של TPU מותקנת גרסת Python שקודמת ל-3.11, צריך להתקין את Python 3.11 כדי להריץ את הגרסה העדכנית של JAX.

1. בודקים איזו גרסת Python פועלת במכונות הווירטואליות של TPU:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   אם הגרסה קודמת ל-Python 3.11, צריך להתקין את Python 3.11:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. יוצרים סביבה וירטואלית:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. מתקינים את JAX בסביבה הווירטואלית:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

הרצת קוד JAX בפרוסה

כדי להריץ קוד JAX בפרוסת TPU, צריך להריץ את הקוד בכל מארח בפרוסת ה-TPU. הפונקציה jax.device_count() function call מפסיקה להגיב עד שהיא נקראת בכל מארח בפלח. בדוגמה הבאה אפשר לראות איך מריצים חישוב JAX על פרוסת TPU.

הכנת הקוד

יוצרים קובץ בשם example.py בכל מכונה:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

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

מריצים את התוכנית example.py בכל TPU VM בפרוסת ה-TPU:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

הפלט אמור להיראות כך:

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]

הסרת המשאבים

כדי לא לצבור חיובים בחשבון Google Cloud על המשאבים שבהם השתמשתם בדף הזה, אתם צריכים למחוק את הפרויקט Google Cloud יחד עם המשאבים.

לחלופין, אם אתם רוצים לשמור את הפרויקט, אתם יכולים למחוק רק את ה-MIG ואת כל מכונות ה-VM בקבוצה באמצעות הפקודה gcloud compute instance-groups managed delete:

gcloud compute instance-groups managed delete quickstart-tpu-mig --region=us-east5

המאמרים הבאים

• מידע על משאבי TPU ב-Compute Engine
• יצירת פלח TPU עם כמה מארחים באמצעות אפשרויות צריכה אחרות
• מידע על מופעי TPU בקבוצות של מופעי מכונות מנוהלים (MIG)