אימון מודלים של למידת מכונה בקנה מידה גדול ב-GKE באמצעות Multi-Tier Checkpointing

Standard

בדף הזה מוסבר איך להשתמש ב-Multi-Tier Checkpointing כדי לאחסן ולנהל באופן מהימן נקודות ביקורת במהלך אימון של מודל למידת מכונה ב-GKE. האחסון והניהול של נקודות הבדיקה הם קריטיים למשימות אימון בקנה מידה גדול, שמוגדרות כמשימות שמשתמשות ביותר מאלפי צמתים. הפרעות למשימות בקנה מידה גדול הן תכופות (יכולות לקרות מדי שעה) וההתאוששות מהן יכולה להיות איטית.

יתרונות

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

• ניהול מלא של נתוני נקודות ביקורת, כולל גיבויים, שכפול ושחזור אוטומטי של עומסי העבודה הבאים:
  • הפעלת אימון של JAX באמצעות Orbax לניהול מצב ב-TPU.
  • עומסי עבודה של PyTorch במעבדי GPU.
• שחזור מהיר של משימות אימון מנקודת ביקורת שמאוחסנת בצומת המקומי. אפשר גם לשחזר באמצעות נקודות ביקורת שמאוחסנות בצומת אחר באשכול האימון.
• שחזור מהיר של משימות אימון מנקודת ביקורת שמאוחסנת בגיבוי של Cloud Storage בתרחישים של המקרה הגרוע ביותר, שבהם אין נקודות ביקורת בתוך האשכול.

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

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

• מפעילים את ממשק ה-API של Google Kubernetes Engine.
הפעלת Google Kubernetes Engine API
• כדי להשתמש ב-Google Cloud CLI למשימה הזו, מתקינים ואז מאתחלים את ה-CLI של gcloud. אם כבר התקנתם את ה-CLI של gcloud, הריצו את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שבגרסאות קודמות של ה-CLI של gcloud אי אפשר להריץ את הפקודות שבמסמך הזה.
  הערה: אם כבר התקנתם את ה-CLI של gcloud, הקפידו להגדיר את compute/region המאפיין. אם אתם משתמשים בעיקר באשכולות אזוריים, הגדירו במקום זאת את compute/zone. אם מגדירים מיקום ברירת מחדל, אפשר להימנע משגיאות ב-CLI של gcloud כמו One of [--zone, --region] must be supplied: Please specify location. יכול להיות שתצטרכו לציין את המיקום בפקודות מסוימות אם המיקום של האשכול שונה מברירת המחדל שהגדרתם.

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

דרישות

כדי להשתמש ב-Multi-Tier Checkpointing, צריך אשכול GKE בגרסה 1.32.4-gke.1415000 ואילך.

מגבלות

• אין תמיכה באשכולות Autopilot.

הגדרת צומתי GKE לשימוש ב-Multi-Tier Checkpointing

בקטע הזה מוסבר איך להגדיר צומתי GKE באשכולות חדשים וקיימים.

הגדרת צמתים באשכול חדש

1. יוצרים אשכול שבו מופעלים Multi-Tier Checkpointing, מנהל התקן ה-CSI של Cloud Storage FUSE ואיחוד זהויות של עומסי עבודה ל-GKE. אם אתם משתמשים בפרוסות TPU לעומס העבודה של למידת המכונה, תצטרכו לשנות את הפקודה ליצירת אשכול כך שתכלול את ההגדרות של מאגר צמתים של פרוסות TPU.

   gcloud container clusters create CLUSTER_NAME \
       --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
       --node-locations=NODE_LOCATION \
       --workload-pool=PROJECT_ID.svc.id.goog \
       --cluster-version=CLUSTER_VERSION
       --location=CLUSTER_LOCATION \
       --machine-type=MACHINE_TYPE \
       --num-nodes=NUM_NODES
   

   מחליפים את הערכים הבאים:

   • ‫CLUSTER_NAME: שם האשכול.
   • ‫NODE_LOCATION: האזור של צמתי האשכול. כאן נמצא נפח ה-TPU שלכם.
   • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
   • ‫CLUSTER_VERSION: גרסת האשכול. ‫1.32.4-gke.1415000 היא הגרסה המינימלית שנתמכת.
   • ‫CLUSTER_LOCATION: האזור שבו רוצים ליצור את האשכול.
   • ‫MACHINE_TYPE: סוג המכונה שמשמש לצמתים שמריצים רכיבים כמו בקר JobSet ובקר multi-tier checkpointing. לאימון בקנה מידה גדול, מומלץ להשתמש ב-e2-standard-4 מכונות לפחות. לא תשתמשו בסוג המכונה הזה לאימון מודלים. במקום זאת, תיצרו מאגרי צמתים נפרדים למטרה הזו, בדרך כלל באמצעות משפחות של מכונות וירטואליות שעברו אופטימיזציה למאיצים.
   • ‫NUM_NODES: מספר הצמתים שייווצרו בכל אחד מהאזורים של האשכול.

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

כדי להשתמש ב-Multi-Tier Checkpointing באשכול קיים, מפעילים אותו יחד עם מנהל התקן ה-CSI של Cloud Storage FUSE ואיחוד הזהויות של עומסי העבודה ל-GKE באמצעות הפקודות הבאות. גרסת האשכול הקיימת צריכה להיות חדשה יותר מ-1.32.3-gke.1170000.

1. הפעלת איחוד זהויות של עומסי עבודה ל-GKE:

   gcloud container clusters update CLUSTER_NAME \
       --workload-pool=PROJECT_ID.svc.id.goog \
       --location=CLUSTER_LOCATION
   

   מחליפים את הערכים הבאים:

   • ‫CLUSTER_NAME: שם האשכול.
   • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
   • ‫CLUSTER_LOCATION: האזור של האשכול.

2. מפעילים את Multi-Tier Checkpointing ואת מנהל התקן ה-CSI של Cloud Storage FUSE:

   gcloud container clusters update CLUSTER_NAME \
       --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
       --location=CLUSTER_LOCATION
   

הגדרת הרשאות לשימוש ב-Multi-Tier Checkpointing

בקטע הזה מוסבר איך להגדיר הרשאות לשימוש ב-Multi-Tier Checkpointing.

איך מעניקים גישה לקטגוריות של Cloud Storage

נפחי האחסון הזמניים שמשמשים את מנהל ההתקן של ה-CSI של Multi-Tier Checkpointing חייבים להשתמש בקטגוריות קיימות של Cloud Storage.

כדי לשמור נקודות ביקורת בקטגוריה של Cloud Storage, צריך לתת ל-Multi-Tier Checkpointing גישה לקטגוריה. כדי להשתמש בתכונה Multi-Tier Checkpointing, צריך להקצות את תפקיד ה-IAM ‏Storage Object User ‏ (roles/storage.objectUser) ב-bucket לחשבון השירות של Kubernetes.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

מחליפים את הערכים הבאים:

• ‫GCS_BUCKET: השם של הקטגוריה של Cloud Storage שממנה רוצים להעביר את הנתונים.
• ‫PROJECT_ID: Google Cloud מזהה הפרויקט.
• ‫PROJECT_NUMBER: מזהה ייחודי שנוצר באופן אוטומטי לפרויקט. כדי למצוא את הערך הזה, עיינו במאמר יצירה וניהול של פרויקטים.

(אופציונלי) מעניקים גישה לחשבון השירות שמוגדר כברירת המחדל של Compute Engine

אם המכונות של Compute Engine צריכות גישת קריאה לקטגוריה של Cloud Storage, צריך לתת את תפקיד ה-IAM 'צפייה באובייקט אחסון' (roles/storage.objectViewer) לחשבון השירות שמוגדר כברירת מחדל ב-Compute Engine.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:[e-mail] \
    --role roles/storage.objectViewer

פריסת הבקר של JobSet

בקר JobSet אחראי לניהול של משימות אצווה שמריצות את אימון המודל ב-GKE, והקצאת המשאבים שלו מותאמת לטיפול בעומס העבודה בצורה יעילה. מוודאים שמפעיל משימת האימון פורס את JobSet ומשתמש בו.

כדי להגדיל את בקשת הזיכרון ל-1Gi, את מגבלת הזיכרון ל-2Gi ואת בקשת המעבד ל-1 בקונטיינר של המנהל בפריסת JobSet, מריצים את פקודת התיקון הבאה:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

הפעלת דרייבר CSI של Multi-Tier Checkpointing

בקטע הזה מוסבר איך לאתחל את מנהל ההתקן של CSI לסימון נקודות ביקורת (checkpointing) מרובות רמות בצמתים שבהם יפעלו עומסי העבודה. דרייבר ה-CSI אחראי לטיפול באחסון ובניהול של נקודות ביקורת במהלך תהליך אימון המודל.

יצירת CheckpointConfiguration

‫CheckpointConfiguration הוא משאב בהתאמה אישית ב-Kubernetes שמציין מאפיינים לפריסת מנהל ההתקן Multi-Tier Checkpointing CSI. המשאב הזה הוא בהיקף אשכול.

1. יוצרים את המניפסט הבא של checkpoint.yaml.

   apiVersion: checkpointing.gke.io/v1
   kind: CheckpointConfiguration
   metadata:
     name: MTC_CONFIG_NAME-configuration
   spec:
       cloudStorageBucketName: GCS_BUCKET
       nodeSelector:
           node.kubernetes.io/instance-type: MACHINE_TYPE
       tolerations:
       - key: TOLERATION_KEY
           operator: Exists
           effect: NoSchedule
       inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
       gcsFuseMountOptions:
       - implicit-dirs
       - metadata-cache:negative-ttl-secs:0
       - metadata-cache:ttl-secs:-1
       - metadata-cache:stat-cache-max-size-mb:-1
       - metadata-cache:type-cache-max-size-mb:-1
       - file-system:kernel-list-cache-ttl-secs:0
       - read_ahead_kb=1024
       - write:enable-streaming-writes:true
   

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

   • ‫MTC_CONFIG_NAME: השם של CheckpointConfiguration. השם הזה הוא גלובלי לאשכול ולא ספציפי לעבודה.
   • ‫GCS_BUCKET: השם של הקטגוריה ב-Cloud Storage שבה יישמרו נתוני נקודות הבדיקה. משתמשים בקטגוריה שהגדרתם בשלב הגדרת קטגוריה של Cloud Storage עם הרשאות.

   • ‫MACHINE_TYPE: סוג המכונה של המאיצים התואמים. הערך יכול להיות אחד מהערכים הבאים:

     • ‫TPU v5p: ct5p-hightpu-4t
     • TPU v5e: ct5e-hightpu-4t
     • ‫TPU v6e: ct6e-standard-4t
     • ‫Ironwood (TPU7x): tpu7x-standard-4t
     • מעבדים גרפיים NVIDIA H100 80GB (סדרה A3):
       • A3 High: a3-highgpu-8g
       • A3 Mega: a3-megagpu-8g
       • A3 Ultra: a3-ultragpu-8g

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

   • ‫TOLERATION_KEY: השדה הזה מאפשר לתזמן את מנהל ההתקן של CSI בצמתים עם taints תואמים. למידע נוסף על אופן הפעולה של taints בסוגים שונים של מאיצים, קראו את הדפים הבאים:

     • יחידות TPU
     • מעבדים גרפיים

   • ‫IN_MEMORY_VOLUME_SIZE: הגודל של מטמון נקודות הבדיקה בזיכרון. מציינים את הכמות והיחידה (לדוגמה, ‎200 Gi).הערך הזה צריך להיות:

     • גודל נקודת הבדיקה המקומית של יחידות TPU כפול 2.2
     • גודל נקודת הבדיקה המקומית של מעבדי GPU עם עמית אחד כפול 6.6.

2. מחילים את המניפסט:

   kubectl apply -f checkpoint.yaml
   

3. בודקים שמנהל ההתקן של CSI פועל:

   kubectl get pod -n gke-managed-checkpointing
   

   הפלט אמור להיראות כך: יהיו כמה רשומות, אחת לכל צומת מואץ.

   NAME                                                          READY   STATUS    RESTARTS   AGE
   multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s
   

הסרת מנהל ההתקן של Multi-Tier Checkpointing CSI

כדי לבטל את הפריסה של מנהל ההתקנים של CSI לנקודות ביקורת מרובות רמות, צריך למחוק את המשאב CheckpointConfiguration. הבקר של Multi-Tier Checkpointing מסיר את מנהל ההתקן של CSI מהצמתים. הפעולה הזו מסירה את דיסקי ה-RAM ומפנה זיכרון לעומסי עבודה אחרים. לדוגמה:

kubectl delete -f checkpoint.yaml

ניהול שמירת הנתונים ומנגנון איסוף הזבל בגיבויים של Cloud Storage

אתם אחראים להטמעת מדיניות שמירת נתונים לגיבויים של נקודות ביקורת ב-Cloud Storage. ב-Multi-Tier Checkpointing, גיבויי נקודות הבדיקה נכתבים רק ב-Cloud Storage, והם אף פעם לא משתנים או נמחקים.

יש הרבה כלים בקוד פתוח שיכולים לטפל בשמירה ובאיסוף, כולל:

• backup-warden
• logrotate

בדוגמה הבאה נעשה שימוש ב-backup-warden, כאשר הספרייה backup נטענת למיקום גיבוי שמשתמש ב-Cloud Storage FUSE:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

עדכון המניפסט של JobSet של עומס העבודה

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

לדוגמה, כדי להרחיב את JobSet לדוגמה מתוך פריסת TPU Multislices ב-GKE, מבצעים את השלבים הבאים:

1. מוסיפים את השורות הבאות למאגר jax-tpu.

   volumeMounts:
   - name: checkpoint
     mountPath: CHECKPOINT_DIR
   

   מחליפים את CHECKPOINT_DIR בנתיב לספריית נקודות הבדיקה. זה המיקום שבו נוצר ה-replicator.yaml, ושבו מתבצעת פעולת השמירה של נקודת הבדיקה באמצעות Multi-Tier Checkpointing. מידע נוסף זמין במאמר שילוב של Multi-Tier Checkpointing באפליקציה.

2. מוסיפים את השורות הבאות לשדה spec.template.spec במפרט המשימה.

   volumes:
   - name: checkpoint
     csi:
       driver: multitier-checkpoint.csi.storage.gke.io
   

שילוב של Multi-Tier Checkpointing באפליקציה

כדי לשתף מידע על מיקומי נקודות הביקורת ועל המוכנות לשכפול, צריך לשנות את האפליקציה כך שתשתמש בפרוטוקול הבא כדי לתקשר עם Multi-Tier Checkpointing.

השילוב הזה כבר מוטמע במסגרות ובספריות הבאות:

1. ‫MaxText (באמצעות Orbax): ל-MaxText יש תמיכה מובנית ב-Multi-Tier Checkpointing‏ (MTC). מידע נוסף זמין במאמרי העזרה בנושא MaxText MTC.
2. כל מודל שמשתמש ב-Orbax: השילוב מורכב משני רכיבים עיקריים: מודול ההפעלה של MTC להגדרה ראשונית (לדוגמה, ראו MaxText), וReplicator Checkpoint Manager לשמירה ולשחזור של נקודות ביקורת מדיסק RAM מקומי (לדוגמה, ראו MaxText).
3. דוגמה להפניה ל-PyTorch: הדוגמה הזו זמינה לעומסי עבודה שמשתמשים ב-NeMo עם PyTorch.

הפעלה

בקטע הזה מתוארים השלבים הראשוניים שהאפליקציה צריכה לבצע כדי להשתמש ב-Multi-Tier Checkpointing.

Replicator הוא רכיב ליבה של Multi-Tier Checkpointing, שפועל בכל צומת כחלק ממנהל ההתקן של CSI. ה-Replicator מנהל את השכפול של נקודות הבדיקה בשכבות האחסון, מדיסק ה-RAM המקומי לצומתי עמיתים ולאחסון חיצוני כמו Cloud Storage.

קובץ replicator.yaml משמש כמישור בקרה דינמי בין משימת האימון של ה-ML (קוד המסגרת) לבין רכיב המשכפל. אפליקציית ה-ML שלכם יוצרת את הקובץ הזה באופן פרוגרמטי בנפח אחסון מקומי (RAMDisk), שאפשר לגשת אליו גם דרך משימת האימון וגם דרך שירות המשכפל. קובץ המניפסט הזה מאפשר למסגרת למידת המכונה לספק הוראות להגדרת זמן ריצה ולניהול מחזור החיים לרפליקטור, שונות מפרמטרים סטטיים של התשתית (לדוגמה, תדירות ההעלאה ל-Cloud Storage) שהוגדרו במהלך הגדרת ה-Backend.

דוגמה קונקרטית לאינטראקציה הזו מופיעה במאמר:

• פרויקט MaxText, שמשתמש בארכיטקטורה הזו ל-JAX ב-Cloud TPU.
• דוגמה לארכיטקטורה של PyTorch, שמשתמשת בארכיטקטורה הזו עם PyTorch ו-NVIDIA NeMO במעבדי GPU.

במהלך ההפעלה, האפליקציה צריכה לבצע את הפעולות הבאות:

1. מחכים עד שהקובץ replicator.yaml לא קיים, מה שמציין שהכלי Replicator מוכן להגדרה על ידי האפליקציה. הקובץ replicator.yaml נוצר במיקום CHECKPOINT_DIR שהגדרתם בקטע עדכון המניפסט של JobSet של עומס העבודה.

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

2. האפליקציה או משימת ה-ML יוצרות את הקובץ replicator.yaml עם הגדרות דומות לאלה:

   Orbax

   job-name: orbax
   framework: orbax
   assume-data-parallelism: 3
   node-rank: 0
   nodes: 32
   peer-ranks: [1, 16] or peers-per-node: 2
   backup-interval-minutes: 30
   

   PyTorch

   job-name: nemo
   framework: pytorch.distributed
   node-rank: 0
   nodes: 32
   peer-ranks: [1, 16] or peers-per-node: 2
   backup-interval-minutes: 30
   

   ההגדרה לדוגמה הזו כוללת את השדות הבאים:

   • ‫name: השם של משימת האימון.
   • ‫framework: ה-framework של למידת המכונה שבו נעשה שימוש במשימת האימון.
   • ‫node-rank: המזהה הייחודי של הצומת הנוכחי במשימת האימון המבוזרת. זהו דירוג הצומת שיוצר את הקובץ הזה. לכל צומת שמשתתף בהרצה יהיה דירוג משלו.
   • ‫nodes: המספר הכולל של הצמתים שמשתתפים במשימת האימון המבוזרת. הערך הזה מגיע מהמטא-נתונים של ה-Pod. גם משימת האימון של ה-ML יכולה לראות את הערך הזה.
   • ‫peer-ranks או peers-per-node: שתי דרכים חלופיות לציין את טופולוגיית הרפליקציה. רק אחד משני הפרמטרים האלה צריך להופיע.
     • ‫peer-ranks: דירוגים מפורשים של צומתי עמיתים שאליהם צריך לשכפל את נתוני נקודת הבדיקה של הצומת הנוכחי. כך אפשר לשלוט בצורה מדויקת באילו צמתים ספציפיים משמשים כשותפי רפליקציה.
     • ‫peers-per-node: מספר הצמתים העמיתים לכל צומת שהכלי Replicator צריך לבחור באופן אוטומטי לשכפול.
   • ‫backup-interval-minutes: התדירות, בדקות, שבה מתבצע גיבוי של נקודות ביקורת ל-Cloud Storage. מומלץ להגדיר את הערך הזה ל-30 דקות או יותר.

3. מחכים עד שהמערכת תמחק את קובץ ה-replicator.yaml החדש. ההודעה הזו מציינת שהכלי לשכפול הופעל מחדש וביצע ניקוי. השלב הזה מאפשר לכם להימנע מקבצים זמניים או לא עדכניים בנפח האחסון המקומי כשהאפליקציה מבצעת את השלבים בקטע הבא.

שחזור מנקודת השמירה התקינה האחרונה (LKG)

1. אחרי שהרפליקטור מאותחל, Multi-Tier Checkpointing יוצר קישור סימבולי אחד לכל TPU או GPU worker. הקישורים הסמליים האלה נוצרים באותו נפח אחסון מקומי שמוטען כמו קובץ replicator.yaml, שבו העבודה שומרת נקודות ביקורת.

   הקישורים הסימבוליים הם בפורמט <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

2. משחזרים כל worker מהקובץ .restore המתאים. דוגמה מופיעה בקטע הבא, שבו מוסבר על מנהל נקודות ביקורת משוכפל של Orbax.

כברירת מחדל, נקודת הבדיקה LKG היא נקודת הבדיקה המלאה האחרונה שאפשר למצוא באמצעות Multi-Tier Checkpointing. עדיף להשתמש בנקודות ביקורת בתוך האשכול (בדיסק RAM המקומי ובצמתים מקבילים), ולהשתמש בגיבויים של Cloud Storage רק כשאין נקודת ביקורת מלאה בתוך האשכול. כדי לשנות את הבחירה הזו, אפשר לעיין במאמר בנושא שחזור מגיבוי ספציפי.

שחזור מגיבוי ספציפי

אתם יכולים לשחזר עבודה מגיבוי ספציפי ב-Cloud Storage במקום מנקודת הבדיקה האחרונה הידועה (LKG) – לדוגמה, כדי לבטל את ההדרכה ולחזור למצב קודם. לפני שהעבודה מתחילה, צריך לבצע את השלבים הבאים:

1. מוצאים את הגיבוי שרוצים לשחזר ממנו. התכונה Multi-Tier Checkpointing כותבת גיבויים לקטגוריה שהגדרתם ב-CheckpointConfiguration, בספרייה שנקראת על שם המשימה. כל גיבוי הוא ספריית משנה שנקראת על שם זמן ה-UTC שבו בוצע הגיבוי, בפורמט YYYY-MM-DD_HH-MM:

   gcloud storage ls gs://GCS_PATH/JOB_NAME/
   

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

   • ‫GCS_PATH: הנתיב ב-Cloud Storage שבו נקודות הבקרה בכמה רמות כותבות את הגיבויים. הערך הזה הוא הקטגוריה שציינתם בשדה cloudStorageBucketName של CheckpointConfiguration, כולל כל ספרייה בקטגוריה.
   • ‫JOB_NAME: השם של משימת האימון, כמו שצוין בשדה job-name בקובץ replicator.yaml.

2. בספריית המשימות בקטגוריה, יוצרים קובץ בשם restore.version. הקובץ צריך להכיל רק את השם של ספריית הגיבוי, לדוגמה:

   echo "BACKUP_NAME" | gcloud storage cp - gs://GCS_PATH/JOB_NAME/restore.version
   

   מחליפים את BACKUP_NAME בשם של ספריית הגיבוי, לדוגמה, 2025-06-30_14-30.

3. מתחילים את המשימה.

כשהקובץ restore.version קיים, נקודות הביקורת בכמה רמות פועלות באופן הבא:

• כל הצמתים משוחזרים מהגיבוי שצוין. נקודות ביקורת בתוך האשכול לא נלקחות בחשבון, גם אם הן עדכניות יותר מהגיבוי.
• קובץ restore.version נמחק אוטומטית אחרי שהמשימה שומרת את נקודת הבדיקה הראשונה שלה אחרי השחזור. השינוי חל רק פעם אחת: אם המשימה מופעלת מחדש מאוחר יותר, היא משוחזרת מנקודת הבדיקה של LKG כרגיל.
• אם הגיבוי שצוין לא קיים או שהוא לא שלם, העבודה מתחילה ללא נקודת ביקורת. אם אתם משתמשים ב-Multi-Tier Checkpointing, המערכת לא חוזרת לנקודות ביקורת באשכול או לגיבויים אחרים.

שמירת נקודת הציון

האפליקציה מבצעת את השלבים האלה כמה פעמים בזמן שמשימת האימון מתקדמת. פעולות השמירה מתבצעות במיקום CHECKPOINT_DIR שהגדרתם בעדכון המניפסט של JobSet של עומס העבודה.

Orbax

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

מידע נוסף על השימוש במנהל נקודות הבדיקה של Orbax replicator זמין בקובץ checkpointing של MaxText. דוגמה לאינטראקציה של שירות משכפל מופיעה בקובץ MaxText max_utils.

PyTorch

משתמשים ב-InClusterLocalCheckpointIO כ-pytorch_lightning.CheckpointIO מותאם אישית כדי להפעיל את התכונה 'יצירת נקודות בקרה מבוזרות' עם אחסון מקומי. הפקודה הבאה לדוגמה מפעילה נקודות בקרה בכמה רמות, באמצעות הטמעה לדוגמה שמבוססת על NVIDIA NeMo framework:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

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

• ‫CHECKPOINT_DIR: הנתיב לספריית נקודות הבדיקה.
• ‫JOB_NAME: השם של עומס העבודה של משימת האימון.

שדרוגי אשכולות

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

אם אתם רוצים שהמפרט של daemonset יישאר ללא שינוי, אתם לא צריכים לעשות כלום.

פתרון בעיות

בקטע הזה מופיע מדריך לפתרון בעיות שקשורות לנקודות ביקורת מרובות רמות. לפתרון בעיות כלליות באחסון, קראו את המאמר פתרון בעיות ב-Cloud Storage ב-GKE.

התכונה "נקודות ביקורת בכמה רמות" לא מופעלת

השגיאה הבאה מציינת ש-Multi-Tier Checkpointing לא מופעל באשכול:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

יכול להיות שתיתקלו בשגיאה הזו אחרי שתריצו את kubectl apply -f checkpoint.yaml בשלב יצירת CheckpointConfiguration.

כדי לפתור את הבעיה, בודקים אם הפעלתם את Multi-Tier Checkpointing באשכול באמצעות הפקודה הבאה:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

אם האפשרות Multi-Tier Checkpointing מופעלת, הפלט אמור להיראות כך:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

אם התכונה Multi-Tier Checkpointing לא מופעלת, צריך לעדכן את האשכול כדי להפעיל אותה.

מנהל ההתקן של CSI ליצירת נקודות ביקורת בכמה רמות לא יכול לטעון אמצעי אחסון

יכול להיות שתיתקלו בבעיה הזו אם מנהל ההתקן של CSI לא מצליח לטעון את עוצמת הקול של Cloud Storage. יכול להיות שיהיו כמה שורות דומות לזו.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

כדי לפתור את הבעיה, בודקים את האירועים של ה-Pod של מנהל התקן CSI, כמו בדוגמה הבאה:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

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

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

• מידע נוסף על פריסת TPU Multislice ב-Google Kubernetes Engine
• איך מבצעים אופטימיזציה של מנהל התקן ה-CSI של Cloud Storage FUSE ל-Google Kubernetes Engine כדי לשפר את הביצועים
• אפשרויות לנקודות ביקורת ב-Orbax