במדריך הזה נסביר איך לפרוס מודל שפה גדול (LLM) ולהפעיל אותו באמצעות מאגר צמתים של פרוסת TPU מרובת-מארחים ב-Google Kubernetes Engine (GKE) עם Saxml לארכיטקטורה יעילה וניתנת להרחבה.
רקע
Saxml היא מערכת ניסיונית שמשמשת כפלטפורמה למסגרות Paxml, JAX ו-PyTorch. אפשר להשתמש ב-TPU כדי להאיץ את עיבוד הנתונים באמצעות המסגרות האלה. כדי להדגים את הפריסה של TPU ב-GKE, במדריך הזה מוצג מודל הבדיקה LmCloudSpmd175B32Test בגודל 175B. GKE פורס את מודל הבדיקה הזה בשני מאגרי צמתים של פרוסות TPU v5e עם טופולוגיה של 4x8.
כדי לפרוס את מודל הבדיקה בצורה נכונה, טופולוגיית ה-TPU הוגדרה על סמך גודל המודל. בהתחשב בכך שמודל של N מיליארד 16 ביט דורש בערך פי 2 (2xN) GB של זיכרון, מודל LmCloudSpmd175B32Test של 175B דורש בערך 350 GB של זיכרון. לשבב TPU v5e יחיד יש 16 GB. כדי לתמוך ב-350GB, ל-GKE נדרשים 21 שבבי TPU מדגם v5e (350/16= 21). על סמך מיפוי ההגדרות של TPU, ההגדרה הנכונה של TPU למדריך הזה היא:
- סוג המכונה:
ct5lp-hightpu-4t - טופולוגיה:
4x8(32 מספר שבבי TPU)
כשפורסים TPU ב-GKE, חשוב לבחור את טופולוגיית ה-TPU הנכונה להצגת מודל. למידע נוסף, אפשר לעיין במאמר בנושא תכנון ההגדרה של TPU.
מטרות
המדריך הזה מיועד למהנדסי MLOps או DevOps או לאדמינים של פלטפורמות שרוצים להשתמש ביכולות התיזמור של GKE כדי להציג מודלים של נתונים.
במדריך הזה מוסבר איך:
- מכינים את הסביבה עם אשכול רגיל. ל-cluster יש שני מאגרי צמתים של פרוסות TPU v5e עם טופולוגיה של
4x8. - פורסים את Saxml. Saxml צריך שרת אדמין, קבוצה של Pods שפועלים כשרת מודל, שרת HTTP מוכן מראש ומאזן עומסים.
- משתמשים ב-Saxml כדי להפעיל את ה-LLM.
התרשים הבא מציג את הארכיטקטורה שמוטמעת במדריך הבא:
לפני שמתחילים
- נכנסים לחשבון Google Cloud . אנחנו ממליצים למשתמשים חדשים ב- Google Cloud ליצור חשבון כדי שיוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
Roles required to select or create a project
- Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
-
Create a project: To create a project, you need the Project Creator role
(
roles/resourcemanager.projectCreator), which contains theresourcemanager.projects.createpermission. Learn how to grant roles.
-
Verify that billing is enabled for your Google Cloud project.
Enable the required API, if it is not already enabled.
Roles required to enable APIs
To enable APIs, you need the
serviceusage.services.enablepermission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
Roles required to select or create a project
- Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
-
Create a project: To create a project, you need the Project Creator role
(
roles/resourcemanager.projectCreator), which contains theresourcemanager.projects.createpermission. Learn how to grant roles.
-
Verify that billing is enabled for your Google Cloud project.
Enable the required API, if it is not already enabled.
Roles required to enable APIs
To enable APIs, you need the
serviceusage.services.enablepermission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.-
צריך לוודא שיש לכם בפרויקט את התפקיד או התפקידים הבאים: roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.policyAdmin
בדיקת התפקידים
-
נכנסים לדף IAM במסוף Google Cloud .
כניסה לדף IAM - בוחרים את הפרויקט.
-
בעמודה Principal, מחפשים את כל השורות שמזהות אתכם או קבוצה שאתם נכללים בה. כדי לדעת באילו קבוצות אתם נכללים, פנו לאדמין.
- בודקים את העמודה Role בכל השורות שבהן מצוין או מופיע השם שלכם, כדי לראות אם רשימת התפקידים כוללת את התפקידים הנדרשים.
מתן התפקידים
-
נכנסים לדף IAM במסוף Google Cloud .
כניסה לדף IAM - בוחרים את הפרויקט.
- לוחצים על Grant access.
-
בשדה New principals, מזינים את מזהה המשתמש. בדרך כלל מזהה המשתמש הוא כתובת האימייל של חשבון Google.
- לוחצים על Select a role ומחפשים את התפקיד.
- כדי לתת עוד תפקידים, לוחצים על Add another role ומוסיפים אותם.
- לוחצים על Save.
-
- מוודאים שיש בפרויקט מכסה מספקת ל-Cloud TPU ב-GKE.
הכנת הסביבה
במסוף Google Cloud , מפעילים מכונה של Cloud Shell:
פותחים את Cloud Shellמגדירים את משתני הסביבה שמוגדרים כברירת מחדל:
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION export BUCKET_NAME=PROJECT_ID-gke-bucketמחליפים את הערכים הבאים:
- PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- CONTROL_PLANE_LOCATION: האזור של Compute Engine במישור הבקרה של האשכול. בוחרים את האזור שבו
ct5lp-hightpu-4tזמין.
בפקודה הזו,
BUCKET_NAMEמציין את השם של קטגוריית האחסון Google Cloudשבה יאוחסנו הגדרות השרת של אדמין Saxml.
יצירת אשכול
אפשר להשתמש ב-Cloud Shell כדי:
יוצרים אשכול Standard שמשתמש באיחוד זהויות של עומסי עבודה ל-GKE:
gcloud container clusters create saxml \ --location=${CONTROL_PLANE_LOCATION} \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --cluster-version=VERSION \ --num-nodes=4מחליפים את
VERSIONבמספר הגרסה של GKE. GKE תומך ב-TPU v5e בגרסה 1.27.2-gke.2100 ואילך. מידע נוסף מופיע במאמר זמינות של TPU ב-GKE.יצירת האשכול עשויה להימשך כמה דקות.
יוצרים את מאגר הצמתים הראשון בשם
tpu1:gcloud container node-pools create tpu1 \ --location=${CONTROL_PLANE_LOCATION} \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=4x8 \ --num-nodes=8 \ --cluster=saxmlהערך של הדגל
--num-nodesמחושב על ידי חלוקת טופולוגיית ה-TPU במספר שבבי ה-TPU לכל פרוסת TPU. במקרה הזה: (4 * 8) / 4.יוצרים את מאגר הצמתים השני בשם
tpu2:gcloud container node-pools create tpu2 \ --location=${CONTROL_PLANE_LOCATION} \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=4x8 \ --num-nodes=8 \ --cluster=saxmlהערך של הדגל
--num-nodesמחושב על ידי חלוקת טופולוגיית ה-TPU במספר שבבי ה-TPU לכל פרוסת TPU. במקרה הזה: (4 * 8) / 4.
יצרת את המשאבים הבאים:
- אשכול Standard עם ארבעה צמתים של מעבד.
- שתי קבוצות של צומתי TPU v5e עם טופולוגיה של
4x8. כל מאגר של צמתי TPU מייצג שמונה צמתי TPU עם 4 שבבי TPU כל אחד.
מודל 175B צריך להיות מוגש בפרוסת TPU v5e מרובת מארחים עם 4x8
פרוסת טופולוגיה (32 שבבי TPU v5e) לכל הפחות.
יצירת קטגוריה של Cloud Storage
יוצרים קטגוריה של Cloud Storage לאחסון ההגדרות של שרת האדמין של Saxml. שרת אדמין פעיל שומר מעת לעת את המצב שלו ואת פרטי המודלים שפורסמו.
ב-Cloud Shell, מריצים את הפקודה הבאה:
gcloud storage buckets create gs://${BUCKET_NAME}
הגדרת הגישה של עומסי העבודה באמצעות איחוד זהויות של עומסי עבודה ל-GKE
מקצים לחשבון שירות של Kubernetes לאפליקציה ומגדירים את חשבון השירות של Kubernetes כך שיפעל כחשבון שירות של IAM.
מגדירים את
kubectlכדי לתקשר עם האשכול:gcloud container clusters get-credentials saxml --location=${CONTROL_PLANE_LOCATION}יוצרים חשבון שירות ב-Kubernetes שהאפליקציה תוכל להשתמש בו:
kubectl create serviceaccount sax-sa --namespace defaultיוצרים חשבון שירות ב-IAM לאפליקציה:
gcloud iam service-accounts create sax-iam-saמוסיפים קישור למדיניות IAM לחשבון השירות של IAM כדי לקרוא ולכתוב ב-Cloud Storage:
gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.adminכדי לאפשר ל-ServiceAccount של Kubernetes להתחזות לחשבון השירות של IAM, מוסיפים קשר בין כללי מדיניות של IAM בין שני חשבונות השירות. הקישור הזה מאפשר לחשבון השירות של Kubernetes לפעול בתור חשבון השירות של IAM, כך שחשבון השירות של Kubernetes יוכל לקרוא ולכתוב ב-Cloud Storage.
gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"מוסיפים הערה לחשבון השירות ב-Kubernetes עם כתובת האימייל של חשבון השירות ב-IAM. כך אפליקציית הדוגמה יודעת באיזה חשבון שירות להשתמש כדי לגשת Google Cloud לשירותים. לכן, כשהאפליקציה משתמשת בספריות לקוח סטנדרטיות של Google API כדי לגשת Google Cloud לשירותים, היא משתמשת בחשבון השירות הזה של IAM.
kubectl annotate serviceaccount sax-sa \ iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
פריסת Saxml
בקטע הזה תפרסו את שרת האדמין של Saxml ואת שרת המודל של Saxml.
פריסת שרת האדמין של Saxml
יוצרים את קובץ המניפסט
sax-admin-server.yamlהבא:מחליפים את
BUCKET_NAMEבקטגוריה של Cloud Storage שיצרתם קודם:perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yamlהחלת המניפסט:
kubectl apply -f sax-admin-server.yamlמוודאים שפוד השרת של האדמין פועל:
kubectl get deploymentהפלט אמור להיראות כך:
NAME READY UP-TO-DATE AVAILABLE AGE sax-admin-server 1/1 1 1 52s
פריסת שרת מודלים של Saxml
עומסי עבודה שפועלים בפרוסות TPU מרובות מארחים דורשים מזהה רשת יציב לכל Pod כדי לגלות עמיתים באותה פרוסת TPU. כדי להגדיר את המזהים האלה, משתמשים ב-IndexedJob, ב-StatefulSet עם Service ללא דפדפן GUI או ב-JobSet, שיוצר באופן אוטומטי Service ללא דפדפן GUI לכל המשימות ששייכות ל-JobSet. Jobset הוא API של עומסי עבודה שמאפשר לכם לנהל קבוצה של משימות Kubernetes כיחידה אחת. התרחיש הנפוץ ביותר לשימוש ב-JobSet הוא אימון מבוזר, אבל אפשר להשתמש בו גם להרצת עומסי עבודה של אצווה.
בקטע הבא מוסבר איך לנהל כמה קבוצות של Pods של שרת מודלים באמצעות JobSet.
מתקינים את JobSet בגרסה 0.2.3 ואילך.
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yamlמחליפים את
JOBSET_VERSIONבגרסת ה-JobSet. לדוגמה,v0.2.3.מוודאים שבקר JobSet פועל במרחב השמות
jobset-system:kubectl get pod -n jobset-systemהפלט אמור להיראות כך:
NAME READY STATUS RESTARTS AGE jobset-controller-manager-69449d86bc-hp5r6 2/2 Running 0 2m15sפריסת שני שרתים של מודלים בשתי קבוצות של צמתים עם חלקי TPU. שומרים את קובץ המניפסט
sax-model-server-set:מחליפים את
BUCKET_NAMEבקטגוריה של Cloud Storage שיצרתם קודם:perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-model-server-set.yamlבמניפסט הזה:
-
replicas: 2הוא מספר העותקים של Job. כל משימה מייצגת שרת מודלים. לכן, קבוצה של 8 מכשירי Pod. -
parallelism: 8ו-completions: 8שווים למספר הצמתים בכל מאגר צמתים. - הערך של
backoffLimit: 0חייב להיות אפס כדי לסמן את ה-Job כנכשל אם אחד מה-Pod נכשל. -
ports.containerPort: 8471היא יציאת ברירת המחדל לתקשורת בין המכונות הווירטואליות -
name: MEGASCALE_NUM_SLICESמבטל את ההגדרה של משתנה הסביבה כי GKE לא מריץ אימון Multislice.
-
החלת המניפסט:
kubectl apply -f sax-model-server-set.yamlבודקים את הסטטוס של ה-Pods של שרת הניהול ושרת המודלים של Saxml:
kubectl get podsהפלט אמור להיראות כך:
NAME READY STATUS RESTARTS AGE sax-admin-server-557c85f488-lnd5d 1/1 Running 0 35h sax-model-server-set-sax-model-server-0-0-nj4sm 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-1-sl8w4 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-2-hb4rk 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-3-qv67g 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-4-pzqz6 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-5-nm7mz 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-6-7br2x 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-7-4pw6z 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-0-8mlf5 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-1-h6z6w 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-2-jggtv 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-3-9v8kj 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-4-6vlb2 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-5-h689p 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-6-bgv5k 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-7-cd6gv 1/1 Running 0 24m
בדוגמה הזו יש 16 קונטיינרים של שרתי מודלים:
sax-model-server-set-sax-model-server-0-0-nj4sm ו-sax-model-server-set-sax-model-server-1-0-8mlf5 הם שני שרתי המודלים העיקריים בכל קבוצה.
באשכול Saxml יש שני שרתי מודלים שפרוסים בשתי קבוצות של צמתי פרוסות TPU מדגם v5e עם טופולוגיה של 4x8 בהתאמה.
פריסת שרת HTTP ומאזן עומסים של Saxml
משתמשים בתמונה הבאה של שרת HTTP מוכן מראש. שומרים את קובץ המניפסט
sax-http.yaml:מחליפים את
BUCKET_NAMEבקטגוריה של Cloud Storage שיצרתם קודם:perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yamlהחלת מניפסט
sax-http.yaml:kubectl apply -f sax-http.yamlמחכים עד שמאגר התגים בצד השרת מסוג HTTP Server יסיים את היצירה:
kubectl get podsהפלט אמור להיראות כך:
NAME READY STATUS RESTARTS AGE sax-admin-server-557c85f488-lnd5d 1/1 Running 0 35h sax-http-65d478d987-6q7zd 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-0-nj4sm 1/1 Running 0 24m ...מחכים עד שמוקצית ל-Service כתובת IP חיצונית:
kubectl get svcהפלט אמור להיראות כך:
NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGE sax-http-lb LoadBalancer 10.48.11.80 10.182.0.87 8888:32674/TCP 7m36s
שימוש ב-Saxml
טעינה, פריסה והפעלה של המודל ב-Saxml בפרוסת v5e TPU multihost:
טעינת המודל
מאחזרים את כתובת ה-IP של מאזן העומסים עבור Saxml.
LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}') PORT="8888"טוענים את מודל הבדיקה
LmCloudSpmd175Bבשתי קבוצות של צמתים מסוג v5e TPU slice:curl --request POST \ --header "Content-type: application/json" \ -s ${LB_IP}:${PORT}/publish --data \ '{ "model": "/sax/test/spmd", "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "replicas": 2 }'למודל הבדיקה אין נקודת ביקורת (checkpoint) שעברה כוונון עדין, והמשקלים נוצרים באופן אקראי. טעינת המודל עשויה להימשך עד 10 דקות.
הפלט אמור להיראות כך:
{ "model": "/sax/test/spmd", "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "replicas": 2 }בודקים את מוכנות המודל:
kubectl logs sax-model-server-set-sax-model-server-0-0-nj4smהפלט אמור להיראות כך:
... loading completed. Successfully loaded model for key: /sax/test/spmdהמודל נטען במלואו.
לקבל מידע על המודל:
curl --request GET \ --header "Content-type: application/json" \ -s ${LB_IP}:${PORT}/listcell --data \ '{ "model": "/sax/test/spmd" }'הפלט אמור להיראות כך:
{ "model": "/sax/test/spmd", "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "max_replicas": 2, "active_replicas": 2 }
פרסום המודל
הצגת בקשת פרומפט:
curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
"model": "/sax/test/spmd",
"query": "How many days are in a week?"
}'
בדוגמה הבאה אפשר לראות את התשובה של המודל. יכול להיות שהתשובה הזו לא תהיה משמעותית כי למודל הבדיקה יש משקלים אקראיים.
ביטול הפרסום של המודל
מריצים את הפקודה הבאה כדי לבטל את הפרסום של המודל:
curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
"model": "/sax/test/spmd"
}'
הפלט אמור להיראות כך:
{
"model": "/sax/test/spmd"
}
הסרת המשאבים
כדי להימנע מחיובים בחשבון Google Cloud בגלל השימוש במשאבים שנעשה במסגרת המדריך הזה, אפשר למחוק את הפרויקט שמכיל את המשאבים או להשאיר את הפרויקט ולמחוק את המשאבים בנפרד.
מחיקת המשאבים שנפרסו
כדי למחוק את האשכול שיצרתם במדריך הזה:
gcloud container clusters delete saxml --location ${CONTROL_PLANE_LOCATION}מוחקים את חשבון השירות:
gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.comמוחקים את הקטגוריה של Cloud Storage:
gcloud storage rm -r gs://${BUCKET_NAME}
המאמרים הבאים
- מידע על הגרסאות הנוכחיות של TPU זמין במאמר בנושא ארכיטקטורת המערכת של Cloud TPU.
- מידע נוסף על TPU ב-GKE