הצגת מודלים מסוג LLM באמצעות TPU מרובה מארחים ב-GKE עם JetStream ו-Pathways

רגילה

במדריך הזה נסביר איך להכניס לשימוש בסביבת הייצור מודלים גדולים של שפה (LLM) המתקדמים ביותר (SOTA), כמו Llama 3.1 405B ב-Google Kubernetes Engine ‏(GKE) באמצעות יחידות לעיבוד טנסורים (TPU) בכמה צמתים.

במדריך הזה נדגים איך להשתמש בטכנולוגיות ניידות בקוד פתוח – Kubernetes,‏ JetStream,‏ Pathways on Cloud ו-API‏ LeaderWorkerSet‏ (LWS) – כדי לפרוס ולהפעיל עומסי עבודה של AI/ML ב-GKE, תוך ניצול היתרונות של GKE: שליטה גרנולרית, יכולת הרחבה, עמידות, ניידות וחסכוניות.

רקע

הגודל של מודלים גדולים של שפה (LLM) גדל, והם כבר לא נכנסים לפרוסת TPU אחת של מארח. לביצוע הסקה של ML, אפשר להשתמש ב-Pathways on Cloud כדי להריץ הסקה מרובת מארחים בקנה מידה גדול ב-GKE במספר צמתי TPU שמחוברים זה לזה. במדריך הזה נסביר איך להקצות אשכול GKE עם פרוסות TPU מרובות מארחים, להשתמש בקבצים הבינאריים של Pathways on Cloud, להפעיל את שרת JetStream עם מסגרת MaxText ולשלוח בקשות הסקה מרובות מארחים.

כשמפעילים מודל LLM באמצעות TPU ב-GKE עם JetStream,‏ MaxText ו-Pathways, אפשר ליצור פתרון חזק ומוכן לייצור עם כל היתרונות של Kubernetes מנוהל, כולל יעילות בעלויות, יכולת הרחבה וזמינות גבוהה יותר. בקטע הזה מתוארות הטכנולוגיות העיקריות שבהן נעשה שימוש במדריך הזה.

מידע על מעבדי TPU

יחידות TPU הן מעגלים משולבים לאפליקציות ספציפיות (ASIC) שפותחו על ידי Google כדי להאיץ למידת מכונה ומודלים של AI שנבנו באמצעות frameworks כמו TensorFlow,‏ PyTorch ו-JAX.

לפני שמשתמשים ב-TPU ב-GKE, מומלץ להשלים את תוכנית הלימודים הבאה:

1. מידע על הזמינות של גרסאות TPU עדכניות זמין בארכיטקטורת המערכת של Cloud TPU.
2. מידע נוסף על TPUs ב-GKE

במדריך הזה נסביר איך להפעיל את מודל Llama 3.1-405B. ‫GKE פורס את המודל בצמתים של TPU v6e עם כמה מארחים, עם טופולוגיות של TPU שהוגדרו על סמך דרישות המודל להצגת הנחיות עם זמן אחזור נמוך.

תוכניות לימודים ב-Cloud

‫Pathways היא שכבת תזמור בקנה מידה גדול למאיצים. התכונה 'נתיבי למידה' נועדה במפורש לאפשר לכם לחקור מערכות חדשות ורעיונות למחקר למידת מכונה, תוך שמירה על ביצועים מתקדמים במודלים הנוכחיים. ‫Pathways מאפשרת לתהליך יחיד של לקוח JAX לתאם חישובים על פני פרוסות TPU גדולות אחת או יותר, וכך לייעל חישובי ML שמתבצעים על מאות או אלפי שבבי TPU.

JetStream

‫JetStream הוא מודל קוד פתוח ש-Google פיתחה כדי להסיק מסקנות. ‫JetStream מאפשרת הסקת מסקנות עם ביצועים גבוהים, תפוקה גבוהה ואופטימיזציה של הזיכרון ב-TPU וב-GPU. ‫JetStream מספקת אופטימיזציות מתקדמות של ביצועים, כולל עיבוד רציף של קבוצות נתונים, אופטימיזציות של מטמון KV וטכניקות קוונטיזציה, כדי להקל על פריסת מודלים גדולים של שפה (LLM). ‫JetStream מאפשרת הצגת PyTorch/XLA ו-JAX TPU כדי לבצע אופטימיזציה של הביצועים.

MaxText

‫MaxText הוא הטמעה של מודל שפה גדול (LLM) של JAX, שניתן להתאמה, לשינוי גודל ולביצועים טובים. הוא מבוסס על ספריות קוד פתוח של JAX כמו Flax,‏ Orbax ו-Optax. ההטמעה של מודל שפה גדול (LLM) ב-MaxText, שכוללת רק פענוח, כתובה ב-Python. הוא מסתמך במידה רבה על מהדר XLA כדי להשיג ביצועים גבוהים בלי צורך ליצור ליבות מותאמות אישית.

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

Llama 3.1 405B

‫Llama 3.1 405B הוא מודל שפה גדול (LLM) של Meta שנועד למגוון משימות של עיבוד שפה טבעית (NLP), כולל יצירת טקסט, תרגום ומענה לשאלות. ‫GKE מספק את התשתית הנדרשת לתמיכה בצרכים של אימון מבוזר והצגת מודלים בקנה מידה כזה.

מידע נוסף זמין במאמרי העזרה בנושא Llama.

ארכיטקטורה

בקטע הזה מתוארת ארכיטקטורת GKE שבה נעשה שימוש במדריך הזה. הארכיטקטורה כוללת אשכול GKE Standard שמקצה TPU ומארח רכיבי JetStream ו-Pathways כדי לפרוס את המודל ולהציג אותו.

בתרשים הבא מוצגים רכיבי הארכיטקטורה הזו:

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

• אשכול אזורי של GKE Standard.
• מאגר צמתים של פרוסת TPU מרובת מארחים שמארח את הפריסה של JetStream ואת רכיבי Pathways.
• ‫Pathways resource manager מנהל את משאבי המאיצים ומתאם את ההקצאה של מאיצים לעבודות של משתמשים.
• הקואורדינטות Pathways client מתואמות עם Pathways resource manager כדי לקבוע איפה התוכניות שעברו קומפילציה יוצבו לצורך הפעלה.
• ‫Pathways worker פועל ומבצע חישובים במכונות מאיץ, ושולח נתונים בחזרה לעומס העבודה שלכם דרך שרת ה-proxy של IFRT.
• ‫IFRT proxy client מטמיע את ה-API של Interim Framework Runtime (IFRT) של OSS ומשמש כגשר תקשורת בין עומס העבודה לבין רכיבי Pathways.
• ה-IFRT proxy server מקבל בקשות מה-IFRT proxy client ומעביר אותן ל-Pathways client, ומחלק את העבודה.
• מאגר התגים JetStream-Pathways מספק שרת הסקה מבוסס-JAX שמקבל בקשות הסקה ומקצה את תהליכי הביצוע שלו ל-Pathways workers
• רכיב השירות מפזר את התנועה הנכנסת לכל העותקים המשוכפלים של JetStream HTTP.
• ‫JetStream HTTP הוא שרת HTTP שמקבל בקשות כעטיפה לפורמט הנדרש של JetStream ושולח אותן ללקוח ה-GRPC של JetStream.

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

נכנסים לחשבון 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 the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

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 the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

• צריך לוודא שיש לכם בפרויקט את התפקיד או התפקידים הבאים: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

  בדיקת התפקידים

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

     כניסה לדף IAM
  2. בוחרים את הפרויקט.

  3. בעמודה Principal (חשבון המשתמש), מוצאים את כל השורות שבהן מופיע השם שלכם או של קבוצה שאתם נכללים בה. כדי לברר באילו קבוצות אתם נכללים, פנו לאדמין.

  4. בודקים את העמודה Role בכל השורות שבהן מצוין או מופיע השם שלכם, כדי לראות אם רשימת התפקידים כוללת את התפקידים הנדרשים.

  מתן התפקידים

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

     כניסה לדף IAM
  2. בוחרים את הפרויקט.
  3. לוחצים על Grant access.

  4. בשדה New principals, מזינים את מזהה המשתמש. ‫ בדרך כלל מזהה המשתמש הוא כתובת האימייל של חשבון Google.

     ‫
  5. לוחצים על Select a role ומחפשים את התפקיד.
  6. כדי לתת עוד תפקידים, לוחצים על Add another role ומוסיפים אותם.
  7. לוחצים על Save.

• מוודאים שיש לכם מספיק מכסת שימוש ל-16 שבבי TPU v6e PodSlice Lite. במדריך הזה משתמשים במופעים על פי דרישה.

גישה למודל

כדי לקבל גישה לנקודת הבדיקה Meta Llama 3.1-405B לפריסה ב-GKE, פועלים לפי השלבים הבאים:

1. חתימה על הסכם ההסכמה לרישיון.
2. נכנסים אל דף ההורדות של Meta Llama.
3. קוראים ומאשרים את התנאים וההגבלות של המודל כדי לקבל את כתובת ה-URL שנדרשת להורדת המודל.
4. כדי להוריד את נקודת הבדיקה של המודל, צריך למצוא את מזהה המודל המתאים. רשימה של מודלים נתמכים ומזהי המודלים זמינה במאמרי העזרה בנושא llama CLI. לדוגמה, משתמשים ב-Llama 3.1-405B-Instruct:bf16-mp16 עבור מודל Llama 3.1-405B.

הכנת הסביבה

במדריך הזה משתמשים ב-Cloud Shell כדי לנהל משאבים שמתארחים ב-Google Cloud. ב-Cloud Shell מותקן מראש התוכנה שצריך למדריך הזה, כולל kubectl ו- ה-CLI של gcloud.

כדי להגדיר את הסביבה באמצעות Cloud Shell:

1. ב Google Cloud מסוף, מפעילים סשן של Cloud Shell על ידי לחיצה על Activate Cloud Shell בGoogle Cloud מסוף. תופעל סשן בחלונית התחתונה של Google Cloud המסוף.

2. מגדירים את משתני הסביבה שמוגדרים כברירת מחדל:

   gcloud config set project PROJECT_ID
   gcloud config set billing/quota_project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export CLUSTER_NAME=CLUSTER_NAME
   export BUCKET_NAME=BUCKET_NAME
   export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
   export NODE_LOCATION=NODE_LOCATION
   export CLUSTER_VERSION=CLUSTER_VERSION
   export MACHINE_TYPE=ct6e-standard-4t
   export TPU_TYPE=v6e
   export TOPOLOGY=4x4
   export WORKERS_PER_SLICE=4
   

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

   • ‫PROJECT_ID: Google Cloud מזהה הפרויקט.
   • ‫CLUSTER_NAME: השם של אשכול GKE.
   • ‫BUCKET_NAME: שם הקטגוריה של Cloud Storage. אין צורך לציין את הקידומת gs://.
   • ‫CONTROL_PLANE_LOCATION: האזור ב-Compute Engine שבו נמצאים אשכול GKE, קטגוריית Cloud Storage וצמתי TPU. האזור מכיל אזורים שבהם זמינים סוגי מכונות TPU v6e (לדוגמה, us-east1,‏ us-east5,‏ europe-west4,‏ asia-northeast1 או us-south1).
   • ‫NODE_LOCATION: האזור שבו משאבי ה-TPU זמינים (לדוגמה, us-east1-d).
   • ‫CLUSTER_VERSION: גרסת GKE, שצריכה לתמוך בסוג המכונה שרוצים להשתמש בו. שימו לב: יכול להיות שהגרסה של GKE שמוגדרת כברירת מחדל לא תהיה זמינה ל-TPU שלכם. רשימה של גרסאות GKE מינימליות שזמינות לפי סוג מכונת TPU מופיעה במאמר זמינות של TPU ב-GKE.
   • ‫MACHINE_TYPE: סוג המכונה v6e.
   • ‫TPU_TYPE: קידומת שמשמשת למתן שמות למאגרי צמתים (v6e).
   • ‫TOPOLOGY: הטופולוגיה של TPU v6e.
   • ‫WORKERS_PER_SLICE: מספר הצמתים לכל מאגר צמתים או פרוסת TPU.

יצירה והגדרה של Google Cloud משאבים

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

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

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

יצירת אשכול GKE

1. יוצרים אשכול GKE Standard אזורי:

   gcloud container clusters create ${CLUSTER_NAME} \
       --project=${PROJECT_ID} \
       --cluster-version=${CLUSTER_VERSION} \
       --location=${CONTROL_PLANE_LOCATION} \
       --scopes=cloud-platform \
       --machine-type=n2-standard-32
   

   יצירת האשכול עשויה להימשך כמה דקות.

   מחליפים את CLUSTER_VERSION בגרסת האשכול המתאימה.

2. יוצרים מאגר צמתים אחד של TPU v6e עם טופולוגיה של 4x4 וארבעה צמתים בכל מאגר:

   gcloud container node-pools create multihost-np \
       --project=${PROJECT_ID} \
       --location=${CONTROL_PLANE_LOCATION} \
       --node-locations=${NODE_LOCATION} \
       --cluster=${CLUSTER_NAME} \
       --machine-type=${MACHINE_TYPE} \
       --num-nodes=${WORKERS_PER_SLICE} \
       --tpu-topology=${TOPOLOGY} \
       --scopes cloud-platform \
       --placement-type=COMPACT \
       --workload-metadata=GCE_METADATA
   

   הערה: הסימן --workload-metadata=GCE_METADATA מציין שאיחוד הזהויות של עומסי עבודה ל-GKE מושבת במאגר הצמתים. במאגרי צמתים שמריצים את Pathways עם JetStream ו-MaxText צריך להשבית את איחוד הזהויות של עומסי עבודה ל-GKE.

הגדרת חשבון שירות לגישה לאובייקט ב-Storage

הגדרה של חשבון שירות של Kubernetes שיפעל כחשבון שירות של IAM.

1. יוצרים חשבון שירות ב-IAM לאפליקציה:

   gcloud iam service-accounts create jetstream-pathways
   

2. מוסיפים קשירת מדיניות IAM לחשבון השירות ב-IAM כדי לנהל את Cloud Storage. כדי לאפשר לחשבון השירות של IAM לגשת לקטגוריית האחסון שבה יישמר הצ'קפוינט:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
     --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
     --role roles/storage.objectUser
   
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
     --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
     --role roles/storage.insightsCollectorService
   

3. מוסיפים הערה לחשבון השירות של Kubernetes עם כתובת האימייל של חשבון השירות ב-IAM.

   kubectl annotate serviceaccount default \
   iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com
   

הגדרת אימות של Docker ל-Artifact Registry

מגדירים את Docker לאימות ב-Artifact Registry כדי שיוכל לשלוף את קובצי האימג' של Pathways שברשימת ההיתרים:

gcloud auth login
gcloud auth configure-docker

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

כדי להמיר נקודת ביקורת של Meta Llama 3.1-405B לנקודת ביקורת של הסקת מסקנות int8 שתואמת ל-MaxText, צריך לבצע את השלבים שמפורטים במאמר המרת נקודת ביקורת באמצעות Llama3.1-405B. הפריסה שלכם משתמשת בנקודת הבדיקה עם הדגל load_parameters_path.

יצירת קטגוריה של Cloud Storage לאחסון קבצים זמניים של Pathways

יוצרים קטגוריה של Cloud Storage לאחסון קבצים זמניים של Pathways, כמו מטמון הקומפילציה:

export PATHWAYS_BUCKET=PATHWAYS_BUCKET
gcloud storage buckets create gs://$PATHWAYS_BUCKET

פריסת JetStream-MaxText ו-Pathways

פורסים את שרת המודלים JetStream-MaxText ו-Pathways.

התחברות לאשכול GKE

gcloud container clusters get-credentials "${CLUSTER_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

פריסת LeaderWorkerSet ‏ (LWS) API

LWS הוא משאב מותאם אישית שנועד לפריסה ולניהול של אפליקציות מבוזרות עם שמירת מצב, במיוחד אפליקציות עם ארכיטקטורה של מוביל-עובד. הוא מתאים במיוחד לעומסי עבודה של AI/ML שבהם מודל גדול מחולק ומשרת מספר מכשירים בכמה צמתים.

VERSION=v0.6.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

מחכים עד שהבקר LeaderWorkerSet יהיה זמין באופן מלא:

kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m

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

deployment.apps/lws-controller-manager condition met

מוודאים שהבקר LeaderWorkerSet פועל במרחב השמות lws-system:

kubectl get pod -n lws-system

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

NAME                          READY   STATUS    RESTARTS    AGE
lws-controller-manager-abcd   1/1     Running   0           40s
lws-controller-manager-efgh   1/1     Running   0           40s

פריסת מניפסט עומס העבודה

1. שומרים את קובץ המניפסט הבא בשם jetstream-pathways-llama-3-1-405b-4x4.yaml:

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=1
             - --instance_type=tpuv6e:4x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             env:
             - name: LOG_LEVEL
               value: "INFO"
             args:
             - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
             - model_name=llama3.1-405b
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - async_checkpointing=false
             - steps=1
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=2
             - ici_tensor_parallelism=8
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=6
             - enable_single_controller=true
             - quantization=int8
             - quantize_kvcache=true
             - checkpoint_is_quantized=true
             - enable_model_warmup=true
             imagePullPolicy: Always
             ports:
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 600
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. מגדירים את הערך של השדה load_parameters_path לנתיב של נקודת הבדיקה שנוצר בתהליך ההמרה של נקודת הבדיקה.

   • בנקודת ביקורת bf16, הנתיב צריך להיות דומה ל-gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • עבור נקודת ביקורת של int8, הערך צריך להיות דומה ל-gs://OUTPUT_BUCKET_DIRECTORY/int8.

   מגדירים את הערך של השדה gcs_scratch_location לקטגוריית הנתיבים שיצרתם קודם.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/int8|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
   perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
   

   תיאור של משתנים והגדרות.

   בתיאורים הבאים מוסבר על המשתנים וההגדרות הרלוונטיים ל-JetStream, ל-Pathways ול-LeaderWorkerSet.

   • Jetstream
     • מאגר התגים הזה מבוסס על פרויקט JetStream של OSS. היא גם מכילה את ספריית הלקוח שמבוססת על gRPC של IFRT Proxy, שמספקת קצה עורפי של JAX שיכול להעביר הפעלה של תוכניות TPU לשרת Pathways Proxy.
     • מידע נוסף על דגלי JetStream-MaxText זמין במאמר תיאורים של דגלי JetStream-MaxText Server.
     • כדי להשתמש בנקודת ביקורת שעברה קוונטיזציה, אפשר להגדיר את הדגלים quantization, quantize_kvcache או checkpoint_is_quantized. אם נקודת הבדיקה היא bf16, אפשר להשמיט את הדגלים האלה.
   • ‫Pathways: תיאור של רכיבי Pathways הרלוונטיים מופיע במאמר רכיבי Pathways.
   • ‫LeaderWorkerSet: מידע נוסף על ההגדרות של LeaderWorkerSet זמין בהגדרת ה-API.

החלת מניפסט הפריסה

מחילים את המניפסט כדי לפרוס את השרת:

kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml

שרת המודל אמור להתחיל לפעול.

אימות ההתחלה של שרת המודל

יכול להיות שיידרשו כ-10 עד 20 דקות לשחזור נקודת הבדיקה בדגם 405B. יכול להיות שתצטרכו לחכות עוד זמן במהלך אתחול המודל אם הפעלתם את הדגל enable_model_warmup.

kubectl logs -f jetstream-pathways-0 -c jax-tpu

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

2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
...
...
...
INFO:     Started server process [7]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

Serve Llama 3.1-405b

כדי להכניס לשימוש בסביבת הייצור את מודל Llama 3.1-405b, צריך להגדיר העברה ליציאה אחרת:

kubectl port-forward svc/jetstream-svc 8000:8000

העברת יציאות מאפשרת לכם לגשת לשירות מחוץ לאשכול. אפשר לגשת לפריסת JetStream-Pathways דרך שירות ClusterIP של GKE. אפשר להגיע לשירותי ClusterIP רק מתוך האשכול.

אינטראקציה עם המודל

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

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

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

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

הפעולות הבאות בוצעו בהצלחה:

1. פריסת שרת מודלים של JetStream עם MaxText ו-Pathways ב-GKE באמצעות TPU.
2. נוצרה נקודת ביקורת של Llama 3.1-405B int8 בכתובת gs://BUCKET_NAME.
3. הוצג למשתמש והייתה לו אינטראקציה עם המודל.

הצגה מפורקת

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

• מילוי מראש: העברה קדימה של הנחיית הקלט כדי לאתחל את מטמון זוגות הערכים.

• פענוח: תהליך שיוצר אסימוני פלט באופן מצטבר, אסימון אחד בכל שלב, וערך אחד של מטמון KV בכל איטרציה.

1. מגדירים את משתני הסביבה שמוגדרים כברירת מחדל:

   export NODE_POOL_NAME=dis-v6e-8
   export NODE_POOL_SIZE=2
   export MACHINE_TYPE=ct6e-standard-4t
   export TOPOLOGY=2x4
   export WORKERS_PER_SLICE=2
   

2. יוצרים שני מאגרי צמתים שמשתמשים בצמתים מסוג v6e-8:

   for i in $(seq 1 ${NODE_POOL_SIZE}); do
     gcloud container node-pools create ${NODE_POOL_NAME}-${i}-np \
     --project=${PROJECT_ID} \
     --cluster=${CLUSTER_NAME} \
     --location=${CONTROL_PLANE_LOCATION} \
     --node-locations=${NODE_LOCATION} \
     --machine-type=${MACHINE_TYPE} \
     --num-nodes=${WORKERS_PER_SLICE} \
     --tpu-topology=${TOPOLOGY} \
     --scopes=cloud-platform \
     --workload-metadata=GCE_METADATA
   done
   

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

כדי להמיר נקודת ביקורת של Meta Llama 2-70B לנקודת ביקורת של int8 inference שתואמת ל-MaxText, צריך לבצע את השלבים שמפורטים במאמר המרת נקודת ביקורת באמצעות Llama2-70B. כשמאשרים את התנאים וההגבלות של מטא, בוחרים במודל Llama2-70B. הפריסה שלכם משתמשת בנקודת הבדיקה עם הדגל load_parameters_path.

מחליפים את הפרמטרים הבאים בקובץ checkpoint-job.yaml:

- --meta_url=META_URL
- --model_name=llama-2
- --model_path=Llama-2-70b-chat
- --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b

נקודת הבדיקה תשמש בפריסה עם הדגל load_parameters_path.

פריסת JetStream Pathways עם הגשה מפוצלת

1. שומרים את קובץ המניפסט הבא בשם jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       subGroupPolicy:
         subGroupSize: 2
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --server_port=38681
             - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
             - --xla_jf_auto_cross_replica_sharding=false
             - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
             - --xla_tpu_enable_windowed_einsum_for_all_gather=false
             - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
             - --xla_tpu_enable_experimental_fusion_cost_model=false
             - --xla_tpu_dot_dot_fusion_duplicated=false
             - --xla_tpu_dot_dot_fusion=true
             - --xla_jf_conv_input_fusion=true
             - --xla_jf_conv_output_fusion=true
             - --xla_tpu_rwb_fusion=false
             - --xla_tpu_copy_fusion_pad_unpad_ratio=0
             - --xla_tpu_licm_size_inflation_ratio=1
             - --xla_tpu_copy_elision_analysis_allowance=150000
             - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
             - --xla_tpu_order_dot_after_layout=true
             - --xla_jf_rematerialization_percent_shared_memory_limit=100
             - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
             - --xla_tpu_enforce_prefetch_fifo_order=false
             - --xla_tpu_prefetch_interval_picker_size_override=6000000
             - --xla_tpu_async_copy_bandwidth_scaling_factor=1
             - --xla_tpu_nd_short_transfer_max_chunks=-1
             - --xla_tpu_enable_aggressive_broadcast_priority_update=true
             - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
             - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
             - --xla_tpu_enable_copy_fusion=true
             - --xla_tpu_enable_cross_program_prefetch_freeing=false
             - --xla_tpu_enable_dot_strength_reduction=true
             - --xla_tpu_layout_use_dot_grouping=false
             - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
             - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
             - --xla_tpu_vector_load_fusion_window=1024
             - --xla_tpu_vector_store_fusion_window=256
             - --xla_jf_conv_reshape_fusion=false
             - --xla_tpu_input_conv_multi_users=false
             - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
             - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
             - --xla_tpu_dot_dot_fusion_separable_convs_only=false
             - --xla_tpu_enable_multi_level_nested_loop_fusion=true
             - --xla_tpu_nested_dot_fusion=true
             - --xla_tpu_enable_multi_level_nested_dot_fusion=false
             - --xla_jf_enable_multi_output_fusion=true
             - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
             - --xla_tpu_enable_flash_attention=true
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:       
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=2
             - --instance_type=tpuv6e:2x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             args:
             - MaxText/configs/base.yml
             - tokenizer_path=assets/tokenizer.llama2
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - model_name=llama2-70b
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=1
             - ici_tensor_parallelism=-1
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=27
             - checkpoint_is_quantized=true 
             - quantization=int8
             - quantize_kvcache=true
             - compute_axis_order=0,2,1,3
             - ar_cache_axis_order=0,2,1,3
             - stack_prefill_result_cache=True
             - inference_server=ExperimentalMaxtextDisaggregatedServer_8
             - inference_benchmark_test=True
             - enable_model_warmup=True
             env:
             - name: LOG_LEVEL
               value: "INFO"
             imagePullPolicy: Always
             securityContext:
               capabilities:
                 add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
             ports: 
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 240
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. מגדירים את הערך של השדה load_parameters_path לנתיב של נקודת הבדיקה שנוצר בתהליך ההמרה של נקודת הבדיקה.

   • בנקודת ביקורת bf16, הנתיב צריך להיות דומה ל-gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • עבור נקודת ביקורת של int8, הערך צריך להיות דומה ל-gs://OUTPUT_BUCKET_DIRECTORY/int8.

   מגדירים את הערך של השדה gcs_scratch_location לקטגוריית הנתיבים שיצרתם קודם.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/maxtext/llama-2-70b/int8|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

3. החלת המניפסט:

   kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

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

הפעולות הבאות בוצעו בהצלחה:

1. פריסת שרת מודלים של JetStream עם MaxText ו-Pathways ב-GKE באמצעות TPU ופריסה מפוצלת.
2. נוצרה נקודת ביקורת של Llama 2-70B int8 ב-gs://BUCKET_NAME.
3. הוצג למשתמש והייתה לו אינטראקציה עם המודל.

פתרון בעיות

• אם מוצגת ההודעה Empty reply from server, יכול להיות שההורדה של נתוני המודל מהקונטיינר לא הסתיימה. בודקים שוב את היומנים של ה-Pod כדי לראות את ההודעה Connected, שמציינת שהמודל מוכן להכניס לשימוש בסביבת הייצור.
• אם מוצגת ההודעה Connection refused, צריך לוודא שהעברת הפורטים פעילה.

הסרת המשאבים

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

מחיקת המשאבים שנפרסו

כדי להימנע מחיובים בחשבון Google Cloud על המשאבים שיצרתם במדריך הזה, מריצים את הפקודות הבאות ופועלים לפי ההנחיות:

gcloud container clusters delete ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}

gcloud iam service-accounts delete jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com

gcloud storage rm --recursive gs://${BUCKET_NAME}

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

• איך מריצים מודלים של Gemma ב-GKE ואיך מריצים עומסי עבודה של AI/ML שעברו אופטימיזציה באמצעות יכולות התזמור של פלטפורמת GKE.
• מידע נוסף על TPU ב-GKE
• מעיינים במאגר GitHub של JetStream.
• מעיינים ב-Vertex AI Model Garden.