פריסת TPU Multislices ב-GKE

בדף הזה מוסבר איך לפרוס עומסי עבודה ב-Google Kubernetes Engine ‏ (GKE) באמצעות הגדרת Cloud TPU Multislice כדי לבצע אימון בקנה מידה גדול וחסכוני.

המדריך הזה מיועד למהנדסי למידת מכונה (ML) ולאדמינים ומפעילים של פלטפורמות שרוצים להשתמש באורקסטרציה של קונטיינרים ב-Kubernetes כדי לנהל עומסי עבודה של אימון, כוונון והסקת מסקנות של מודלים בקנה מידה גדול באמצעות TPU. מידע נוסף על תפקידים נפוצים ומשימות לדוגמה שמוזכרים ב Google Cloud תוכן זמין במאמר תפקידים נפוצים של משתמשי GKE ומשימות.

לפני שמגדירים את Multislice ב-GKE, חשוב להכיר את המושגים הבאים:

  1. מבוא ל-Cloud TPU
  2. ארכיטקטורת המערכת של Cloud TPU
  3. מידע על TPUs ב-GKE

מה זה TPU Multislice

‫TPU Multislice הוא ארגון ארכיטקטוני של מכונות וירטואליות ב-TPU slice, שבו שני Cloud TPU slices או יותר מתקשרים דרך רשת מרכז הנתונים (DCN). ‫Multislice מאפשרת אימון מלא, חסכוני ורחב היקף עם סקייל כמעט לינארי של עד עשרות אלפי שבבי TPU. בהגדרה של Multislice,‏ GKE פורס עומס עבודה של Multislice בכמה חלקי TPU. התקשורת בין שבבי TPU בתוך פרוסת TPU מתבצעת באמצעות חיבורים בין-שבביים (ICI). התקשורת בין הפרוסות מתבצעת דרך ה-DCN.

מומלץ להשתמש ב-Multislice אם העבודה גדולה מדי ולא נכנסת לפלח TPU יחיד.

זמינות של ריבוי-פרוסות (Multislice) ב-GKE

  • ‫Standard תומך ב-Multislice בגרסה 1.27.4-gke.900 ואילך.
  • ‫Autopilot תומך ב-Multislice בגרסה 1.29.2-gke.1521000 ואילך.
  • ‫Multislice תומך ב-frameworks של JAX ו-PyTorch. הגרסה המינימלית הנתמכת של JAX היא 2.1.
  • ‫Multislice תומך רק במאגרי צומתי TPU של multi-host TPU slice. לדוגמה, אי אפשר להשתמש ב-Multislice עם טופולוגיה של ct4p-hightpu-4t עם 2x2x1 או עם טופולוגיה של ct5lp-hightpu-4t עם 2x2, כי אלה מאגרי צמתים של פרוסות TPU עם מארח יחיד.
  • התכונה Multislice תומכת רק באימון סינכרוני של כמה בקרים.
  • עומסי עבודה של ריבוי-פרוסות יכולים לפעול רק בפרוסות TPU שחולקות את אותו סוג, גודל וטופולוגיה של TPU.
  • ‫Multislice לא תומך ב-TPU v3.

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

לפני שמתחילים, חשוב לוודא שביצעתם את המשימות הבאות:

  • מפעילים את ממשק ה-API של Google Kubernetes Engine.
  • הפעלת Google Kubernetes Engine API
  • כדי להשתמש ב-CLI של Google Cloud למשימה הזו, צריך להתקין ואז לאתחל את ה-CLI של gcloud. אם התקנתם בעבר את ה-CLI של gcloud, מריצים את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שגרסאות קודמות של ה-CLI של gcloud לא יתמכו בהרצת הפקודות שמופיעות במסמך הזה.

בחירת שיטת הטמעה של GKE

אתם יכולים להכין את מאגרי הצמתים של GKE TPU באופן ידני או לפרוס אותם באופן אוטומטי באמצעות תוכניות של Cluster Toolkit.

יצירה ידנית של מאגר צמתים ב-GKE

בקטע הזה מוסבר איך להכין מאגר צמתים במצב רגיל להרצת עומס עבודה ב-Multislice. אם אתם משתמשים במצב GKE Autopilot, דלגו אל הקטע הפעלת עומס עבודה של Multislice. באשכולות של Autopilot שמופעלת בהם גרסה 1.29.2-gke.1521000 ואילך, יחידות TPU מופעלות כברירת מחדל.

הקטע הזה כולל את השלבים הבאים:

  1. יצירה של שלושה מאגרי צמתים של פרוסות TPU עם כמה מארחים
  2. אימות הסטטוס של מאגר הצמתים
  3. הפעלת עומס עבודה של Multislice

יצירת מאגר צמתים של חלקי TPU

אפשר ליצור יותר ממאגר צמתים אחד של פרוסות TPU עם כמה מארחים. לצורך המדריך הזה, יוצרים שלושה מאגרי צמתים של פרוסות TPU מרובות-מארחים כדי להריץ עומס עבודה של Multislice. השלבים ליצירת מאגר צמתים של פרוסת TPU מרובת-מארחים שונים בהתאם לגרסת ה-TPU שבה אתם משתמשים – Ironwood‏ (TPU7x) או גרסה קודמת של TPU.

‫Ironwood (TPU7x)

אתם יכולים ליצור מאגר צמתים של פרוסת TPU עם כמה מארחים בגרסה Ironwood‏ (TPU7x) באמצעות Google Cloud CLI או Terraform:

gcloud

כדי ליצור מאגר צמתים של פרוסת TPU עם כמה מארחים באמצעות Ironwood‏ (TPU7x), צריך קודם ליצור מדיניות עומס עבודה.

  1. יוצרים מדיניות של עומס עבודה:

    gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=TPU_TOPOLOGY \
        --project=PROJECT_ID \
        --region=REGION
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫WORKLOAD_POLICY_NAME: שם למדיניות של עומס העבודה.
    • ‫TPU_TOPOLOGY: טופולוגיית TPU Ironwood‏ (TPU7x). לדוגמה, 2x2x2. כדי לראות את כל הטופולוגיות הנתמכות של Ironwood ‏ (TPU7x), אפשר לעיין בקטע בנושא טופולוגיה.
    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫REGION: האזור של מדיניות עומס העבודה. מדיניות עומס עבודה היא משאב אזורי, ואפשר לעשות בה שימוש חוזר במאגרי צמתים שחולקים את אותה טופולוגיה.
  2. יוצרים את מאגר הצמתים עם מדיניות עומס העבודה:

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --machine-type=tpu7x-standard-4t \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONE \
      --project=PROJECT_ID \
      --reservation=RESERVATION_NAME \
      --reservation-affinity=specific
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫NODE_POOL_NAME: השם של מאגר הצמתים החדש.
    • ‫CLUSTER_NAME: השם של אשכול GKE.
    • ‫WORKLOAD_POLICY_NAME: השם של מדיניות עומס העבודה שיצרתם.
    • ‫CONTROL_PLANE_LOCATION: המיקום של מישור הבקרה של האשכול ב-Compute Engine. מציינים אזור לאשכולות אזוריים או אזור לאשכולות אזוריים.
    • ‫NODE_ZONE: שם האזור בהתאם לגרסת ה-TPU שבה רוצים להשתמש.

      אופציונלי: אפשר להשתמש באזור AI, כמו us-central1-ai1a. אזורי AI הם מיקומים ייעודיים שעברו אופטימיזציה לעומסי עבודה של AI/ML בתוך Google Cloud אזורים.

      כדי לזהות מיקום זמין, אפשר לעיין במאמר זמינות של TPU ב-GKE.
    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫RESERVATION_NAME: השם של ההזמנה לשימוש.

    בפקודה הזו, הדגל --tpu-topology הוחלף בדגל --placement-policy.

Terraform

  1. ודאו שאתם משתמשים בגרסה 4.84.0 ואילך של ספק google.
  2. יוצרים מדיניות של עומס עבודה:

    resource "google_compute_resource_policy" {
      name   = "WORKLOAD_POLICY_NAME"
      region = CLUSTER_LOCATION
      workload_policy {
        type = "HIGH_THROUGHPUT"
        accelerator_topology = "TPU_TOPOLOGY"
      }
    }
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫WORKLOAD_POLICY_NAME: שם למדיניות של עומס העבודה.
    • ‫CLUSTER_LOCATION: מיקום החישוב של האשכול. מומלץ להשתמש באשכול אזורי כדי לשפר את המהימנות של מישור הבקרה של Kubernetes. אפשר גם להשתמש באשכול אזורי. מידע נוסף זמין במאמר בחירת גרסה וטופולוגיה של TPU.
    • ‫TPU_TOPOLOGY: טופולוגיית TPU Ironwood‏ (TPU7x). לדוגמה, 2x2x2. כדי לראות את כל הטופולוגיות הנתמכות של Ironwood‏ (TPU7x), אפשר לעיין במאמר בנושא תכנון של יחידות TPU.

    מידע נוסף על ההפניה google_compute_resource_policy זמין במאמר Terraform Provider.

  3. מוסיפים את הבלוק הבא להגדרות של Terraform:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        policy_name = WORKLOAD_POLICY_NAME
      }
    }
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫NODE_POOL_RESOURCE_NAME: השם של משאב מאגר הצמתים בתבנית Terraform.
    • PROJECT_ID: מזהה הפרויקט.
    • ‫CLUSTER_NAME: השם של האשכול הקיים שאליו רוצים להוסיף את מאגר הצמתים.
    • ‫POOL_NAME: השם של מאגר הצמתים שרוצים ליצור.
    • ‫NODE_ZONES: רשימה מופרדת בפסיקים של אזור אחד או יותר שבהם GKE יוצר את מאגר הצמתים.

      אופציונלי: אפשר להשתמש באזור AI, כמו us-central1-ai1a. אזורי AI הם מיקומים ייעודיים שעברו אופטימיזציה לעומסי עבודה של AI/ML בתוך Google Cloud אזורים.

    • ‫NUM_NODES: מספר הצמתים במאגר הצמתים. הערך צריך להיות אפס או מכפלה של מספר שבבי ה-TPU חלקי ארבע, כי בכל צומת של פרוסת TPU עם כמה מארחים יש ארבעה שבבים. לדוגמה, אם TPU_TOPOLOGY הוא 4x8, יש 32 צ'יפים, ולכן NUM_NODES צריך להיות 8. מידע נוסף על טופולוגיות של TPU זמין בטבלה שבמאמר בחירת גרסת ה-TPU.
    • ‫TPU_TOPOLOGY: מציין את הטופולוגיה הפיזית שנבחרה עבור פרוסת ה-TPU. הפורמט של הטופולוגיה תלוי בגרסת ה-TPU שבה אתם משתמשים. מידע נוסף על טופולוגיות של TPU זמין בטבלה שבקטע בחירת טופולוגיה.

    אפשר גם להשתמש במשתנים הבאים:

    • ‫RESERVATION_NAME: אם משתמשים בהזמנת TPU, צריך לספק רשימה של תוויות משאבי הזמנה לשימוש כשיוצרים את מאגר הצמתים. מידע נוסף על אכלוס השדה RESERVATION_LABEL_VALUES ב-reservation_affinity זמין במאמר בנושא Terraform Provider.
    • ‫autoscaling: יצירת מאגר צמתים עם הפעלה של שינוי גודל אוטומטי. כש-GKE מבצע שינוי גודל של מאגר צמתים מסוג TPU slice מרובה מארחים, הוא אטומי מגדיל את מאגר הצמתים מאפס לגודל המקסימלי.
      • ‫MAX_NODES: הגודל המקסימלי של מאגר הצמתים. הערך צריך להיות שווה למכפלה של הערכים שמוגדרים ב-TPU_TOPOLOGY ({A}x{B}x{C}) חלקי מספר השבבים בכל מכונה וירטואלית. לדוגמה, אם הערך של TPU_TOPOLOGY הוא 2x2x2, המוצר הוא 8. מכיוון שכל מכונה וירטואלית ב-tpu7x-standard-4t כוללת 4 שבבים, מספר הצמתים הוא 2.
    • ‫spot: מאגר הצמתים שישתמש במכונות וירטואליות במודל Spot לצמתים של חלקי ה-TPU. אי אפשר לשנות את ההגדרה הזו אחרי שיוצרים את מאגר הצמתים. מידע נוסף זמין במאמר בנושא מכונות וירטואליות מסוג Spot.
    • ‫flex_start: מאגר הצמתים שישתמש באפשרות הצריכה flex-start. אי אפשר להגדיר את ההגדרה הזו לערך true אם ההגדרה spot מופעלת.

גרסאות אחרות של TPU

אפשר ליצור מאגר צמתים של פרוסת TPU עם כמה מארחים בגרסאות v3,‏ v4,‏ v5p,‏ v5e ו-Trillium ‏ (v6e) באמצעות Google Cloud CLI,‏ Terraform או Google Cloud המסוף.

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      [--num-nodes=NUM_NODES] \
      [--spot \]
      [--flex-start \]
      [--enable-autoscaling \
        --max-nodes MAX_NODES]
      [--reservation-affinity=specific \
      --reservation=RESERVATION_NAME] \
      [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
      [--placement-type=COMPACT]

מחליפים את מה שכתוב בשדות הבאים:

  • ‫POOL_NAME: השם של מאגר הצמתים החדש.
  • ‫CONTROL_PLANE_LOCATION: שם האזור בהתאם לגרסת ה-TPU שבה רוצים להשתמש. כדי לזהות מיקום זמין, אפשר לעיין במאמר זמינות של TPU ב-GKE.
  • ‫CLUSTER_NAME: שם האשכול.
  • ‫NODE_ZONES: רשימה מופרדת בפסיקים של אזור אחד או יותר שבהם GKE יוצר את מאגר הצמתים.

    אופציונלי: אפשר להשתמש באזור AI, כמו us-central1-ai1a. אזורי AI הם מיקומים ייעודיים שעברו אופטימיזציה לעומסי עבודה של AI/ML בתוך Google Cloud אזורים.

  • ‫MACHINE_TYPE: סוג המכונה לשימוש בצמתים. מידע נוסף על סוגי המכונות הזמינים מופיע במאמר בחירת גרסת ה-TPU.
  • ‫TPU_TOPOLOGY: הטופולוגיה הפיזית של פרוסת ה-TPU. הפורמט של הטופולוגיה תלוי בגרסת ה-TPU. מידע נוסף על טופולוגיות של TPU זמין בטבלה שבקטע בחירת טופולוגיה.

    מידע נוסף זמין במאמר בנושא טופולוגיה.

    אפשר גם להשתמש בדגלים הבאים:

  • ‫NUM_NODES: מספר הצמתים במאגר הצמתים. הערך צריך להיות אפס או מכפלת הערכים שמוגדרים ב-TPU_TOPOLOGY ({A}x{B}x{C}) חלקי מספר השבבים בכל מכונה וירטואלית. ב-TPU v4 וב-TPU v5e עם כמה מארחים, מספר הצ'יפים בכל מכונה וירטואלית הוא ארבעה. לכן, אם TPU_TOPOLOGY הוא 2x4x4 (TPU v4 עם ארבעה שבבים בכל מכונה וירטואלית), אז NUM_NODES הוא 32 חלקי 4, כלומר 8. אם לא מציינים את הדגל הזה, מספר הצמתים מחושב ומוגדר כברירת מחדל על סמך הטופולוגיה וסוג המכונה.

  • ‫RESERVATION_NAME: השם של ההזמנה ש-GKE משתמש בה כשיוצרים את מאגר הצמתים. אם לא מציינים את הדגל הזה, GKE משתמש במאגרי צמתים זמינים של פרוסות TPU. מידע נוסף על הזמנות של TPU זמין במאמר הזמנה של TPU.

  • ‫--spot: מגדיר את מאגר הצמתים לשימוש במכונות וירטואליות במודל Spot עבור צמתים של פרוסת TPU. אי אפשר לשנות את ההגדרה הזו אחרי שיוצרים את מאגר הצמתים. מידע נוסף זמין במאמר בנושא מכונות וירטואליות מסוג Spot.

  • ‫--flex-start: הגדרת מאגר הצמתים לשימוש במכונות וירטואליות מסוג Flex-start. מכונות וירטואליות עם הפעלה גמישה נוצרות באמצעות אפשרות הצריכה flex-start, שנתמכת ב-GKE מגרסה ‎1.33.0-gke.1712000 ואילך.

  • ‫--enable-autoscaling: יצירת מאגר צמתים עם הפעלה של שינוי גודל אוטומטי. כש-GKE מבצע שינוי גודל של מאגר צמתים של פרוסת TPU מרובת מארחים, הוא מבצע שינוי גודל של מאגר הצמתים מאפס לגודל המקסימלי באופן אטומי.

    • ‫MAX_NODES: הגודל המקסימלי של מאגר הצמתים. חובה להשתמש בדגל --max-nodes אם מציינים את --enable-autoscaling, והוא צריך להיות שווה למכפלה של הערכים שמוגדרים ב-TPU_TOPOLOGY ({A}x{B}x{C}) חלקי מספר השבבים בכל מכונה וירטואלית.
  • ‫--node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: מציין ל-GKE שמאגר הצמתים של פרוסת ה-TPU עם כמה מארחים הוא אוסף. משתמשים בדגל הזה אם התנאים הבאים מתקיימים:

    • מאגר הצמתים מריץ עומסי עבודה של היקשים.
    • מאגר הצמתים משתמש ב-TPU Trillium.
    • אי אפשר לתזמן איסוף של מכונות וירטואליות במודל Spot.

    מידע נוסף על ניהול תזמון איסוף נתונים זמין במאמר ניהול תזמון איסוף נתונים בפרוסות TPU מרובות מארחים.

  • ‫--placement-type=COMPACT: יצירת מאגר צמתים עם מיקום קומפקטי מופעל. חובה להשתמש באפשרות הזו עם הדגל --tpu-topology. מידע נוסף זמין במאמרים בנושא יצירת מדיניות מיקום קומפקטית וטופולוגיה של TPU.

Terraform

  1. ודאו שאתם משתמשים בגרסה 4.84.0 ואילך של ספק google.
  2. מוסיפים את הבלוק הבא להגדרות של Terraform:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        type = "COMPACT"
        tpu_topology = TPU_TOPOLOGY
      }
    }
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫NODE_POOL_RESOURCE_NAME: השם של משאב מאגר הצמתים בתבנית Terraform.
    • PROJECT_ID: מזהה הפרויקט.
    • ‫CLUSTER_NAME: השם של האשכול הקיים שאליו רוצים להוסיף את מאגר הצמתים.
    • ‫POOL_NAME: השם של מאגר הצמתים שרוצים ליצור.
    • ‫CLUSTER_LOCATION: מיקום החישוב של האשכול. מומלץ להשתמש באשכול אזורי כדי לשפר את המהימנות של מישור הבקרה של Kubernetes. אפשר גם להשתמש באשכול אזורי. מידע נוסף זמין במאמר בנושא בחירת גרסה וטופולוגיה של TPU.
    • ‫NODE_ZONES: רשימה מופרדת בפסיקים של אזור אחד או יותר שבהם GKE יוצר את מאגר הצמתים.

      אופציונלי: אפשר להשתמש באזור AI, כמו us-central1-ai1a. אזורי AI הם מיקומים ייעודיים שעברו אופטימיזציה לעומסי עבודה של AI/ML בתוך Google Cloud אזורים.

    • ‫NUM_NODES: מספר הצמתים במאגר הצמתים. הערך צריך להיות אפס או מכפלה של מספר שבבי ה-TPU חלקי ארבע, כי בכל צומת של פרוסת TPU עם כמה מארחים יש 4 שבבים. לדוגמה, אם TPU_TOPOLOGY הוא 4x8, יש 32 צ'יפים, ולכן NUM_NODES צריך להיות 8. מידע נוסף על טופולוגיות של TPU זמין בטבלה שבמאמר בחירת גרסת ה-TPU.
    • ‫TPU_TOPOLOGY: מציין את הטופולוגיה הפיזית של חלוקת ה-TPU. הפורמט של הטופולוגיה תלוי בגרסת ה-TPU שבה אתם משתמשים. מידע נוסף על טופולוגיות של TPU זמין בטבלה שבקטע בחירת טופולוגיה.

    אפשר גם להשתמש במשתנים הבאים:

    • ‫RESERVATION_NAME: אם משתמשים בהזמנת TPU, זוהי רשימת התוויות של משאבי ההזמנה שבהן צריך להשתמש כשיוצרים את מאגר הצמתים. מידע נוסף על אכלוס השדה RESERVATION_LABEL_VALUES ב-reservation_affinity זמין במאמר בנושא Terraform Provider.
    • ‫autoscaling: יצירת מאגר צמתים עם הפעלה של שינוי גודל אוטומטי. כש-GKE מבצע שינוי גודל של מאגר צמתים של פרוסת TPU מרובת מארחים, הוא מבצע שינוי גודל של מאגר הצמתים מאפס לגודל המקסימלי באופן אטומי.
      • ‫MAX_NODES: הגודל המקסימלי של מאגר הצמתים. הערך הזה צריך להיות שווה למכפלה של הערכים שמוגדרים ב-TPU_TOPOLOGY ({A}x{B}x{C}) חלקי מספר הצ'יפים בכל מכונה וירטואלית).
    • ‫spot: מאפשר למאגר הצמתים להשתמש במכונות וירטואליות במודל Spot עבור צמתי פרוסת ה-TPU. אי אפשר לשנות את ההגדרה הזו אחרי שיוצרים את מאגר הצמתים. מידע נוסף זמין במאמר בנושא מכונות וירטואליות מסוג Spot.
    • ‫flex_start: מגדיר את מאגר הצמתים לשימוש באפשרות הצריכה flex-start. אי אפשר להגדיר את הערך true אם האפשרות spot מופעלת.

המסוף

כדי ליצור מאגר צמתים עם מעבדי TPU:

  1. נכנסים לדף Google Kubernetes Engine במסוף Google Cloud .

    מעבר אל Google Kubernetes Engine

  2. ברשימת האשכולות, לוחצים על שם האשכול שרוצים לשנות.

  3. לוחצים על הכרטיסייה Nodes (צמתים).

  4. לוחצים על Create user-managed node pool (יצירת מאגר צמתים בניהול המשתמש).

  5. בקטע פרטי מאגר צמתים, מסמנים את התיבה ציון מיקומי צמתים.

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

  7. בחלונית הניווט, לוחצים על צמתים.

  8. בקטע Machine Configuration (הגדרת המכונה), בוחרים באפשרות TPUs.

  9. בתפריט הנפתח סדרה, בוחרים באחת מהאפשרויות הבאות:

    • ‫CT3: TPU v3, מכשיר עם מארח יחיד
    • ‫CT3P: TPU v3, multi-host pod slice
    • ‫CT4P: TPU v4
    • ‫CT5LP: TPU v5e
    • ‫CT5P: TPU v5p
    • ‫CT6E: TPU Trillium (v6e)
  10. בתפריט הנפתח Machine type (סוג המכונה), בוחרים את שם המכונה שרוצים להשתמש בה לצמתים. כדי להגדיר את סוג המכונה ואת טופולוגיית ה-TPU שיוצרים מאגר צמתים של פרוסת TPU מרובת מארחים, אפשר להיעזר בטבלה בחירת גרסת ה-TPU.

  11. בתפריט הנפתח TPU Topology, בוחרים את הטופולוגיה הפיזית של פרוסת ה-TPU.

  12. בתיבת הדו-שיח נדרשים שינויים, לוחצים על ביצוע שינויים.

  13. מוודאים שסוג דיסק האתחול הוא דיסק מתמיד סטנדרטי או דיסק מתמיד שמבוסס על SSD.

  14. אופציונלי: מסמנים את תיבת הסימון Enable nodes on spot VMs כדי להשתמש במכונות וירטואליות מסוג Spot עבור הצמתים במאגר הצמתים.

  15. לוחצים על יצירה.

אימות הסטטוס של מאגר הצמתים

  1. קבלת פרטי כניסה כדי להשתמש ב-kubectl כדי לגשת לאשכול:

    gcloud container clusters get-credentials CLUSTER_NAME \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫CLUSTER_NAME: שם האשכול.
    • PROJECT_ID: מזהה הפרויקט.
    • ‫CONTROL_PLANE_LOCATION: המיקום של מישור הבקרה של האשכול ב-Compute Engine. מציינים אזור לאשכולות אזוריים או אזור לאשכולות אזוריים.
  2. כדי לראות את הצמתים של חלקי ה-TPU, משתמשים בפקודה kubectl ב-Cloud Shell:

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
       -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫ACCELERATOR_TYPE: סוג המאיץ של מאגר הצמתים של פרוסת ה-TPU. הערך הזה צריך להיות מיפוי לסוג המכונה שבה השתמשתם במהלך יצירת מאגר הצמתים. לדוגמה, אם יצרתם מאגר צמתים עם סוג המכונה ct5lp-hightpu-4t, משתמשים ב-tpu-v5-lite-podslice. כדי לדעת באיזה סוג של מאיץ כדאי להשתמש, אפשר לעיין בטבלה שבמאמר תכנון השימוש ב-TPU ב-GKE.
    • ‫TPU_TOPOLOGY: הטופולוגיה הפיזית של פרוסת ה-TPU.

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

     NAME                                    STATUS   ROLES    AGE    VERSION
     gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000
    

הפעלת עומס עבודה של Multislice

בקטע הזה מריצים עומס עבודה של JAX שמציג את המספר הכולל של שבבי TPU בפרוסת ה-TPU, ואז יוצאים.

כדי להריץ עומס עבודה של JAX:

  1. יוצרים את קובץ המניפסט tpu-multislice.yaml הבא:

    טייס אוטומטי

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 2 # This value defines a "Multislice" setup, by creating 2 independent TPU slices.
          template:
            spec:
              parallelism: 4
              completions: 4
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 4x4
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                     limits:
                        google.com/tpu: 4 # This container requires 4 TPU chips
    

    במניפסט הזה, בוחרים tpu-v5-lite-podslice עם טופולוגיה של 4x4. ההגדרה הזו אומרת:

    • בטופולוגיה יש 16 צ'יפים בסך הכול (4*4=16).
    • ‫tpu-v5-lite-podslice נפרס ב-ct5lp-hightpu-4t מכונות עם ארבעה צ'יפים לכל מכונה וירטואלית. כלומר, כדי לקבל טופולוגיה של 4x4, צריך ארבע מכונות וירטואליות (16/4=4). לכן, בשדות parallelism ו-completions צריך להזין את הערך 4.
    • כל קונטיינר ב-Pod מבקש ארבעה שבבי TPU, שזהה למספר השבבים לכל מכונה וירטואלית. המידע הזה מוגדר במניפסט בשורה google.com/tpu: 4.
    • השדה replicas מוגדר ל-2 כדי ליצור שני פלחים.

    מידע נוסף על הערכים האפשריים של המשתנים האלה זמין בטבלה שבמאמר תכנון TPU ב-GKE.

    רגילה

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: NUM_SLICES
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                       google.com/tpu: NUM_CHIPS
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫NUM_SLICES: מספר מאגרי הצמתים של פרוסות TPU. במקרה הזה, NUM_SLICES שווה ל-3.
    • ‫ACCELERATOR_TYPE: סוג המאיץ של מאגר הצמתים של פרוסת ה-TPU. הערך הזה צריך להיות ממופה לסוג המכונה שבו השתמשתם במהלך יצירת מאגר הצמתים. לדוגמה, אם יצרתם מאגר צמתים עם סוג המכונה ct5lp-hightpu-4t, משתמשים ב-tpu-v5-lite-podslice. כדי לדעת באיזה סוג של מאיץ כדאי להשתמש, אפשר לעיין בטבלה במאמר תכנון השימוש ב-TPU ב-GKE.
    • ‫TPU_TOPOLOGY: הטופולוגיה הפיזית של פרוסת ה-TPU. לדוגמה, 4x4x4 או 2x2 בהתאם לגרסת ה-TPU.
    • ‫NUM_NODES: מספר הצמתים במאגר הצמתים. הערך צריך להיות אפס או מכפלת הערכים שמוגדרים ב-TPU_TOPOLOGY ({A}x{B}x{C}) חלקי מספר שבבי ה-TPU בכל מכונה וירטואלית. ב-TPU v4 עם כמה מארחים, מספר שבבי ה-TPU בכל מכונה וירטואלית הוא ארבעה. ב-TPU v5e עם כמה מארחים, מספר שבבי ה-TPU בכל מכונה וירטואלית הוא אחד, ארבעה או שמונה. לכן, אם TPU_TOPOLOGY הוא 2x4x4 (TPU v4 עם ארבעה שבבי TPU בכל מכונה וירטואלית), אז NUM_NODES הוא 32/4, כלומר 8.
    • ‫NUM_CHIPS: מספר הצ'יפים לכל מכונה וירטואלית.
  2. החלת המניפסט:

    kubectl apply -f tpu-multislice.yaml
    
  3. מוודאים שעומס העבודה התקבל:

    kubectl get jobsets
    

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

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. עוקבים אחרי הסטטוס של ה-Pods שהוקצו:

    kubectl get pods
    

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

     NAME                                READY   STATUS      RESTARTS   AGE
     multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
     multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
     multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
     multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
     multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
     multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s
    

    ‫multislice-job JobSet מתזמן, יוצר ואז מריץ את קבוצות ה-Pod עד לסיום. שמות ה-Pod הם בפורמט <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. הקידומת jobsetName קובעת את JobSet שאליו ה-Pod שייך.

  5. אופציונלי: מסירים את עומס העבודה של JAX:

    kubectl delete -f tpu-multislice.yaml
    

שימוש ב-Cluster Toolkit כדי ליצור משאבים ולשלוח עומסי עבודה של multislice

  1. הגדרת Cluster Toolkit
  2. כדי לפרוס אשכול GKE multislice TPU, פועלים לפי ההוראות במאמר פריסת אשכול GKE TPU 7x או פריסת אשכול GKE TPU v6e (למידע על תוכניות בסיסיות של TPU v4, ‏ v5e ו-v5p, אפשר לעיין במאגר Cluster Toolkit GitHub).

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

    • מגדירים את num_slices למספר פרוסות ה-TPU שרוצים להקצות (ב-Multislice, מציינים ערך גדול מ-1, כמו 2 או 4).
    • מגדירים את reservation לשם הזמנת ה-TPU.
  3. שולחים עומס עבודה של כמה פרוסות באמצעות הפקודה gcluster job submit. מידע נוסף על אפשרויות לשליחת משימות זמין במדריך לשליחת משימות ב-Cluster Toolkit.

    gcluster job submit \
        --name=WORKLOAD_NAME \
        --cluster=DEPLOYMENT_NAME \
        --project=PROJECT_ID \
        --location=ZONE \
        --compute-type=COMPUTE_TYPE \
        --topology=TOPOLOGY \
        --num-slices=NUM_SLICES \
        --image=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
        --command="python -c 'import jax; print(\"TPU cores:\", jax.device_count())'"
    

    מחליפים את מה שכתוב בשדות הבאים:

    • ‫WORKLOAD_NAME: השם של עומס העבודה.
    • ‫DEPLOYMENT_NAME: השם של פריסת Cluster Toolkit (שהוא גם השם של אשכול GKE).
    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫ZONE: האזור שבו נמצא האשכול.
    • ‫COMPUTE_TYPE: סוג המחשוב שתואם לסוג המאיץ של מאגר הצמתים, לדוגמה, v6e-16 או tpu7x-standard-4t.
    • ‫TOPOLOGY: הטופולוגיה הפיזית שתואמת למאגר הצמתים, לדוגמה, 2x4 או 2x2x2.
    • ‫NUM_SLICES: מספר פרוסות ה-TPU להקצאה לעומס העבודה (בהתאם להגדרת פרוסת האשכול). לדוגמה, 2 או 4.

הגדרת הגדרות נוספות

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

שיפור ביצועי הרשת באמצעות hostNetwork

כדי לשפר את הביצועים של הרשת בין חלקי TPU, מומלץ להפעיל את hostNetworking. משתמשים ב-hostNetwork: true במפרט של ה-Pod כדי לדלג על כל מחסנית הרשת של Kubernetes ולאפשר ל-Pods של Kubernetes להשתמש ברשת המארח ישירות לתקשורת בין מכונות וירטואליות.

כדי להפעיל את hostNetworking, מוסיפים את שתי השורות הבאות למפרט ה-Pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

כדי להמשיך להשתמש ב-podHostnames לחיפוש צמתים של עובדים באמצעות hostNetwork, צריך להגדיר את dnsPolicy: ClusterFirstWithHostNet. זה חשוב כשמריצים משימות אימון עם הפעלה אוטומטית של המשך האימון, וצריך להשתמש באותם שמות כדי לטעון מחדש את אותן נקודות ביקורת.

אם אתם משתמשים ב-TPU Trillium‏ (v6e) או ב-Ironwood‏ (TPU7x) (גרסת Preview) וה-Pods שלכם משתמשים ב-hostNetworking, אתם צריכים להתקין את DaemonSet הבא כדי לכוונן את /proc/sys/net/ipv4/tcp_rmem בצומת.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

שיפור ביצועי הרשת ללא hostNetwork ב-TPU Trillium או ב-Ironwood (TPU7x)

אם אתם משתמשים ב-TPU Trillium או ב-Ironwood‏ (TPU7x) (גרסת Preview) ואין אפשרות להשתמש ב-hostNetworking ב-Pods שלכם, כדאי להטמיע GKE DRANET כדי להקצות מכשירי רשת לביצועי הרשת הטובים ביותר. התכונה GKE DRANET מאפשרת לבקש משאבי רשת עם ביצועים גבוהים בשביל ה-Pods, תוך עקיפה של Kubernetes ושל GKE Dataplane V2.

סוג המכונה ct6e-standard-4t מגובה על ידי שני כרטיסי NIC פיזיים. ‫Kubernetes דורש vNIC אחד שלא ניתן להעביר לקבוצות Pod. לכן, לכל צומת צריכים להיות שלושה vNICs כדי לאפשר ל-Pods גישה ישירה לשני vNICs, וכך להשיג את הביצועים הטובים ביותר של שני ה-NICs הפיזיים.

כדי להפעיל את התכונה DRANET ולהקצות מכשירי רשת ל-ct6e-standard-4t, מבצעים את השלבים הבאים:

  1. יוצרים עוד שני VPC באמצעות הפקודות הבאות:

    gcloud compute networks subnets create ADDITIONAL_SUBNET_1 \
        --project=PROJECT_ID \
        --network=ADDITIONAL_VPC_1 \
        --location=CONTROL_PLANE_LOCATION
    
    gcloud compute networks subnets create ADDITIONAL_SUBNET_2 \
        --project=PROJECT_ID \
        --network=ADDITIONAL_VPC_2 \
        --location=CONTROL_PLANE_LOCATION
    
  2. יצירת אשכול GKE שמשתמש ב-GKE Dataplane V2.

  3. כדי ליצור מאגר צמתים עם הרשתות שנוצרו ועם GKE DRANET מופעל, מוסיפים את התווית cloud.google.com/gke-networking-dra-driver=true:

    gcloud container node-pools create NODE_POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --machine-type=MACHINE_TYPE \
      --num-nodes=NUM_NODES \
      --additional-node-network network=ADDITIONAL_VPC_1,subnetwork=ADDITIONAL_SUBNET_1 \
      --additional-node-network network=ADDITIONAL_VPC_2,subnetwork=ADDITIONAL_SUBNET_2 \
      --node-labels=cloud.google.com/gke-networking-dra-driver=true
    
  4. אחרי אימות הסטטוס של מאגר הצמתים, יוצרים עומס עבודה עם מכשירי הרשת הנוספים על ידי החלת ResourceClaimTemplate:

    apiVersion: resource.k8s.io/v1
    kind: ResourceClaimTemplate
    metadata:
      name: two-netdev
    spec:
      spec:
        devices:
          requests:
          - name: req-netdev
            exactly:
              deviceClassName: netdev.google.com
              allocationMode: ExactCount
              count: 2
    ---
    apiVersion: v1
    kind: Pod
    metadata:
      name: sample-netdevice-pod-1
    spec:
      containers:
      - name: sample-netdevice-pod
        image: busybox
        command: ["sleep", "infinity"]
        resources:
          claims:
          - name: netdev
          limits:
            google.com/tpu: 4
      restartPolicy: Always
      resourceClaims:
      - name: netdev
        resourceClaimTemplateName: two-netdev
    
  5. מוודאים שהממשקים זמינים ב-Pod:

    kubectl exec -it pods/sample-netdevice-pod-1 -- ip a
    
שיטה מומלצת:

כדי לשפר את הביצועים של הרשת, משתמשים בממשקים eth1 ו-eth2 במקום בממשק eth0. כדי לעשות את זה, מוסיפים את export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" למפרט של עומס העבודה.

הפעלת הרישום ביומן

יומנים שנוצרו על ידי קונטיינרים שפועלים בצמתים של GKE, כולל צמתים של פרוסות TPU, מוצגים ב-Logs Explorer אם הפעלתם את רישום היומנים של מערכת GKE באשכול.

אתם יכולים לראות את היומנים מ-GKE באמצעות Logs Explorer עם המסנן הבא כדי לראות את יומני הקונטיינרים של עומס העבודה:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

משתמשים במסנן הבא עבור פרוסת TPU ועובדים:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

מידע נוסף מופיע במאמר בנושא צפייה ביומני TPU ב-GKE.

הפעלת מדדים נוספים

בנוסף למדדי ה-TPU הכלליים, יש 4 מדדים נוספים של זמן ריצה של TPU שספציפיים ל-multislice. המדדים האלה זמינים ב-GKE גרסה ‎1.29.1-gke.1016000 ואילך. עומס העבודה של TPU חייב להשתמש בגרסה 0.4.24 של JAX

אלה המדדים הזמינים של פילוח רב-ממדי:

  • זמני השהיה בהעברה ב-DCN (רשת מרכזי נתונים): התפלגות של זמני השהיה בהעברת נתונים ברשת עבור תנועה מרובת פרוסות.
  • זמני תגובה מצטברים: התפלגות של זמן התגובה המצטבר מקצה לקצה לתנועה של כמה פרוסות.
  • זמני השהיה בהעברה ממארח למכשיר: חלוקת זמני ההשהיה בהעברה ממארח למכשיר לכל נתח נתונים בתנועה של כמה פרוסות.
  • זמני השהיה בהעברה ממכשיר למארח: חלוקת זמני ההשהיה בהעברה ממכשיר למארח לכל קבוצת נתונים בתנועה מרובת-פרוסות.

המדדים האלה נמצאים בסכימת הקונטיינר של Kubernetes ‏ (k8s_container):

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

פרוסת TPU לעומת Multislice

בטבלה הבאה מפורטים ההבדלים בארגון הארכיטקטוני של TPU slice ו-Multislice:

TPU slice Multislice
קישוריות הדדית עומס העבודה פועל בפלח TPU יחיד. כל שבבי ה-TPU בפרוסת TPU מחוברים באמצעות ICI. עומס העבודה פועל בכמה חלקי TPU. התקשורת בתוך פרוסת נתונים מתבצעת באמצעות ICI. התקשורת בין הפרוסות מתבצעת באמצעות DCN.
מאגרי צמתים נתמכים פרוסת TPU במארח יחיד ופרוסת TPU בכמה מארחים קבוצות של חלקי TPU עם כמה מארחים
סוג העומס המומלץ ‫IndexedJob או JobSet JobSet

לפנות משאבים

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

מחיקת הפרויקט

    כדי למחוק Google Cloud פרויקט:

    gcloud projects delete PROJECT_ID

מחיקת משאבים בודדים

מוחקים את אשכול GKE:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

מחיקת משאבים ב-Cluster Toolkit

אם השתמשתם ב-Cluster Toolkit כדי ליצור את האשכול, צריך למחוק את הפריסה כדי להימנע מחיובים נוספים:

gcluster destroy DEPLOYMENT_NAME

מחליפים את DEPLOYMENT_NAME בשם של ספריית הפריסה שנוצרה כשפרסתם את תוכנית האב.

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