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

במדריך הזה נסביר איך להפעיל מודלים גדולים של שפה (LLM) מתקדמים, כמו 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. מידע נוסף על TPU ב-GKE

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

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

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

JetStream

‫JetStream הוא מודל קוד פתוח ש-Google פיתחה כדי להסיק מסקנות. ‫JetStream מאפשרת הסקת מסקנות עם ביצועים גבוהים, תפוקה גבוהה ואופטימיזציה של הזיכרון ב-TPU וב-GPU. ‫JetStream מספק אופטימיזציות מתקדמות של ביצועים, כולל חלוקה רציפה למנות (batching), אופטימיזציות של מטמון 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) של מטא שנועד למגוון משימות של עיבוד שפה טבעית (NLP), כולל יצירת טקסט, תרגום ומענה לשאלות. ‫GKE מספק את התשתית הנדרשת לתמיכה בצרכים של אימון מבוזר והצגת מודלים בקנה מידה כזה.

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

ארכיטקטורה

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

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

ארכיטקטורה של אשכול GKE עם מאגר צומתי 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.

    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 serviceusage.services.enable permission. 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.

    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.

    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 serviceusage.services.enable permission. 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.

    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 מסוף כדי להפעיל סשן של Cloud Shell. סשן יופעל בחלונית התחתונה של 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

הגדרת חשבון שירות לגישה לאובייקט ב-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 כדי שיוכל למשוך את קובצי האימג' של תוכניות הלימודים שנוספו לרשימת ההיתרים:

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

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

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

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

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

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

שחזור נקודת הביקורת בדגם 405B עשוי להימשך כ-10 עד 20 דקות. יכול להיות שתצטרכו לחכות עוד זמן במהלך אתחול המודל אם הפעלתם את הדגל 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. כשמאשרים את התנאים וההגבלות של Meta, בוחרים את המודל 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}

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